Naive attention builds the full n×n matrix of Query·Key scores and writes it to VRAM. FlashAttention never does — it tiles the work into blocks small enough to live in on-chip SRAM, so the matrix is never materialized. Same FLOPs, bit-identical answer; the win is memory movement. Drag the sequence length and tile size and watch 32 MiB collapse to 56 KiB.
Attention (Part 1 lesson 6) scores every token's Query against every token's Key — for n tokens that's an n×n matrix of scores. Naive attention writes that whole matrix to VRAM, reads it back to run softmax, then reads it a third time to blend the Values. At n = 8192 that's a 128 MiB matrix shuttled several times — per head, per layer.
FlashAttention never builds it. It cuts Q, K, V into tiles small enough to fit in on-chip SRAM (lesson 1: ~192 KB per SM), computes attention one tile-pair at a time, and keeps a running max and running sum — an online softmax — so each new tile corrects the partial answer without ever seeing a whole row.
The output is bit-identical to naive attention. The only thing that changed is where the numbers live: the quadratic matrix stays on-chip and is thrown away tile by tile. It's the roofline lesson applied to attention — stop moving O(n²) bytes through slow VRAM.
TRY → Start at n = 4096, B = 64, fp16: the greyed matrix is 32 MiB of scores that naive attention writes to VRAM; Flash keeps just the lit 56 KiB tile on-chip — 4096× less. Push n to 16384 and the matrix balloons to 512 MiB while the tile never grows. Now drag B up to 256: the tile blows past the 192 KiB SRAM budget and a rose won't fit badge appears — switch to int8 and it fits again. Every figure recomputes from n² vs (3·B·d + B²).
Naive vs one tile. head dimension d = 128, fp16 (2 bytes/score). The naive score matrix is scoreMatrixBytes(4096) = 4096² · 2 = 33,554,432 B = 32 MiB — written to VRAM per head, per layer, then read back twice. One Flash tile needs room for its Q block, K block, V block and the small scores block: flashTileBytes(64) = (3·64·128 + 64²)·2 = (24,576 + 4,096)·2 = 57,344 B = 56 KiB — well under the 192 KiB SRAM budget.
| n | naive matrix (VRAM) | tiles/axis (B=64) | one tile (SRAM) | reduction (n/B)² |
|---|---|---|---|---|
| 1024 | 2 MiB | 16 | 56 KiB | 256× |
| 4096 | 32 MiB | 64 | 56 KiB | 4096× |
| 8192 | 128 MiB | 128 | 56 KiB | 16384× |
| 16384 | 512 MiB | 256 | 56 KiB | 65536× |
The tile size is fixed by B and d, not n — so as the sequence grows, the naive matrix grows quadratically while the on-chip tile stays 56 KiB. At n = 4096 there are flashTilesPerAxis(4096,64) = 64 tiles per axis ⇒ 64² = 4096 tile-pairs, and K/V are re-streamed once per Q-block, so flashKvPasses = 64.
The online softmax — why tiling stays exact. Softmax needs the row's maximum and the sum of all exponentials, but a tile only sees part of the row. So Flash carries two running numbers and patches them up as each new tile arrives:
m'.
running sum ℓthe accumulated denominator Σ e^(score − m). When m changes, the old sum is rescaled by e^(m − m') so it's measured against the new max, then the new tile's contribution is added.
running outputthe partial blend of Values is rescaled by the same e^(m − m') factor and the new tile's weighted Values are added in.
No tile ever needs the full row, so the n×n matrix is never assembled — yet the final m, ℓ and blend are exactly what a one-shot softmax over the whole row would give. The answer is identical to Part 1 lesson 6; only the memory schedule changed (the same "tricks don't change the output" point as the roofline).
n×n score matrix to VRAM — it tiles the work in on-chip SRAM. It's an IO trick: memory movement, not math.Naive attention materializes an O(n²) score matrix in slow VRAM (32 MiB per head at n=4096). FlashAttention tiles Q·K·V into blocks that fit in fast on-chip SRAM (56 KiB), uses an online softmax to stay exact, and never writes the matrix — a 4096× smaller memory footprint for the identical output.
(n/B)².