Flash attention and efficient attention: Visual walkthrough

Step 1 — What are we even computing? Attention as a weighted average

Let us name the players, and never use a symbol before it is a picture.

$q$ — one query vector (a row). It is the token that is "asking".
$k_{1}, k_{2}, \dots, k_{N}$ — the key vectors, one per token. Each key is what a token "advertises".
$v_{1}, v_{2}, \dots, v_{N}$ — the value vectors, the actual content each token would contribute.
$N$ — how many tokens there are (the sequence length).
$d$ — how many numbers are inside each vector (the dimension).

The score between the query and key $j$ is their dot product $s_{j} = q \cdot k_{j}$ . A dot product is big when two arrows point the same way — so a big $s_{j}$ means "this token matches what I'm asking for".

PICTURE. Below: the query arrow, four key arrows, and the weights they earn. The output is the amber blended arrow.

Figure — Flash attention and efficient attention

The trouble is hiding in plain sight: the denominator needs every score before you can finish. That is the whole problem of this page.

Step 2 — Why we can't just start summing: the denominator blocks us

WHAT we want instead: pour keys through one block at a time, and never hold more than a small tile in fast memory.

WHY it seems impossible: division by a total you haven't finished. If you commit to a weight now, a later key could change the total and make your weight wrong.

PICTURE. The full row of scores vs. the tile we are allowed to see. Everything greyed-out is what we are forbidden to store.

The escape: don't compute final weights early. Keep a running numerator and a running denominator, and fix them up whenever new information arrives. That fix-up is the heart of the algorithm.

Step 3 — The overflow danger, and the max-shift that defuses it

Before we tile, one hazard. Scores can be large; $e^{50}$ overflows a computer's float. So everyone uses the safe softmax: subtract the biggest score first.

(To keep symbols light we fold $1/ d$ into the scores from here on: read $s_{j}$ as "already scaled".)

WHY this matters for tiling: each tile only knows the max of its own tile. Different tiles will use different reference points $m$ . When we merge tiles we must convert them to a common reference — that conversion is a single multiply, and it is the one clever line in Flash Attention.

PICTURE. Left: raw exponentials skyrocketing. Right: after shifting by $m$ every bar sits at or below $1$ .

Step 4 — Book-keeping per tile: three running numbers

For each tile of keys we compute exactly three things and carry them forward.

The final answer, once all tiles are done, is one division: $O = \tilde{O} / ℓ$ . We delay the division to the very end — that is how we dodge the "need the total first" trap.

For a single fresh tile $b$ with its own scores $x^{(b)}$ : $m^{(b)} = max (x^{(b)}), ℓ^{(b)} = \sum_{j \in b} e^{x_{j} - m^{(b)}}, \tilde{O}^{(b)} = \sum_{j \in b} e^{x_{j} - m^{(b)}} v_{j}$

$m^{(b)}$ — this tile's local reference point.
$ℓ^{(b)}$ — this tile's local denominator.
$\tilde{O}^{(b)}$ — this tile's local value-blend (not yet divided).

PICTURE. One tile arriving; the three quantities it produces shown as a max-marker, a sum-bar, and a small blended vector.

Step 5 — The merge, derived symbol by symbol

Now the key move. We hold old state $(m^{old}, ℓ^{old}, \tilde{O}^{old})$ and a new tile $(m^{(b)}, ℓ^{(b)}, \tilde{O}^{(b)})$ . How do we combine them without re-reading old scores?

First, the new common reference is the bigger of the two maxima: $m^{new} = max (m^{old}, m^{(b)})$

Now the crux. A sum measured against $m^{old}$ can be re-based to $m^{new}$ by multiplying by one factor. Watch why:

j \sum e^{x_{j} - m^{new}} = j \sum e^{(x_{j} - m^{old}) + (m^{old} - m^{new})} = rescale factor e^{m^{old} - m^{new}} = ℓ^{old} j \sum e^{x_{j} - m^{old}}

Every term picked up the same extra factor $e^{m^{old} - m^{new}}$ , so it factors out. That single multiply "moves" a whole sum to the new reference point — no old scores needed.

PICTURE. Two sums at different reference heights; the amber rescale arrow pulls the lower one up to the shared baseline before they add.

Step 6 — Run it: the full example, tile by tile

Take the parent note's numbers ( $N = 4$ , $d = 2$ , tile size $2$ ), following query row $0$ only. Scores are already scaled by $1/ 2$ .

Tile 1 (keys 0,1): scores $x^{(1)} = [1/ 2, 0] \approx [0.707, 0]$ .

$m^{(1)} = 0.707$
exponents $[e^{0}, e^{- 0.707}] = [1, 0.493]$ , so $ℓ^{(1)} = 1.493$
$\tilde{O}^{(1)} = 1 \cdot [1, 1] + 0.493 \cdot [2, 0] = [1.986, 1.0]$

