ATTENTION KERNELS

The roofline's "bandwidth" is HBM bandwidth. On an H100 there are several tiers above it:

Shared memory is about 6× the bandwidth of HBM and 16× lower latency, but there is only 228 KB of it per SM. That constraint — fast but tiny — is what shapes every high-performance attention kernel.

Now count what naive attention costs in HBM traffic. For one head, sequence length N, head dimension d, in fp16:

1. read Q, K                       2 * N * d * 2 bytes
2. write S = QK^T                  N^2 * 2 bytes
3. read S                          N^2 * 2 bytes
4. write P = softmax(S)            N^2 * 2 bytes
5. read P, read V                  N^2 * 2 + N * d * 2 bytes
6. write O                         N * d * 2 bytes
                                   ---------------------
                                   ~4 N^2 * 2 bytes of traffic for the score matrix

At N = 8192, d = 128: the N × d terms are about 2 MB each, while each N² term is 134 MB. The score matrix dominates by roughly 60×, and you move it four times.

For a full model — 32 heads, 32 layers — that is 4 × 134 MB × 32 × 32 = 549 GB of HBM traffic for one 8k-token prefill, just shuttling intermediate scores. At 3.35 TB/s that is 164 ms of pure memory movement, and none of it is arithmetic.

The FLOPs, meanwhile, are 4 N² d per head — real work that must happen. The problem is not the arithmetic. It is that the intermediate never needed to exist in HBM.

This is a general lesson worth extracting: the roofline says "minimize bytes from DRAM", and one of the most powerful ways to do that is to restructure an algorithm so intermediates stay in a faster tier. That is what kernel fusion is, and attention is its most valuable application.

The obstacle to tiling attention is softmax. It needs a denominator that sums over every score in the row, so it appears to require the whole row before producing any output.

The resolution is to compute it incrementally and correct as you go. This is worth deriving, because the exactness of FlashAttention rests entirely on it.

Standard safe softmax over a vector x:

m = max(x)                       subtract the max so exp() cannot overflow
l = sum_i exp(x_i - m)
y_i = exp(x_i - m) / l

Now suppose x arrives in two blocks, x1 then x2. Process x1:

m1 = max(x1)
l1 = sum exp(x1 - m1)

Then x2 arrives and may contain a larger value. The new global max is m = max(m1, m2), and the correction is one multiplication:

l = l1 * exp(m1 - m)  +  l2 * exp(m2 - m)

The factor exp(m1 - m) rescales the old partial sum from base m1 to base m. Since exp(a - m1) × exp(m1 - m) = exp(a - m) exactly, nothing is approximated. It is the same number, re-expressed.

The same correction applies to the accumulated output. If you have a partial output O1 computed against the old normalization, rescale it the same way before adding the new block's contribution:

