Skip to content

Writing Custom CUDA Kernels

Overview

When the GPU is idle waiting on the CPU (launch-bound) or no PyTorch op expresses your algorithm, write a CUDA kernel— a function that runs in parallel on GPU threads— and expose it through a CUDAExtension. The win: fusion (Ch 04-05), custom algorithms (flash-style attention, quantization), and elimination of host-device round-trips.

  • Kernels: __global__ functions; grid/block launch config; threadIdx/blockIdx indexing.
  • Memory: registers → shared memory → global memory; coalescing matters.
  • Compile with torch.utils.cpp_extension.load(..., with_cuda=True) or CUDAExtension.
  • Wrap in torch.autograd.Function for gradients.

Rule: correctness first, then profile with ncu. A naive kernel is often slower than PyTorch's own.

-

Minimal CUDA Extension

// mycu.cpp
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void add_kernel(const float* a, const float* b, float* out, int n) {
 int i = blockIdx.x * blockDim.x + threadIdx.x;
 if (i < n) out[i] = a[i] + b[i];
}

torch::Tensor cuda_add(torch::Tensor a, torch::Tensor b) {
 auto out = torch::empty_like(a);
 int n = a.numel();
 int threads = 256;
 int blocks = (n + threads - 1) / threads;
 add_kernel<<<blocks, threads>>>(a.data_ptr<float>(), b.data_ptr<float>(),
 out.data_ptr<float>(), n);
 return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 m.def("add", &cuda_add);
}

from torch.utils.cpp_extension import load
cu = load(name="mycu", sources=["mycu.cpp"], with_cuda=True)

x = torch.randn(1024, device='cuda')
y = torch.randn(1024, device='cuda')
print(cu.add(x, y).shape)

-

Autograd for a Custom Kernel

import torch, mycu

class CudaAddFn(torch.autograd.Function):
 @staticmethod
 def forward(ctx, a, b):
 return mycu.add(a.contiguous(), b.contiguous())
 @staticmethod
 def backward(ctx, g):
 return g, g # d/da (a+b) = 1, d/db = 1 (elementwise)

x = torch.randn(8, device='cuda', requires_grad=True)
CudaAddFn.apply(x, torch.randn(8, device='cuda')).sum().backward()
print("grad ok:", x.grad.sum().item())

Always .contiguous() before data_ptr— strided tensors will corrupt reads.

-

Launch Configuration Rules

threads = 256 # 1D typical; power-of-two
blocks = (n + threads - 1) // threads
kernel<<<blocks, threads>>>(...)
  • Coalescing: consecutive threads → consecutive addresses. The #1 perf lever.
  • Grid-stride loop: for huge tensors, loop within kernel to avoid blocks > 2^31.
__global__ void add_stride(const float* a, const float* b, float* o, int n) {
 int i = blockIdx.x * blockDim.x + threadIdx.x;
 int stride = gridDim.x * blockDim.x;
 for (; i < n; i += stride)
 o[i] = a[i] + b[i];
}
  • Shared memory for reductions/tiling: __shared__ float smem[256];.
  • Non-blocking: after launch, code proceeds; call torch.cuda.synchronize() only in tests.

When a Custom Kernel Wins

Situation Verdict
Elementwise chains (fusion) Triton/Inductor already fuses → skip hand CUDA
Custom math (quant, top-k, search) hand CUDA wins
Attention/Flash patterns use F.scaled_dot_product_attention first
Memory-bound reduction over huge arrays hand CUDA + vectorized loads
Learning/experiments torch.compile + Triton first

Modern guidance: reach for Triton inside torch.compile or write Triton kernels directly (triton python) before raw CUDA, unless you need low-level control (shared memory tuning, atomics).

-

Debugging Kernels

  1. Check launch bounds valid (blocks/threads positive).
  2. Use cudaGetLastError() after launch in debug builds.
  3. compute-sanitizer --tool memcheck python test.py for OOB/race.
  4. Reference-check vs a CPU/python implementation on random small inputs.
  5. Add -lineinfo -G for source-level ncu/nsight attribution.
def reference_check(fn_cuda, fn_py, shape=(8, 8), seed=0):
 torch.manual_seed(seed)
 a = torch.randn(*shape, device='cuda')
 b = torch.randn(*shape, device='cuda')
 assert torch.allclose(fn_cuda(a, b), fn_py(a, b), atol=1e-5)
 print("reference check passed ")

-

Key Takeaways

  • Kernels = __global__ + launch config + data_ptr I/O; wrap with autograd Function.
  • Coalescing + grid-stride loops are the two big perf levers.
  • .contiguous() inputs before passing pointers; check bounds.
  • Don't hand-write what Inductor/Triton/SDPA already fuse— verify need via profile.
  • Debug with compute-sanitizer + reference checks; profile with ncu.

-