GPU kernel engineering

We start with vendor libraries and compiler output, and write custom kernels only where profiling shows a bottleneck.

Most of the time in a well-served model is spent in kernels someone else has already tuned: cuBLAS and CUTLASS for matrix multiplies, the attention kernels shipped with the serving engine, and whatever torch.compile generates for the rest. A custom kernel is a maintenance cost that has to pay for itself, so it is written only after a profiler has shown which operation dominates and by how much.

Why kernels matter

A GPU has two limits: how fast it can do arithmetic, and how fast it can move data from its main memory (HBM). Which limit binds depends on the operation's arithmetic intensity, the number of floating-point operations it performs per byte it moves.

• If intensity is below the GPU's ratio of compute to bandwidth, the operation is memory-bound: the arithmetic units wait for data.
• If it is above, the operation is compute-bound: memory keeps up and the arithmetic is the limit.

A roofline plots attainable performance against intensity. It rises along a sloped line (bandwidth × intensity) until it meets a flat ceiling (peak compute). Reading where a kernel sits on that plot says what kind of optimisation can help it.

In a language model:

For memory-bound work, the main lever is kernel fusion: doing several operations in one kernel so intermediate results stay in registers or on-chip memory instead of making a round trip to HBM.

CUDA fundamentals

• The hierarchy. Threads are grouped into warps of 32 that execute in lockstep, warps into thread blocks, and blocks are scheduled onto streaming multiprocessors (SMs). A block's threads can share fast on-chip shared memory and synchronise with each other.
• Memory coalescing. When the 32 threads of a warp read 32 consecutive addresses, the hardware serves them in a few wide transactions. Scattered addresses cost many more transactions for the same data.
• Shared-memory banks. Shared memory is divided into 32 banks. If several threads in a warp hit different addresses in the same bank, the accesses are serialised. Padding a tile's row length by one element is the classic fix.
• Occupancy. The number of warps resident on an SM, limited by registers and shared memory per block. More warps hide memory latency, but more is not always faster: a kernel that uses many registers per thread to keep data on-chip can beat one tuned for occupancy.
• Tensor cores. Dedicated units for small matrix multiplies at reduced precision. Hopper added asynchronous data movement (the Tensor Memory Accelerator) and warpgroup-level matrix instructions, so loading the next tile and computing on the current one overlap. Blackwell adds tensor memory, a dedicated on-chip store for matrix operands and accumulators.
• Libraries and DSLs. CUTLASS provides C++ templates for high-performance matrix kernels; the CuTe DSL exposes its layout abstractions from Python; cuTile Python is NVIDIA's tile-based Python programming model.

Triton and compilers

• Triton lets a kernel be written at the level of blocks rather than individual threads. The programmer says what each program instance does to a block of data; the compiler handles coalescing, shared-memory staging and much of the scheduling.
• Autotuning. triton.autotune tries a set of configurations (block sizes, number of warps, pipeline stages) and caches the fastest per input shape.
• TorchInductor. The default backend of torch.compile generates Triton kernels for fused elementwise and reduction work automatically, which is often all the fusion a model needs.
• Helion. A higher-level PyTorch kernel language that compiles to Triton, currently in beta.
• Where Triton gives less control. Exact warp specialisation, register allocation and some Hopper and Blackwell features are easier to reach from CUDA, CUTLASS or the CuTe DSL. Peak-performance attention and GEMM kernels on the newest hardware are usually written there.

Attention and matrix kernels

• The FlashAttention line. FlashAttention computes attention block by block without ever writing the full score matrix to HBM, using the online softmax shown in the worked example. Later versions improved parallelism and work partitioning (FlashAttention-2), used Hopper's asynchronous features and FP8 (FlashAttention-3), and FlashAttention-4 targets Hopper and Blackwell.
• FlexAttention. A PyTorch API in which variations of attention (causal masks, sliding windows, document masks, score modifications) are written as small Python functions and compiled into a fused kernel, instead of writing a new kernel for each.
• Paged-attention kernels. Attention that reads the KV cache from non-contiguous blocks, matching how serving engines store it.
• Mixture-of-experts. Tokens routed to different experts form many small matrix multiplies of uneven size. Grouped GEMMs run them in one launch, and fused MoE kernels combine routing, the expert computation and the weighted sum.
• Low-precision GEMMs. FP8 and NVFP4 matrix multiplies need per-block scale factors applied inside the kernel. DeepGEMM is one published example of FP8 GEMMs with fine-grained scaling.

Portability

• ThunderKittens is a small library of tile primitives for writing fast kernels; HipKittens ports the approach to AMD GPUs.
• TileLang is a tile-level language that compiles to several GPU backends.
• Mojo is a language from Modular aimed at writing portable high-performance kernels.
• ROCm and HIP. HIP mirrors much of the CUDA API, and tools exist to translate CUDA source to HIP. It is not a drop-in replacement: AMD data-centre GPUs schedule 64-thread wavefronts rather than 32-thread warps, have different memory hierarchies and matrix instructions, and a kernel tuned for one vendor usually needs retuning for the other.