O = O1 * (l1 * exp(m1 - m) / l)  +  (new block's contribution)

In FlashAttention this is arranged so the division by l happens once at the end, and the inner loop just tracks unnormalized accumulators along with m and l.

Why exactness matters so much. There is a long line of approximate attention methods — Linformer, Performer, Reformer, sparse patterns — that reduce complexity by changing what is computed. They all face the same adoption problem: you cannot drop them into a trained model without retraining, and you cannot be sure what capability you gave up. FlashAttention computes bit-comparable attention (modulo floating-point reassociation) and is therefore a pure drop-in. You can apply it to any existing checkpoint with no evaluation required.

That is the whole reason it won. Not the speedup — a 2–4× speedup with an asterisk about quality is a much harder sell than a 2–4× speedup with no asterisk at all.

Put the pieces together. The algorithm splits Q, K and V into blocks sized so a working set fits in shared memory, then loops:

for each block of K, V:                     # loaded into SRAM once
    for each block of Q:                    # loaded into SRAM
        S_ij  = Q_i @ K_j^T                 # in SRAM, small
        m_new = max(m_i, rowmax(S_ij))
        P_ij  = exp(S_ij - m_new)           # in SRAM
        l_new = l_i * exp(m_i - m_new) + rowsum(P_ij)
        O_i   = O_i * (l_i * exp(m_i - m_new) / l_new) + (P_ij @ V_j) / l_new
        m_i, l_i = m_new, l_new
    # O_i, m_i, l_i live in registers/SRAM across the inner loop
write O                                     # the only large HBM write

The [N, N] score matrix is never assembled. Only [block, block] tiles exist, in shared memory, and they are discarded as soon as they are consumed.

The traffic analysis: HBM accesses drop from O(N² + Nd) to O(N²d²/M), where M is the SRAM size. With d = 128 and M ≈ 100 KB, the factor d²/M is well under 1, so the reduction is substantial — the paper reports 9× fewer HBM accesses on GPT-2 and 2–4× end-to-end speedups.

Memory goes from quadratic to linear in sequence length, because you only ever hold blocks. This is what made long-context training practical at all.

The trade is recomputation. During the backward pass you need the attention probabilities, which you did not store. FlashAttention recomputes them from Q, K, V on the fly. That is more FLOPs than the naive version — and it is still faster, because the FLOPs were never the bottleneck. It is the clearest possible demonstration of this course's thesis: trading arithmetic for bandwidth is a winning trade, even at a large arithmetic premium.

FlashAttention-2 (2023) reworked the same algorithm around GPU scheduling rather than the memory hierarchy: reduce non-matmul FLOPs (rescaling operations are much slower per FLOP than tensor-core matmuls on modern GPUs), swap the loop order so the query block is outer, and parallelize over sequence length as well as batch and heads. Roughly another 2×, reaching 50–73% of theoretical peak matmul throughput.

FlashAttention-3 (2024) targets Hopper specifically: warp specialization (separate warps produce and consume data), asynchronous TMA copies overlapped with computation, interleaving the matmul and softmax so the non-tensor-core softmax work hides behind tensor-core work, and FP8 support. The trend across all three is instructive — the first was an algorithmic insight, the next two were about exploiting specific hardware asynchrony. Extracting the last factor of two now requires knowing the machine intimately.

FlashAttention parallelizes over batch, heads, and query blocks. At decode time the query length is 1, so the third source of parallelism disappears.

Concretely: batch 8, 32 heads, query length 1, KV length 8192. You have 8 × 32 = 256 independent work units for a GPU with 132 SMs — and each unit must sequentially scan 8,192 keys. Occupancy is poor and the long serial scan dominates. Meanwhile at batch 1 you have only 32 work units for 132 SMs: the GPU is more than 75% idle.

The fix is FlashDecoding, and it is a split-K reduction. Partition the KV sequence into chunks, compute a partial attention result for each chunk in parallel, then combine.

1. split the 8192 keys into, say, 16 chunks of 512
2. each chunk computes, independently and in parallel:
      - partial output O_c
      - its running max m_c and sum l_c
3. combine the 16 partials using exactly the online-softmax rescaling
   from earlier -- the same algebra, applied across chunks instead of
   within a loop

Parallelism goes from batch × heads to batch × heads × chunks — a 16× increase in this example, which is the difference between an idle GPU and a busy one. Reported speedups for long-context decode are substantial, up to several times for large KV lengths.

Note how neatly the online softmax pays off twice. It was introduced to allow sequential tiling within one kernel; the identical algebra allows parallel partial results to be combined afterwards. That is a sign the abstraction was the right one.

Three further decode-specific realities:

GQA changes the arithmetic. With 8 KV heads serving 32 query heads, the 4 query heads sharing a KV head can read it once from HBM and use it four times. A GQA-aware decode kernel gets 4× the arithmetic intensity on the attention step. A kernel that naively replicates KV heads to match query heads throws that away — and this is a real and common performance bug.

Paged KV changes the memory access pattern. Under PagedAttention the cache is not contiguous, so the kernel must gather through a block table. vLLM's paged attention kernels handle this; a stock FlashAttention kernel cannot read a paged cache directly. The two features have to be co-designed, which is why they arrived together.

Attention at decode is memory-bound no matter what. Each cached key element is read and used in one multiply-accumulate, so intensity is O(1) — and unlike weights, batching does not amortize it, because every sequence has its own cache. No kernel fixes this. The only fixes are storing less (Module 6) or reading it less often.

FlashAttention is the most valuable instance of a general principle: if two operations are performed back to back, do not send the intermediate to HBM.

Consider RMSNorm followed by a linear projection at decode time, batch 1, d_model = 4096, fp16:

unfused:
  RMSNorm kernel:  read 8 KB, ~3 FLOPs/element, write 8 KB
  matmul kernel:   read 8 KB + weights, write output
                   -> the 8 KB intermediate makes a full HBM round trip

fused:
  one kernel:      read 8 KB, normalize in registers, matmul, write output
                   -> intermediate never leaves the chip

8 KB sounds negligible, and per operation it is. The problem is the count. A decode step for a 32-layer model has, per layer: two RMSNorms, RoPE on Q and K, two residual adds, and a SiLU-and- multiply. That is roughly 7 elementwise operations × 32 layers ≈ 224 kernel launches, each with a round trip and each with ~5 µs of launch overhead.

launch overhead:  224 x 5 us          = 1.12 ms
against a theoretical step of         = 4.5 ms
                                        -> 25% of the budget

Plus the memory traffic, which is small per operation but adds up, and — often worse — each tiny kernel fails to saturate the memory system at all, so its effective bandwidth is far below peak.

The standard remedies:

Fusion. Combine adjacent operations into one kernel. Fuse the norm into the following matmul, fuse SiLU and the elementwise multiply in SwiGLU, fuse the residual add into whatever precedes it. Every serving engine ships hand-written fused kernels for the common patterns, and torch.compile generates them automatically for many cases.

CUDA Graphs. Capture the whole sequence of launches once and replay it as a single unit. This does not reduce memory traffic but eliminates nearly all the per-launch CPU cost. For decode — where the shape of every step is identical — it is close to free and is standard practice. It is frequently the single largest fix when a naive PyTorch loop is 3× off the roofline.

Persistent kernels. One kernel that loops internally over layers, avoiding launches entirely. Highest performance, least flexibility, hardest to write.

The reason this module comes after Module 4 is that the roofline tells you whether fusion will help. If you are compute-bound, an extra HBM round trip may hide behind arithmetic and fusion buys little. At decode you are never compute-bound, so every round trip is exposed and every fusion pays.

FlashAttention makes exact attention fast, but it does not make it sub-quadratic — the FLOPs are still O(N²d) and the KV cache still grows linearly. A parallel research line tried to change the asymptotics.

Linear attention. Replace softmax(QK^T)V with a kernel feature map φ so that φ(Q)(φ(K)^T V) can be computed by associativity in O(Nd²) instead of O(N²d). The state is a fixed [d, d] matrix — constant size, no growing cache. The cost is that softmax's sharp, content-addressable selectivity is replaced by something smoother, and exact retrieval suffers.

State space models. Mamba (2023) is the most successful. It maintains a fixed-size recurrent state updated per token, with input-dependent (selective) state transitions, and a parallel scan that makes training efficient. There is no KV cache at all — decode memory is constant in sequence length, which is an enormous inference advantage.

The trouble is a genuine capability gap. Attention can look up any previous token exactly; a fixed-size state cannot store an arbitrary amount of detail. Empirically, pure SSMs underperform on tasks requiring precise recall from long context — copying, retrieval, in-context learning from many examples. This is not an implementation shortfall, it is information-theoretic: a constant-size state cannot losslessly summarize an unbounded history.

Hybrids won. The field's answer was to interleave: mostly SSM or linear-attention layers, with a few full-attention layers to provide exact recall. Jamba, Zamba and Samba are examples, and the approach is now common in production models. You get most of the memory saving with most of the recall.

Sparse attention is the other direction: keep softmax attention but restrict which positions attend to which. Sliding-window attention (Mistral) bounds the cache at the window size. Global plus local patterns give some long-range connectivity cheaply. DeepSeek's Native Sparse Attention (2025) trains the sparsity pattern rather than fixing it in advance, which addresses the main objection to earlier sparse methods — that a hand-designed pattern is a guess about what the model needs.

The honest summary as of now: full attention with FlashAttention kernels remains the default for quality-critical work, hybrids are a real and growing production option where long context matters more than exact recall, and the sub-quadratic line has not displaced attention so much as found a complementary role. Anyone claiming attention is obsolete is ahead of the evidence.