Tile 2 (keys 2,3): scores $x^{(2)} = [1/ 2, 0] \approx [0.707, 0]$ .

$m^{(2)} = 0.707$
exponents $[1, 0.493]$ , so $ℓ^{(2)} = 1.493$
$\tilde{O}^{(2)} = 1 \cdot [0, 2] + 0.493 \cdot [1, 1] = [0.493, 2.493]$

Merge: here $m^{new} = max (0.707, 0.707) = 0.707$ , so both rescale factors are $e^{0} = 1$ .

$ℓ^{new} = 1.493 + 1.493 = 2.986$
$\tilde{O}^{new} = [1.986, 1.0] + [0.493, 2.493] = [2.479, 3.493]$
$O = \tilde{O}^{new} / ℓ^{new} = [2.479, 3.493] /2.986 \approx [0.830, 1.170]$

This equals what standard attention gives for row $0$ — but the $4 \times 4$ matrix was never stored.

PICTURE. The two tiles' running state, then the merged answer, drawn as accumulating bars.

Step 7 — Degenerate and edge cases (never leave the reader stranded)

Real inputs are not always tidy. Every case still works — here is why.

Formula Edge cases, each resolved

Empty / zero tile ( $v_{j} = 0$ or a masked block): its $\tilde{O}^{(b)} = 0$ and it still adds a valid $ℓ^{(b)}$ ; merge is unaffected.
All scores equal ( $x_{j} = c$ ): after max-shift every exponent is $1$ , so $ℓ = (tile size)$ and the output is the plain average of values — softmax with no preference, as expected.
One dominating score ( $x_{1} ≫$ rest): after shifting by $m = x_{1}$ , all other exponents $\to 0$ , so $O \to v_{1}$ . Attention collapses to a single token — a "hard" lookup.
A very negative masked score ( $x_{j} = - \infty$ , causal mask): $e^{- \infty} = 0$ , contributes nothing to $ℓ$ or $\tilde{O}$ . Masking is free.
First tile (no old state): initialise $m^{old} = - \infty, ℓ^{old} = 0, \tilde{O}^{old} = 0$ . The rescale factor $e^{- \infty - m^{new}} = 0$ wipes the empty history — the merge rule handles the very first step with no special code.

PICTURE. Three mini-panels: uniform scores → flat average; one spike → single-token pickup; masked entry → zero contribution.

Step 8 — Backward pass: the recompute trade

Training needs gradients, which normally means storing $P$ (the $N \times N$ weights). Flash Attention refuses. Instead it stores only the tiny per-row statistics $(m, ℓ)$ and, in the backward pass, recomputes the score tiles from $Q, K$ on the fly — the same idea as gradient checkpointing.

PICTURE. Standard path: fat $N^{2}$ matrix parked in slow HBM. Flash path: only slim $(m, ℓ)$ stored, tiles recomputed in fast SRAM.

The one-picture summary

Read it left to right: keys stream in as tiles; each tile makes a local $(m, ℓ, \tilde{O})$ ; the merge rescales to a shared max and accumulates; one final division at the far right gives the exact same output as full-matrix attention — with linear memory.

Recall Feynman retelling — say it back in plain words

Imagine averaging exam scores, but the students hand papers in one class at a time and you're not allowed to keep the papers. You keep three sticky notes: the highest score so far, a running total of "weights", and a running total of "weight × value". Every time a new class arrives you note its own highest score, its own totals, then — because the two classes measured their weights against different high-score reference points — you shrink the lower-reference class's totals by one multiply so both use the same reference. Add them. At the very end, divide the weighted total by the total weight. You get exactly the answer you'd get if you'd kept every paper — but you never held more than one class's worth at once. Attention does this with keys instead of classes, in fast on-chip memory, and that is why it's both leaner and faster.

Recall Quick self-test

Why can't we compute a softmax weight before seeing all scores? ::: The weight divides by the sum over all scores; that total isn't known until the last score arrives — so we delay the division and keep a running numerator and denominator instead. What are the three running numbers Flash Attention keeps per query? ::: the running max $m$ , the running denominator $ℓ$ , and the running un-normalised output $\tilde{O}$ . Why is a rescale factor needed when merging two tiles? ::: The tiles measured their sums against different maxima; multiplying by $e^{m^{old} - m^{new}}$ re-bases a whole sum to the shared reference in one step. Both $ℓ$ and $\tilde{O}$ get rescaled by the same factors — why? ::: They are the denominator and numerator of the same fraction; if only one were rescaled they'd be measured against different references and the final division would be wrong. Why is the backward pass's extra recompute worth it? ::: It avoids reading/writing the $N^{2}$ matrix from slow HBM; saving that memory traffic outweighs the small arithmetic cost, so it's a net speedup.