Kernels Save Time by Moving Less Data

A fused GeLU runs in 1.1 ms. The same formula written as raw PyTorch math takes 8.1 ms. The hand-written CUDA kernel lands at about 1.8 ms, behind the library it set out to beat. Measure first, then fuse the memory-bound chains, and write a kernel only when the profiler names one.

Measure first or you are guessing

Three traps make a GPU benchmark lie, and each one has a fix that costs nothing.

• The first CUDA call compiles or loads code, so trial one reports startup instead of steady state.
• Python returns as soon as it queues the work, so an unsynchronized timer measures the queue.
• Thermals, operating system activity, and other processes move the number from run to run.

A benchmark that reports the queue makes a slow kernel look free. You then optimize the operation next to it.

One helper answers all three in five steps. Run warmup iterations, synchronize, time several trials, synchronize after each trial, and report the mean. A timing taken any other way is a rumor about the kernel.

The profiler tells you which operation owns the time

The PyTorch profiler splits a workload into Python calls, aten operations, CUDA kernels, and launch overhead. It tells you whether the GPU or the dispatch layer owns the time.

Profile A + B and the shape is clear. The CUDA elementwise kernel finishes in microseconds while aten::add and cudaLaunchKernel own most of the CPU time. Nothing about that operation is worth a rewrite.

Profile torch.cdist and the answer moves. The call breaks into matmuls, elementwise work such as pow and subtract, plus reductions and a square root. The matmuls take more than 70 percent of the time, so the gain lives there and nowhere else.

Nsight Systems puts the CPU threads and the GPU timeline on one screen. NVTX ranges such as forward and backward mark the phases. Two behaviors show up there that the operator table hides.

The CPU runs far ahead of the GPU and queues kernels the GPU executes later. A print of a GPU scalar forces a synchronize, and the CPU then waits for the GPU to drain. Frequent prints and frequent .item() calls in a training loop pay that stall on every iteration.

Python slowness usually does not matter when the GPU owns the time. The exceptions are large pure-Python loops and frequent syncs.

The operator table also shows which of your calls are already fused. A fused built-in appears as one kernel and a naive chain appears as many.

Fusion is where the free speed lives

Fusion combines several operations into one kernel. The intermediate values stay in registers instead of going to global memory. That cuts the traffic and the launch overhead together.

A fast kernel obeys one pattern. Load from DRAM once, reuse the values in registers or shared memory, and write to DRAM once. Every extra trip to global memory is time the arithmetic never sees.

The naive GeLU launches a separate kernel for each multiply, add, tanh, and cube on the tensor. Every kernel reads the tensor from global memory and writes it back. The built-in F.gelu does all of that math inside one kernel with the intermediates held in registers.

The GPU note gave the rule that explains the gap. Arithmetic intensity is FLOPs for each byte moved, and compute grows faster than bandwidth in every hardware generation. An elementwise chain carries almost no arithmetic for each byte, so it waits on memory the whole way.

Fusion buys its speed by cutting trips, and the arithmetic stays exactly the same. An operation whose arithmetic intensity you cannot state is an operation you cannot correct.

The same rule sorts the rest of your model. Large, well-tiled matmuls run compute-bound. Elementwise operations, small matmuls, and work with a poor access pattern run memory-bound. Fusion pays on all three.

A custom kernel in CUDA C++ or Triton runs the same GeLU at about 1.8 ms. torch.compile on the naive code runs it at about 1.47 ms and generates the fused Triton kernel itself. Someone tuned the library kernel for these shapes on this hardware. A first hand-written version starts with none of that tuning.

Triton buys most of CUDA's speed in Python

The CUDA version makes you write both halves by hand. Mark the kernel __global__, compute the global index from the block and thread ids, then bounds-check and store. Set CUDA_LAUNCH_BLOCKING=1 while you debug so the error surfaces at the line that caused it.

Triton moves that work into Python at the block level. You write one program instance for each block, and Triton handles coalescing and the low-level details.

pid = tl.program_id(0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < num_elements
x = tl.load(x_ptr + offsets, mask=mask, other=0)
# Compute GeLU on x as vector math
tl.store(y_ptr + offsets, y, mask=mask)

The Triton GeLU runs at about 1.84 ms, which sits next to the CUDA C++ kernel. The code reads like vectorized Python math, which makes it cheaper to change.

Softmax is harder than GeLU because it needs a row-wise reduction. Subtract the maximum, exponentiate, sum, and divide. The simple design gives one row to one block, with as many blocks as rows. BLOCK_SIZE is a power of two, no smaller than the number of columns.

On a large matrix that design loses. The naive PyTorch softmax takes about 3.7 s and the built-in takes about 1.5 s. torch.compile takes about 1.3 s and the simple Triton kernel about 1.9 s. The compiler searches tilings and fusion choices that a first draft in Triton does not.

torch.compile traces the code into a graph, fuses the elementwise chains, and picks kernels for your shapes and hardware. Most of the fusion gain arrives without a kernel written by you.

Write a custom kernel only when the profiler names it

The accounting note ran one triage on a budget, and the same triage runs on a profile. Spend the effort where the share of the time is large enough to repay it.

A hand-written kernel earns its place when the libraries cover nothing like the pattern. The other case is a tight inner loop the profiler flags, where the compiler misses hardware behavior that matters.

The numerical check is the part people drop. A kernel that is faster and wrong costs more than the operation it replaced. Run the simple baseline next to the kernel on the same input, and compare time and output.

The Builder Test

Profile one training step before you change a line. Write down the top three operations and the share of the time each one owns. The operator breakdown answers which operation, and the timeline answers whether the GPU waited.

If that list surprises you, the optimization you had planned was a guess. Start at the top of the list instead. Re-profile after the change, so the next decision has a number under it. I rank operations by their share of the time before I open an editor.

What Carries

A guess about the hotspot costs more than the measurement that names it. Memory, compute, communication, and data each call for a different tool.

Name the constraint first, then pick the tool that attacks it. One GPU is busy now, so the next constraint appears when the model needs many of them.