Attention and KV Compression

## Why a kernel rewrite mattered at all Modern accelerators are enormously faster at arithmetic than at moving data. A GPU can perform hundreds of floating-point operations in the time it takes to fetch one value from its main memory. Attention, implemented naively, is unusually bad on this axis: it produces a large intermediate — the score matrix relating every query to every key — and then reads and writes that intermediate several times. The operation is therefore **memory-bound**, not compute-bound, and speeding up the arithmetic does nothing. FlashAttention's contribution was to notice this and restructure the computation so that the intermediate never leaves fast on-chip memory. ## The memory hierarchy in one paragraph A GPU has a large pool of high-bandwidth memory (HBM) — gigabytes, relatively slow — and a small pool of on-chip SRAM per compute unit — kilobytes, roughly an order of magnitude faster. Anything that must round-trip through HBM costs far more than the same work done inside SRAM. A naive attention kernel writes the scores to HBM, reads them for softmax, writes the normalised probabilities, reads them again to multiply by values. Each of those passes touches an object that grows quadratically with sequence length. ## Tiling and online softmax FlashAttention loads a block of queries and then streams through blocks of keys and values. For each tile it computes that tile's scores directly in SRAM, applies the exponentials there, and accumulates a partial weighted sum of values. The difficulty is that softmax normalises over the *whole* row, and you do not know the row's maximum or its denominator until you have seen every key. The fix is **online softmax**: carry a running maximum and a running sum of exponentials alongside the partial output. When a new tile arrives with a larger maximum, rescale the accumulated output and denominator by the exponential of the difference, then add the new tile's contribution. At the end the accumulator holds exactly the value the one-shot computation would have produced. This is numerically stable, because the running maximum is always subtracted before exponentiating, and it is exact up to the usual reordering effects of floating-point addition. Only the final output — and small per-row statistics — is written to HBM. The quadratic-sized intermediate is never materialised there at all. ## Recomputation in the backward pass Training needs the attention probabilities again to compute gradients. Storing them would reintroduce the very object the forward pass avoided. FlashAttention instead **recomputes** each tile's scores during the backward pass from the saved queries, keys, values and the small saved statistics. That costs extra arithmetic, which is cheap, in exchange for avoiding memory traffic, which is expensive. This is the same trade as gradient checkpointing, applied inside a single operation. ## What it does and does not change It changes: wall-clock time, achievable sequence length within a given memory budget, and the activation memory needed for training. Longer contexts became practical partly because attention stopped requiring a quadratic-sized buffer. It does not change: the model's weights, its outputs (up to floating-point reassociation), its parameter count, or the asymptotic number of arithmetic operations — every query still scores against every key. This is the point people most often get wrong. FlashAttention is not a way to make attention sub-quadratic. If you need to stop doing the work at all, you need a different mechanism: learned sparsity to skip most pairs, or a recurrent formulation that avoids the pairwise structure entirely. FlashAttention makes the exact computation as cheap as it can be; it does not make it a different computation. Later versions refined the same idea for newer hardware — better work partitioning across compute units, fewer non-matrix operations on the critical path, and support for narrower numeric formats — but the principle is unchanged. ## The composability point Because it is exact and weight-preserving, a FlashAttention-class kernel composes with essentially everything else in this area. Head sharing changes how many key/value heads exist; the kernel still tiles over them. Sparse selection changes which blocks are attended; the kernel runs on the selected blocks. As of mid-2026 an efficient tiled kernel is simply assumed in any serious training or serving stack rather than being a named optimisation you choose. ## How to answer Say "exact, not approximate" in the first sentence — that is the discriminator. Then explain memory-bound versus compute-bound, then tiling with online softmax, then backward recomputation. Close by stating explicitly that the operation count is unchanged, which pre-empts the follow-up about whether it makes attention linear.