Profiling and verification

• Find the bottleneck first. Nsight Systems shows a timeline of kernels, memory copies and CPU activity, which reveals gaps and launch overhead. The PyTorch profiler attributes GPU time to model operations. Nsight Compute then examines one kernel in depth, including where it sits on the roofline.
• Correctness before speed. A new kernel is compared against a higher-precision reference with stated tolerances, on shapes that include awkward sizes, masked inputs and non-contiguous tensors (see the worked example).
• Benchmark hygiene. Warm up before timing, repeat enough times to report a distribution, fix GPU clocks where possible, synchronise before reading timers, and state the exact shapes, dtypes and hardware with every number.

Worked examples

Arithmetic intensity of decode

Illustrative example. This explains a technique; it does not describe any client system.

A matrix-vector product at batch size b, with a d × k weight matrix in 16-bit, does about 2·b·d·k floating-point operations while reading about 2·d·k bytes of weights:

intensity ≈ (2 · b · d · k) / (2 · d · k) = b   operations per byte

An H100's dense 16-bit tensor throughput divided by its memory bandwidth is in the hundreds of operations per byte. NVIDIA's headline tensor figures are quoted with structured sparsity; dense throughput is half. So decode at batch size 1 has an intensity of about 1 and sits far down the memory-bound slope of the roofline, while large prefill matrix multiplies sit under the flat compute roof.

What fusion buys

Illustrative example. This explains a technique; it does not describe any client system.

An unfused residual add followed by RMSNorm launches two kernels:

kernel 1: read x, read residual → write h = x + residual     (HBM write)
kernel 2: read h → compute RMS → write y = h / rms(h) · g     (HBM read)

A fused kernel computes h, keeps it in registers, and writes only y (and h, if the next layer needs it as its residual). That removes one full read and one write of the activation tensor. For memory-bound elementwise work, runtime is roughly bytes moved ÷ bandwidth, so the saving tracks the bytes removed.

Online softmax, the idea behind FlashAttention

Illustrative example. This explains a technique; it does not describe any client system.

Softmax needs the maximum and the sum over a whole row, which seems to require the whole row at once. The online form keeps a running maximum m and running sum l per row, and updates them one block of scores s at a time:

m_new = max(m, max(s))
l_new = e^(m − m_new) · l + Σ e^(s − m_new)
o_new = e^(m − m_new) · o + Σ e^(s − m_new) · v

After the last block, the output is o / l. The result is exact, not an approximation. Because the full score matrix is never written to HBM, memory traffic grows linearly with sequence length rather than quadratically.

A Triton kernel, in outline

Illustrative example. This explains a technique; it does not describe any client system.

A row-wise fused softmax: one program per row, a mask for row lengths that are not a power of two, the row maximum subtracted for numerical stability.

import triton
import triton.language as tl

@triton.autotune(
    configs=[triton.Config({}, num_warps=w, num_stages=s) for w in (2, 4, 8) for s in (1, 2)],
    key=["n_cols"],
)
@triton.jit
def softmax_kernel(out_ptr, in_ptr, in_stride, out_stride, n_cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK)
    mask = offs < n_cols
    x = tl.load(in_ptr + row * in_stride + offs, mask=mask, other=-float("inf"))
    x = x - tl.max(x, axis=0)
    num = tl.exp(x)
    tl.store(out_ptr + row * out_stride + offs, num / tl.sum(num, axis=0), mask=mask)

# launch: BLOCK = triton.next_power_of_2(n_cols); grid = (n_rows,)

Here the block covers a whole row, so the autotuner searches warp count and pipeline stages per row length. For kernels that tile a larger problem, block sizes are part of the search as well.

Verifying a kernel

Illustrative example. This explains a technique; it does not describe any client system.

Compare against a reference computed in higher precision with torch.testing.assert_close, which passes when

|actual − expected| ≤ atol + rtol · |expected|

Test shapes that are not multiples of the block size, masked and non-contiguous inputs, very large and very small values, and repeated runs to check determinism. Then benchmark after warm-up, over enough repetitions, at fixed clocks, with the shapes stated.

Sources

• NVIDIA: GPU performance background for deep learning (opens in a new tab)
• CUDA C++ Best Practices Guide (opens in a new tab)
• NVIDIA: roofline analysis in Nsight Compute (opens in a new tab)
• CUTLASS CuTe DSL (opens in a new tab), cuTile Python (opens in a new tab)
• Triton programming guide (opens in a new tab), Helion (opens in a new tab), FlexAttention (opens in a new tab)
• Dao et al. (2022), FlashAttention (opens in a new tab); FlashAttention repository (opens in a new tab); FlashAttention-4 (opens in a new tab)
• DeepGEMM (opens in a new tab)
• ThunderKittens (opens in a new tab), HipKittens (opens in a new tab), TileLang (opens in a new tab)
• AMD: what is HIP (opens in a new tab)
• PyTorch profiler recipe (opens in a new tab), torch.testing default tolerances (opens in a new tab)
• NVIDIA H100 specifications (opens in a new tab)