Writing Fast Kernels: Fusion, CUDA, and Triton

From op to kernel

Every PyTorch operation you call — `x.relu()`, `a + b`, a softmax — eventually launches one or more kernels: programs that run on the GPU's thousands of threads. The trouble is that a chain of ops launches a chain of kernels, and each one reads its inputs from HBM and writes its output back. A residual block of five element-wise ops touches memory ten times to do almost no arithmetic — a textbook bandwidth-bound disaster. The fix is to make the GPU do more work between memory trips.

Fusion: the highest-leverage move

Kernel fusion merges several operations into one kernel so that intermediate values live in registers or shared memory and are never written back to HBM. Fuse the five element-wise ops above and you read the input once, compute everything on-chip, and write the result once — memory traffic drops from ten trips to two, and arithmetic intensity rises by the same factor. Fusion is the single most effective bandwidth optimization in deep learning, which is why so much compiler and kernel engineering is, at heart, an effort to fuse more aggressively.

The CUDA execution model

To write a fused kernel by hand you confront the GPU's real machine model, and CUDA kernel optimization is the discipline of mapping your problem onto it well. Threads run in lockstep groups of 32 called warps; warps form thread blocks that share a fast scratchpad of shared memory; blocks are scheduled onto streaming multiprocessors. Three levers dominate performance. Occupancy: keep enough warps resident to hide memory latency. Coalescing: arrange so that the 32 threads in a warp read 32 contiguous addresses, turning many small reads into one wide transaction. Shared-memory tiling: stage a tile of data on-chip and reuse it across threads instead of re-reading HBM.

# tiled matmul, conceptually
for tile_k in range(0, K, TILE):
    load A[:, tile_k:tile_k+TILE] into shared
    load B[tile_k:tile_k+TILE, :] into shared
    __syncthreads()                 # barrier: tile is ready
    accumulate partial products from shared, not HBM
    __syncthreads()

Tiling raises reuse: each loaded tile feeds many multiply-accumulates before it is evicted.

I=\frac{\text{FLOPs}}{\text{bytes}}\approx\frac{2B^{3}}{8B^{2}}=\frac{B}{4}

Why tiling raises reuse, quantified: a B-by-B tile loaded once supplies about B-cubed multiply-accumulates, so arithmetic intensity grows linearly with the tile size B.

Triton: kernels you can actually maintain

Hand-written CUDA is powerful and miserable to maintain. Triton hits a sweet spot: you write Python that operates on whole tiles (blocks of elements), and the compiler handles the per-thread mapping, coalescing, and shared-memory allocation for you. You still think in tiles and intensity, but you stop hand-managing 32-thread warps. In practice a Triton kernel reaches a large fraction of expert CUDA performance in a fraction of the code, which is why it has become the default way researchers ship custom kernels.

@triton.jit
def fused_gelu_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr):
    off = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = off < n
    x = tl.load(x_ptr + off, mask=mask)      # one HBM read
    y = x * 0.5 * (1.0 + tl.erf(x * 0.7071)) # compute on-chip
    tl.store(y_ptr + off, y, mask=mask)      # one HBM write

A whole element-wise activation as one fused, coalesced Triton kernel — two memory trips, no Python loop over threads.

Case study and when to stop

FlashAttention is the canonical payoff. It fuses the entire attention computation — scores, softmax, and the value-weighted sum — into one kernel, tiling over the sequence so the N-by-N score matrix never touches HBM, and uses an online-softmax trick to handle the reduction without materializing the full row. FlashAttention-2 then re-tunes the work partitioning across warps to push occupancy higher still. The result is both faster and more memory-frugal, enabling far longer contexts.

\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V

The entire computation FlashAttention fuses into one tiled kernel — scores, softmax, and the value-weighted sum, never materializing the full N-by-N matrix.