Case study: attention, mapped

Write both as loop nests and the only differences are tile factors and bindings.

Standard, three separate kernels with an $S \times S$ tensor bound to HBM:

S = Q @ K.T          # writes S×S to HBM
P = softmax(S)       # reads S×S, writes S×S
O = P @ V            # reads S×S

Traffic: $O(S^2)$ words written and $O(S^2)$ read, twice. At $S = 8192$ and FP16, that is roughly 400 MB of HBM traffic per head per layer, for $4S^2d$ FLOPs.

FlashAttention, one kernel with the scores bound to SRAM:

for i in range(S / Br):                 # tile the query rows
  load Q_i (Br × d) into SRAM
  m, l, O_i = -inf, 0, 0                # running max, running sum, accumulator
  for j in range(S / Bc):               # tile the keys — the reduction
    load K_j, V_j (Bc × d) into SRAM
    S_ij = Q_i @ K_j.T                  # Br × Bc — lives in SRAM only
    m_new = max(m, rowmax(S_ij))
    P_ij  = exp(S_ij - m_new)
    l     = l * exp(m - m_new) + rowsum(P_ij)     # rescale the running sum
    O_i   = O_i * exp(m - m_new) + P_ij @ V_j     # rescale the accumulator
    m     = m_new
  write O_i / l                         # Br × d to HBM

Compare against Module 5’s four decisions:

Every arithmetic operation on the critical path is the same. The exponentials, the multiplies, the matmuls — identical count. What changed is where one tensor lives and how the loops are cut, which is the definition of a mapping.

Two consequences that follow only from seeing it this way, and are the reason the distinction matters:

• The technique is general. Any operation with a large materialised intermediate consumed immediately is a candidate — fused MLP blocks, fused loss functions with large logits, chunked linear attention, and the same trick applied to convolutions. Once it is “re-bind the intermediate and make the reduction streamable” rather than “a clever attention algorithm”, you go looking for other instances.
• It gets better with hardware, in a predictable direction. SRAM per SM grew from 164 KB on A100 to 228 KB on H100; TMA made the tile loads cheaper; warp specialisation let the loads and the math proceed on different warps. FlashAttention 2 and 3 are largely re-mappings onto those new capabilities rather than new mathematics — which is exactly what you would predict for a result whose content is a mapping.

The one genuinely new ingredient is online softmax, and it is worth being precise that it is not an approximation: the rescaling by $\exp(m - m_{\text{new}})$ makes the streamed result exactly equal to the two-pass result in exact arithmetic, and better-conditioned than naive softmax in floating point.

Derive it first, because the derivation is where the consequences come from.

At decode, each sequence contributes one token. A weight matrix of $P$ parameters must be read in full — $2P$ bytes at FP16 — and each parameter participates in one multiply-accumulate per sequence, so $2P \cdot B$ FLOPs for batch $B$. The intensity is

$\frac{2PB}{2P} = B \text{ FLOP per byte}$

Independent of model size, independent of layer shape. The batch size is the arithmetic intensity. Against an H100’s ridge of ~295, you need ~295 concurrent sequences to be compute-bound; at $B = 1$ you are running at roughly 0.3% of peak, and there is nothing wrong with the kernel.

What follows, and each of these is a real design consequence:

• Token rate at small batch is a pure bandwidth calculation. Time per token $\approx$ model bytes / HBM bandwidth. A 70 B model in FP8 is 70 GB; on 3.35 TB/s that is ~21 ms per token, ~48 tokens/s, and no amount of compute changes it. This one division predicts single-stream decode performance across essentially all hardware, and it is the first calculation to do when someone quotes a token rate.
• Batching is the only lever, so the serving system’s real job is assembling batches. This is why continuous batching — admitting new sequences into a running batch at token granularity instead of waiting for a whole batch to finish — was such a large practical win. It is not a scheduling nicety; it is the mechanism that raises arithmetic intensity.
• KV cache capacity therefore sets throughput. Batch is limited by how many KV caches fit in the memory the weights left over. The arithmetic above — ~7 sequences on an 80 GB H100 with a 70 B FP8 model and GQA at 8 K context — is why every serving optimisation that shrinks the KV cache translates directly into tokens per second. That is a memory-capacity problem wearing a throughput costume.
• Quantizing weights helps twice. Fewer bytes to stream per token (directly faster), and more room for KV cache (bigger batch, higher intensity). Compounding, which is why 4-bit weights can deliver more than 4× on decode.
• Large arrays make it worse. A bigger MAC array raises the ridge point, so the batch needed to saturate it grows, while the KV cache that limits batch has not shrunk. Decode gets structurally harder every generation even as the datasheet improves — Module 2’s warning, arriving with numbers.

