Kernel Fusion, Softmax, and Attention

Consider a PyTorch expression that adds a bias to 67 million activation values and then applies GELU. Eager execution can launch one CUDA kernel to write the biased values to GPU high-bandwidth memory and another kernel to read those values, apply GELU, and write the final output. Nsight Systems shows two launch rectangles, while Nsight Compute reports a full-tensor write and read for an array that exists only between the two operations. A single CUDA or Triton kernel can calculate both expressions while each value remains in a register.

Kernel fusion is the established name for combining such producer and consumer operations into one GPU kernel. Fusion can remove launches and intermediate memory traffic, but the combined kernel may keep more values alive, allocate more registers, lose a library GEMM, or duplicate work needed by another consumer. The chapter measures the removed bytes and the new resource costs before treating fusion as an improvement.

Transformer attention supplies a larger real example. A direct implementation computes a score for every query-key pair, writes the score matrix to high-bandwidth memory, reads it for softmax, writes probabilities, and reads them again to combine value vectors. For a batch of eight sequences, 32 heads, and sequence length 4096, a single score tensor contains more than four billion elements. A profiler or memory snapshot can therefore show gigabytes of intermediate storage even though the model ultimately needs only the normalized weighted value for each query.

In the profiles, a central processing unit (CPU) launches kernels on a graphics processing unit (GPU), and large tensors reside in high-bandwidth memory (HBM). Attention contains matrix multiplication (GEMM) for query-key scores and probability-value products. During inference, later chapters retain key-value (KV) vectors from earlier tokens so one-token decode does not recompute the complete prompt.

The chapter first derives a stable running softmax from small rows whose numbers can be checked by hand. The chapter then follows query, key, and value tiles through on-chip storage to explain FlashAttention, an exact attention algorithm that avoids materializing the complete score and probability matrices. Separate lessons show why training, prompt prefill, and one-token decode use different attention schedules even though all three implement the same mathematical attention equation.