Flash Attention — the n² flag comes down

Now the older flag, planted in Topic 7: attention compares every token with every other token — O(n²). For years this hard-capped context windows. The fix that broke the cap is the most instructive optimization story in the field, because what it optimizes is not what anyone expected.

The surprise: Flash Attention does not reduce the n² math. All the same comparisons still happen. What it eliminates is the n² memory traffic.

To see why that's the win, one more piece of GPU anatomy — the memory hierarchy inside the GPU. VRAM (the HBM chips) is the big pool: tens of GB at ~1.5–3 TB/s. But on the chip itself sits SRAM — a tiny scratchpad, ~20 MB total, at roughly 10× the bandwidth. Whiteboard and filing cabinet: the whiteboard is lightning-fast and fits almost nothing; the cabinet holds everything and requires a walk.

Now watch standard attention at 32K context. The score matrix — every token against every token — is 32,768² ≈ 1.1 billion entries, ~2 GB in FP16, for a single head in a single layer. Nowhere near fitting on the whiteboard, so the naive algorithm computes it in pieces, files the entire giant matrix into the cabinet (HBM), then walks back to retrieve it all for the softmax, files the result, retrieves it again to multiply by V… The GPU spends its time walking to the cabinet, cooks idle, while the actual arithmetic — which tensor cores could do almost instantly — waits on memory. Attention was never compute-bound. It was IO-bound, and nobody had organized the work around that fact.

Flash Attention (Tri Dao et al., 2022) reorganizes it. Two ideas:

1. Tiling: chop Q, K, V into small blocks that fit on the whiteboard. Load a block of Q and a block of K/V into SRAM, compute their piece of attention entirely on-chip, accumulate into the running output, discard, next block. The full n×n matrix is never written to HBM — it never exists anywhere.
2. Online softmax — the trick that makes tiling legal. Softmax seems to need all scores before converting any to percentages (you need the total to divide by). The online variant processes blocks sequentially while keeping a running maximum and running sum, retroactively rescaling earlier contributions whenever a new block shifts the statistics. Like computing a class average while grading papers one stack at a time — keep running totals, adjust as you go, and the final number is exactly what grading all at once would give.

That last word is the part people disbelieve: Flash Attention is exact. Not an approximation, not a quality trade — the mathematically identical output, computed in a smarter order. The results: attention 2–4× faster, and attention's memory footprint drops from O(n²) to O(n) — which, together with the KV cache tricks of Topic 32, is precisely what unlocked the 128K-context era. Today it's simply the default — baked into PyTorch, every serving engine, every training framework (your Lesson 6 capstone used it without asking). Successors (FlashAttention-2, -3) refine the same idea for newer chips.

And the general lesson outranks the specific technique: on modern hardware, arithmetic is nearly free; moving data is expensive. The biggest speedups of the past few years — Flash Attention, quantization's speed side, kernel fusion — are all fundamentally about touching memory less. Optimization stopped being about doing less math and became about fewer trips to the cabinet. Carry that lens into the next topic, where it becomes a formula.

---