01 // THE BOTTLENECK
A model can look fast on a desktop GPU and still miss the budget on an embedded device. In one detector pipeline, the network was not the only problem. Python overhead, small CUDA launches, and accidental cudaMemcpy calls were taking enough time to matter.
PyTorch is great for trying ideas quickly, but in 2018 its eager path had real overhead. One tensor op is cheap. Hundreds of tiny ops in a real-time pipeline are not. The fix was to move the hot path into one compiled operator.
"The fastest code is the code that never runs. Fuse your operators, eliminate your copies, and get Python out of the hot path."
PyTorch's C++ extension API made that possible: write the operator in C++/CUDA, bind it with pybind11, and call it from Python like a normal torch op. This post walks through the kernel, memory handling, TorchScript export, and benchmarks.
02 // EXTENSION ARCHITECTURE
PyTorch provides two mechanisms for building C++ extensions: setuptools (ahead-of-time compilation via setup.py) and JIT compilation (on-the-fly via torch.utils.cpp_extension.load()). Both use the same underlying pipeline — they invoke nvcc for CUDA sources and your system C++ compiler for host code, then link the result into a shared library that pybind11 exposes as a Python module.
The JIT path is convenient while developing: call load() once and the compiled .so is cached in ~/.cache/torch_extensions. For deployment, setuptools is cleaner because the extension can be built as a wheel with pinned CUDA and PyTorch ABI versions.
The binding layer uses pybind11 to marshal torch::Tensor objects between Python and C++. This is a zero-copy operation — the C++ side receives a pointer to the same GPU memory that the Python tensor owns. No data moves. The only contract is that you must not resize or free the tensor on the C++ side while Python still holds a reference.
03 // WRITING THE CUDA KERNEL
Let's build something concrete: a fused attention-weighted NMS (Non-Maximum Suppression) kernel for a single-shot object detector. In the standard PyTorch pipeline, this requires three separate kernel launches — attention score computation, element-wise multiplication, and IoU-based suppression. Fusing them into a single kernel eliminates two global memory round-trips and two kernel launch overheads (~5μs each on Jetson).
__global__ void fused_attn_nms_kernel(
const float* __restrict__ boxes, // [N, 4] xyxy format
const float* __restrict__ scores, // [N]
const float* __restrict__ attn_map, // [N, N] pairwise attention
int* __restrict__ keep_mask, // [N] output: 1=keep, 0=suppress
const int N,
const float iou_threshold
)Each thread handles one candidate box. We use 1D grid indexing with 256 threads per block. Shared memory stores the current "pivot" box coordinates for IoU computation, avoiding redundant global memory reads within a warp.
const int threads = 256;
const int blocks = (N + threads - 1) / threads;
const int shared_mem = threads * sizeof(float) * 4; // xyxy per thread
fused_attn_nms_kernel<<<blocks, threads, shared_mem>>>(
boxes_ptr, scores_ptr, attn_ptr, mask_ptr, N, 0.45f
);Inside the kernel, each thread loads its box into shared memory, then iterates over higher-scoring boxes. The IoU computation uses the standard intersection-over-union formula, but all box coordinates are read from __shared__memory rather than global DRAM — a 10× bandwidth advantage on the TX2's memory hierarchy.
__shared__ float s_boxes[256 * 4];
int tid = threadIdx.x;
int gid = blockIdx.x * blockDim.x + tid;
if (gid < N) {
s_boxes[tid*4+0] = boxes[gid*4+0]; // x1
s_boxes[tid*4+1] = boxes[gid*4+1]; // y1
s_boxes[tid*4+2] = boxes[gid*4+2]; // x2
s_boxes[tid*4+3] = boxes[gid*4+3]; // y2
}
__syncthreads();04 // MEMORY MANAGEMENT
The single most common performance bug in PyTorch inference code is implicit memory transfers. Every .cpu() call, every .item() on a GPU tensor, every Python-side conditional that touches tensor data — these all trigger synchronous cudaMemcpy calls that stall the GPU pipeline. In a custom extension, you control the memory lifecycle explicitly.
Contiguous vs strided. PyTorch tensors can have non-contiguous memory layouts after operations like transpose() or narrow(). CUDA kernels require contiguous memory. Always call .contiguous() before passing tensors to your extension — or better, assert contiguity inside the C++ binding and throw a clear error.
Pinned memory. For CPU-to-GPU transfers (e.g., ingesting camera frames), allocate host tensors with pin_memory=True. Pinned (page-locked) memory enables DMA transfers that bypass the CPU cache, achieving the full PCIe bandwidth:
CUDA streams. Overlapping compute and copy helps throughput. Use separate streams for copy and compute, then synchronize at pipeline boundaries, not after every small operation:
at::cuda::CUDAStream compute_stream = at::cuda::getStreamFromPool();
at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool();
// Launch async copy on copy_stream
frame_gpu.copy_(frame_pinned, /*non_blocking=*/true);
// Record event on copy_stream, wait on compute_stream
auto event = copy_stream.record_event();
compute_stream.wait_event(event);
// Launch kernel on compute_stream — zero stall05 // TORCHSCRIPT INTEGRATION
A custom CUDA extension is useful, but Python still owns orchestration unless the model is exported. To run from C++, register the op with TorchScript and load the model through LibTorch.
Registration binds your C++ function into the TorchScript operator registry under a custom namespace. Once registered, the op is callable from both Python (torch.ops.my_ops.fused_attn_nms) and inside traced/scripted models.
static auto registry = torch::RegisterOperators()
.op("my_ops::fused_attn_nms",
&fused_attn_nms_forward);
// Now callable as:
// torch.ops.my_ops.fused_attn_nms(boxes, scores, attn, iou_thresh)The serialized model.pt file contains the full computation graph as TorchScript IR, plus the serialized weights. Custom ops are referenced by their registered name — the LibTorch runtime resolves them at load time from the linked .so. This means the same model file works on x86 and aarch64, as long as the extension is compiled for the target architecture.
06 // BENCHMARKING
Profiling CUDA code requires GPU-side timing. Wall-clock time.time() in Python measures CPU time plus synchronization stalls — not actual kernel execution time. Use torch.cuda.Event with enable_timing=True to measure elapsed GPU time between two events, and NVIDIA Nsight Systems for full pipeline analysis with kernel overlap visualization.
We benchmarked three implementations of the attention-weighted NMS pipeline on both a Titan V (Volta, data center) and a Jetson TX2 (Pascal, embedded). Input: 2000 candidate boxes from a single-shot detector.
Sequential IoU in Python, per-box attention lookup
14.2 ms (Titan V)
87.5 ms (TX2)
Three separate CUDA kernels, implicit syncs
2.1 ms (Titan V)
11.3 ms (TX2)
Single kernel launch, shared memory IoU
0.34 ms (Titan V)
1.8 ms (TX2)
The fused kernel is 6.2× faster than vectorized PyTorch and 48.6× faster than the Python loop on the TX2. The speedup comes from three sources: eliminating kernel launch overhead (3 launches → 1), avoiding intermediate global memory writes (the attention-weighted scores never leave registers), and leveraging shared memory for the IoU computation.
Roofline analysis confirms the kernel is compute-bound on Titan V (hitting 78% of peak FLOPS) and memory-bound on TX2 (saturating the 59.7 GB/s LPDDR4 bandwidth). This informs optimization strategy: on the TX2, further gains come from reducing memory traffic via quantization; on Volta, from increasing arithmetic intensity via loop unrolling.
07 // DEPLOYMENT
Getting a custom extension from a developer workstation to an embedded device in a car or robot involves several non-obvious steps. The extension must be compiled for the target architecture with the exact CUDA toolkit and PyTorch ABI that the device runs. Cross-compilation is the practical approach.
Use torch.utils.cpp_extension.CUDAExtension in your setup.py. Pin torch and cuda versions in the wheel metadata. Build separate wheels for each CUDA arch: sm_62 (TX2), sm_72 (Xavier), sm_70 (Volta).
NVIDIA provides the JetPack SDK with aarch64 cross-compilation toolchains. Set TORCH_CUDA_ARCH_LIST="6.2" and point CMake at the aarch64 sysroot. The compiled .so links against the Jetson's local libcudart at runtime.
For ops that TensorRT natively supports (convolutions, BatchNorm), let TensorRT handle optimization. For custom ops like our fused NMS, register them as TensorRT plugins. The plugin API receives GPU pointers and returns GPU pointers — the same pattern as our CUDA kernel, just wrapped in TensorRT's IPluginV2 interface.
Base image: nvcr.io/nvidia/l4t-pytorch:r32.2-pth1.0. Install your wheel with pip install. Mount /dev/nvhost-* and /dev/nvmap for GPU access. The container runs identically on every Jetson with the same JetPack version — reproducible deployment.
08 // CONCLUSION
The pattern is straightforward: prototype in Python, profile the slow operators, fuse the hot path in CUDA, then export through TorchScript + LibTorch when Python is no longer acceptable on the target.
In 2018, this still took real CUDA knowledge. But once the kernel was correct, the result was simple to call and much easier to fit into a fixed latency budget.
PyTorch C++ extension tutorial: pytorch.org/tutorials/advanced/cpp_extension — TorchScript documentation: pytorch.org/docs/stable/jit — NVIDIA Jetson developer zone: developer.nvidia.com/embedded-computing.