And the disaggregation conclusion that the industry reached from this: prefill and decode want opposite things, so run them on separate pools of hardware, transferring the KV cache between them. Prefill wants compute; decode wants bandwidth and capacity. Serving them on one machine means one of the two is always misprovisioned.

Both, and being clear about which part is which is the point.

What they do. Multi-head attention gives every query head its own key and value heads, so the KV cache scales with the head count. Multi-query attention shares one KV head across all queries, cutting the cache by the head count — 64× for a 64-head model. Grouped-query attention interpolates, sharing each KV head across a group of, say, 8 queries. Multi-head latent attention compresses K and V into a low-rank latent that is cached instead, decompressing on use. All three attack the same quantity: bytes of KV cache per token.

Why the pressure exists at all. From the arithmetic above, KV capacity limits batch, batch is the arithmetic intensity, and intensity is throughput. A model that halves its KV cache roughly doubles the batch that fits and therefore roughly doubles decode throughput. That is an enormous return for a model-architecture change, and far larger than any plausible kernel optimisation. So the pressure is entirely legitimate.

The case that this is good codesign. The constraint is real and permanent — it follows from the beachfront limit and the memory hierarchy, not from a fixable software deficiency — so a model design that ignores it is leaving a factor of several on the table for no benefit. The accuracy cost of GQA is genuinely small at reasonable group sizes, which means the original MHA design was spending memory on a degree of freedom the model did not need. Finding that out is exactly what codesign is for: information flowed from the hardware constraint back into the model design and improved both. MLA goes further and arguably improves quality per byte outright.

The case that it is a symptom. The model is now shaped by a hardware artefact, and hardware artefacts move. If HBM capacity per FLOP improved sharply, or if a memory technology changed the constraint, we would have a generation of models carrying a compromise for a bottleneck that no longer binds. There is precedent — the entire convolution-efficiency literature of 2016–2019 optimised for a reuse structure that transformers then made irrelevant. And there is a subtler cost: the constraint is strongest at inference, so the model is being shaped by deployment economics in ways that are hard to separate from what makes it good.

The resolution, which is the general lesson: codesign is legitimate when it removes a degree of freedom the model did not need, and dangerous when it removes one the model did need but whose value is hard to measure. GQA looks like the first. It is worth noticing that the way you tell them apart is not available at design time — you find out later, which is the same epistemic problem Module 8 has with cost models and Module 14 has with the whole loop.

Because the binding constraint was memory fragmentation, not memory bandwidth or compute, and fragmentation is an allocator problem.

The situation before. A serving system does not know how long a sequence will be, so it allocates KV cache for the maximum context length up front — 8 K, 32 K, whatever the model supports — for every sequence in flight. A request that generates 200 tokens holds an allocation sized for 8192. The utilisation of the KV region is then the ratio of actual to maximum length, which in real traffic is often under 30%. Since KV capacity sets the batch and batch sets the intensity, two thirds of your throughput was lost to an allocation policy.

PagedAttention applies virtual memory’s answer. Split the KV cache into fixed-size blocks (16 or 32 tokens), allocate blocks on demand as a sequence grows, and keep a per-sequence block table mapping logical positions to physical blocks. The attention kernel reads through the block table instead of assuming contiguity. Internal fragmentation drops to at most one partial block per sequence, so utilisation goes from ~30% to over 90%, batch grows by 2–4×, and throughput follows.

Two aspects deserve emphasis:

• The kernel got slightly slower and the system got much faster. Indirection through a block table costs a lookup and loses some coalescing (Module 3’s granularity, again). That is a small constant against a 2–4× batch increase. The instructive part is that a purely local measurement — kernel microbenchmark — would have rejected the change, which is a general hazard when optimising a system by measuring its parts.
• Sharing falls out for free. Once blocks are indirected, two sequences with a common prefix can point at the same physical blocks with copy-on-write. Shared system prompts, beam search and parallel sampling all become nearly free in memory. This is fork/COW, arriving in an inference server for the same reasons it arrived in an operating system.

The generalisable lesson is the one this whole Part is arguing: find the binding resource before optimising, and be willing for the answer to be unglamorous. Everyone working on inference at that time was optimising kernels — the FLOPs and the bandwidth — while the actual constraint was an allocator wasting 70% of a capacity that determined throughput. Nothing in a roofline plot or a kernel profile shows you that. It shows up only in a resource accounting of the whole system, which is why Module 8 insisted on modelling the workload rather than the kernel.