The same attention, tiled so the n×n score matrix is never stored. Stream the keys and values in blocks; keep a running max m and a running normalizer ℓ, and rescale the accumulator by exp(mold−mnew) every time the max moves. The output is bit-for-bit exact — not an approximation — while the memory drops from quadratic to linear. Down the center, data flows: the vectors go in, both engines run, the proven result comes out. The blue team builds it; the red team tries to break it.
source Dao, Fu, Ermon, Rudra, Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) — arxiv.org/abs/2205.14135. Rendered, not quoted.
Softmax needs its max and its sum, which classically means two passes over the whole row. The online recurrence folds both into one streaming pass. Carry three registers per query: max m, normalizer ℓ, accumulator acc (a d-vector).
For each new block of scores with block-max m′: set mnew=max(m,m′), correction c=exp(m−mnew), then
ℓ ← ℓ·c + Σ exp(sⱼ−mnew)
acc ← acc·c + Σ exp(sⱼ−mnew)·Vⱼ
out = acc / ℓ (once, at the end)
That single factor c retro-corrects everything already accumulated under the old, smaller max. Drop it and the sums are weighted wrong — window 6 does exactly that.
The softmax is the normalizer that turns scores into a distribution — but naïvely it wants the whole row at once to find its max and sum. Compute it online, in one pass with a running max, and the row never has to exist all at the same time.
That is the whole hinge: the attention made memory-linear without changing a single output value. The neighbour's identity becomes this sphere's engine.
The blue team's live check: re-run flash and naive on the same input and confirm the outputs agree to 1e-10, the flash path stays finite where the unstable naive path overflows, and flash's peak score memory is a single block. If red drops the rescale, this badge is where it shows.
Attention over a short sequence: n queries, keys and values, each a d-dim vector (d=4, fixed-seed PRNG). One score row is sⱼ = (Q·Kⱼ)/√d. The naive path forms the full n×n score matrix; the flash path only ever holds one block of it.
| path | scores held | output |
|---|---|---|
| naive | n × n | softmax(QKᵀ)·V |
| flash (tiled) | block size Bₐ | identical, exact |
Same arithmetic, different order — that is what you feed the panel below.
Moderate logits: the naive (no max-subtraction) path still survives. Both paths finite; flash matches naive exactly.
Change any control — naive and flash are recomputed live from the vectors, never looked up. The outputs are compared element-by-element.
What the machine produces, proven: for every setting, flash = naive to 1e-10 — exact, not approximate (the whole point of the paper). Peak score memory: naive holds — numbers, flash holds —. In the EXTREME regime the naive-without-max path is — while flash stays finite.
The blue team's witness (left) confirms these live; the red team (right) tries to make them wrong.
And it is only exact because addition here is treated as associative — in true fixed-order floating point the tiled sum can differ from the naive sum in the last bit. This toy holds to 1e-10; a real kernel documents ULP-level differences. "Exact" means the algorithm computes the same softmax, not that every rounding matches byte-for-byte.
"FlashAttention is an approximation, like Linformer / Performer." Cut. Those change the math to get sub-quadratic. Flash changes only the order of operations — same softmax, same output. It is exact.
"It makes attention O(n) in compute." Cut. O(n) in memory. FLOPs stay O(n²). The speedup is fewer reads/writes to slow memory, not fewer multiplies.
"The online rescaling is just an optimization you can skip." Kept, corrected. Skip the exp(mold−mnew) factor and the answer is simply wrong — window 6 proves it.
The red team's move: drop the accumulator rescaling on a max update — forget the exp(mold−mnew) factor. The online softmax now weights old blocks under the wrong max, so flash no longer equals naive. The blue team's witness (window 7) is watching.
Drop exp(mold−mnew) and the 1e-10 identity fails — the witness recomputes, finds flash ≠ naive, and turns red. Nothing is faked; the attack is real and it is caught.