Flash attention and efficient attention: Worked examples

This page is a drill deck. The parent note built the ideas (tiling, incremental softmax, recomputation). Here we hammer every shape the maths can take: normal blocks, blocks with different maxima, a degenerate all-equal row, a masked (causal) block, a single-block "edge" case, a memory word-problem, and an exam twist.

Throughout, HBM means High-Bandwidth Memory — the large but comparatively slow GPU RAM (tens of GB) where big matrices must otherwise be parked. The whole point of Flash Attention is to avoid writing the $N \times N$ matrix out to HBM.

Before anything, one reminder of the two running numbers Flash Attention carries per query row. If you have not met them, read them here first:

Now the exact recipe every example below follows. Learn this once and every cell is just plugging numbers in.

Recall Why never store the full matrix?

What object does standard attention write to slow GPU memory (HBM)? ::: The full $N \times N$ scores matrix $S$ and probabilities $P$ . What does Flash Attention keep instead, per query row? ::: Just three small things — a max $m$ , a sum $ℓ$ , and an unnormalized output vector $\tilde{O}$ . Why divide by $ℓ$ only at the very end and not per block? ::: Because $ℓ$ is not complete until the last block; dividing early would use a partial denominator.

The scenario matrix

Every worked problem below is tagged with the cell it covers. Together they fill the table.

Cell	Scenario class	What can go wrong / why it matters
A	Two blocks, equal maxima	Rescale factors are all $1$ — the "easy" case, warm-up
B	Two blocks, different maxima	Rescale factors $\neq = 1$ — this is the whole point of $m^{new}$
C	Degenerate: all scores equal	softmax collapses to uniform $1/ n$ ; check no divide-by-zero
D	Masked / causal block (future keys $= - \infty$ )	Whole entries drop out; $e^{- \infty} = 0$
D2	Fully masked row (every key $= - \infty$ )	$ℓ = 0$ → divide-by-zero danger; needs an explicit guard
E	Single block (edge case)	Flash must reduce exactly to plain softmax
F	Limiting: one score $\to + \infty$	softmax becomes a one-hot; output $\to$ one value vector
G	Word problem: memory in bytes	The real-world payoff — $O (N^{2}) \to O (N)$
H	Exam twist: does block order change the answer?	Tests understanding of associativity of the merge

Cell A — two blocks, equal maxima

Worked example A: warm-up merge with equal maxima

One query row with scores split into two blocks: $x^{(1)} = [2, 1], x^{(2)} = [2, 0]$ Value vectors (1-D for clarity): $v = [10, 20, 30, 40]$ (first two go with block 1, last two with block 2). Compute the final attention output using the incremental merge.

Forecast: the two block maxima are both $2$ , so guess — will any rescale factor differ from $1$ ?

Block 1. $m^{(1)} = max (2, 1) = 2$ . Exponents $[e^{0}, e^{- 1}] = [1, 0.3679]$ , so $ℓ^{(1)} = 1.3679$ . Why this step? Subtracting the block max makes the largest term exactly $1$ — no overflow. Unnormalized output $\tilde{O}^{(1)} = 1 \cdot 10 + 0.3679 \cdot 20 = 17.3576$ . Why this step? The block carries values weighted by their (unnormalized) exponentials, not raw values.
Block 2. $m^{(2)} = max (2, 0) = 2$ . Exponents $[e^{0}, e^{- 2}] = [1, 0.1353]$ , so $ℓ^{(2)} = 1.1353$ . Why this step? Same max-shift rule, applied independently to this block — each block is safe on its own before we ever merge. $\tilde{O}^{(2)} = 1 \cdot 30 + 0.1353 \cdot 40 = 35.4134$ .
Merge. $m^{new} = max (2, 2) = 2$ . Both rescale factors $e^{2 - 2} = 1$ . Why this step? Equal maxima → identical reference point → nothing to correct, so we simply add the block sums and block outputs (the generic merge rule with both factors $= 1$ ). $ℓ^{new} = 1 \cdot 1.3679 + 1 \cdot 1.1353 = 2.5032$ . $\tilde{O}^{new} = 1 \cdot 17.3576 + 1 \cdot 35.4134 = 52.7710$ .
Normalize. $O = 52.7710/2.5032 = 21.0812$ . Why this step? Only now do we divide by the (complete) denominator $ℓ^{new}$ to turn unnormalized weights into a true softmax average.

Verify: compute plain softmax over all four scores $[2, 1, 2, 0]$ at once and take $\sum p_{i} v_{i}$ . Same answer → the tiling was exact.

Cell B — two blocks, different maxima

This is the cell people get wrong, because they forget to rescale block 1 when block 2 reveals a bigger max. The figure below is a bar chart of the unnormalized exponential weight of each score.

Figure — Flash attention and efficient attention

How to read the figure. The horizontal axis lists the four scores in two groups; the vertical axis is the exponential weight (before normalizing). The blue bars are block 1 as it was first computed, using its own local max $m^{(1)} = 1$ : the top score gives $e^{0} = 1$ , the other $e^{- 1} = 0.3679$ . The pink bars are block 2, computed with its own max $m^{(2)} = 3$ (again $1$ and $0.3679$ ). The yellow arrow marks the moment we discover the true global max lives in block 2, so $m^{new} = 3$ . The yellow bars show what block 1 becomes after we multiply it by $e^{m^{(1)} - m^{new}} = e^{1 - 3} = 0.1353$ : the tall blue $1$ shrinks to $0.1353$ and the $0.3679$ shrinks to $0.0498$ . Those shrunken yellow bars are the ones we actually sum. So the picture literally says: block 1's bars must be pushed down to the same baseline as block 2 before adding.

Worked example B: the rescale actually fires

Scores $x^{(1)} = [1, 0]$ (block 1), then $x^{(2)} = [3, 2]$ (block 2). Values $v = [10, 20, 30, 40]$ .

Forecast: block 1 sets $m^{(1)} = 1$ . Block 2 has a bigger entry ( $3$ ). Guess: will block 1's stored sum be too big or too small once we adopt the new global max?

Block 1. $m^{(1)} = 1$ . Exponents $[e^{0}, e^{- 1}] = [1, 0.3679]$ , $ℓ^{(1)} = 1.3679$ , $\tilde{O}^{(1)} = 1 \cdot 10 + 0.3679 \cdot 20 = 17.3576$ . Why this step? We do not yet know a bigger score is coming, so block 1 legitimately measures itself against its own max $m^{(1)} = 1$ — this is the "provisional reference point" (the blue bars in the figure) that the merge will later correct.
Block 2. $m^{(2)} = max (3, 2) = 3$ . Exponents $[e^{0}, e^{- 1}] = [1, 0.3679]$ , $ℓ^{(2)} = 1.3679$ , $\tilde{O}^{(2)} = 1 \cdot 30 + 0.3679 \cdot 40 = 44.7152$ . Why this step? Block 2 shifts by its own max $3$ (the pink bars); notice this is larger than block 1's max, which is exactly the situation the rescale step exists to handle.
Merge with correction. $m^{new} = max (1, 3) = 3$ . Why this step? The two blocks used different reference points ( $1$ vs $3$ ). Apply the generic merge rule: re-express both against the common baseline $m^{new} = 3$ . Block 1's exponentials were computed with the smaller subtraction, so they are inflated by exactly $e^{m^{(1)} - m^{new}} = e^{1 - 3} = e^{- 2} = 0.1353$ ; multiply them down (blue bars → yellow bars). This factor is the entire reason the merge formula exists.
- Block 1 factor: $e^{1 - 3} = 0.1353$ . Block 2 factor: $e^{3 - 3} = 1$ .
- $ℓ^{new} = 0.1353 \cdot 1.3679 + 1 \cdot 1.3679 = 1.5530$ .
- $\tilde{O}^{new} = 0.1353 \cdot 17.3576 + 1 \cdot 44.7152 = 47.0641$ .
Normalize. $O = 47.0641/1.5530 = 30.3053$ . Why this step? Divide by the corrected denominator to recover the exact softmax average.

Verify: plain softmax over $[1, 0, 3, 2]$ , then $\sum p_{i} v_{i}$ , gives $30.3053$ . If we had forgotten the $0.1353$ factor and naively added the raw block sums, we'd get a different, wrong number biased toward the small-score block. The correction is what makes it exact.

Cell C — degenerate: all scores equal

Worked example C: every key equally relevant

All four scores equal: $x = [5, 5, 5, 5]$ , split $[5, 5] ∣ [5, 5]$ . Values $v = [10, 20, 30, 40]$ .

Forecast: if every score is identical, what should each attention weight be? Guess before computing.

Block 1. $m^{(1)} = 5$ , exponents $[e^{0}, e^{0}] = [1, 1]$ , $ℓ^{(1)} = 2$ , $\tilde{O}^{(1)} = 1 \cdot 10 + 1 \cdot 20 = 30$ . Why this step? After subtracting the shared max, every exponent is $e^{0} = 1$ — perfectly uniform.
Block 2. $m^{(2)} = 5$ , $ℓ^{(2)} = 2$ , $\tilde{O}^{(2)} = 30 + 40 = 70$ . Why this step? Identical logic — the flat scores give $e^{0} = 1$ again, so both blocks look the same internally.
Merge. $m^{new} = 5$ , factors both $1$ . $ℓ^{new} = 4$ , $\tilde{O}^{new} = 100$ . Why this step? The block maxima are equal ( $5 = 5$ ), so $e^{5 - 5} = 1$ — no correction needed and we simply add.
Normalize. $O = 100/4 = 25$ . Why this step? Dividing by $ℓ^{new} = 4$ turns the four equal weights into $\frac{1}{4}$ each — a plain average of the values.

Verify: uniform weights $= \frac{1}{4}$ each, so $O = \frac{1}{4} (10 + 20 + 30 + 40) = 25$ . ✓ And note $ℓ^{new} = 4 \neq = 0$ — no divide-by-zero even in the flat case, because at least one exponential is $e^{0} = 1$ .

Cell D — masked / causal block

In GPT-style decoders a query at position $t$ may not look at future keys $j > t$ . Flash Attention implements this by setting those scores to $- \infty$ before the exponential.

Worked example D: causal mask kills future keys

Query is at position 1. Raw scores against keys $0, 1, 2, 3$ are $[2, 1, 4, 3]$ , but positions $2, 3$ are in the future → masked to $- \infty$ . Blocks: $[2, 1] ∣ [- \infty, - \infty]$ . Values $v = [10, 20, 30, 40]$ .

Forecast: with the future fully masked, guess which value vectors can possibly appear in the output.

Block 1 (visible). $m^{(1)} = 2$ , exponents $[e^{0}, e^{- 1}] = [1, 0.3679]$ , $ℓ^{(1)} = 1.3679$ , $\tilde{O}^{(1)} = 10 + 0.3679 \cdot 20 = 17.3576$ . Why this step? These are the allowed (past+present) keys, handled by the ordinary max-shift rule.
Block 2 (masked). Scores $- \infty$ . $e^{- \infty} = 0$ for every entry, so $ℓ^{(2)} = 0$ , $\tilde{O}^{(2)} = 0$ . Why this step? $- \infty$ is the softmax way of saying "impossible" — its weight is exactly $0$ , no rounding needed.
Merge. Take $m^{new} = max (2, - \infty) = 2$ . Block 2 factor $e^{- \infty - 2} = 0$ . Why this step? The global max is set by the only real scores (block 1). The masked block contributes zero sum and zero output, so multiplying it by any factor still leaves nothing — the future cannot leak in. $ℓ^{new} = 1 \cdot 1.3679 + 0 = 1.3679$ , $\tilde{O}^{new} = 17.3576$ .
Normalize. $O = 17.3576/1.3679 = 12.6893$ . Why this step? Divide by the denominator built only from visible keys — the masked keys never entered it.

Verify: softmax over only the two visible scores $[2, 1]$ gives weights $[0.7311, 0.2689]$ ; $0.7311 \cdot 10 + 0.2689 \cdot 20 = 12.6893$ . ✓ Future values $30, 40$ never leak in — exactly what causal masking must guarantee.

Recall Masking mechanics

How does Flash Attention exclude a forbidden key? ::: Set its score to $- \infty$ so $e^{score - m} = 0$ ; it contributes nothing to $ℓ$ or $\tilde{O}$ .

Cell D2 — the fully masked row (divide-by-zero danger)

Worked example D2: every key is forbidden

Query at position 0 in a strictly causal setup (a token may attend only to earlier positions, not itself). Position 0 has no earlier tokens, so every score is $- \infty$ : blocks $[- \infty, - \infty] ∣ [- \infty, - \infty]$ . Values $v = [10, 20, 30, 40]$ .

Forecast: with $ℓ = 0$ the normalize step $O = \tilde{O} / ℓ$ divides by zero. What should a correct implementation return?

Both blocks. Every $e^{- \infty} = 0$ , so $ℓ^{(1)} = ℓ^{(2)} = 0$ and $\tilde{O}^{(1)} = \tilde{O}^{(2)} = 0$ . Why this step? Nothing is attendable, so there is no exponential mass anywhere.
Merge. $m^{new} = max (- \infty, - \infty) = - \infty$ ; $ℓ^{new} = 0$ , $\tilde{O}^{new} = 0$ . Why this step? The generic merge rule still runs, but every term is zero — the running state never leaves its empty start.
Guard before normalize. Detect $ℓ^{new} = 0$ and do not divide. Return the convention $O = 0$ (zero vector). Why this step? $0/0$ is undefined; real Flash Attention kernels either (a) initialize $m = - \infty, ℓ = 0$ and output $0$ when nothing was attended, or (b) never emit such a row because a token always at least attends to itself. The lesson: the divide is only safe because a well-formed mask always leaves $\geq 1$ visible key, giving $ℓ \geq e^{0} = 1 > 0$ .

Verify: $ℓ^{new} = 0$ exactly, so the guard fires and the output is the zero vector by convention. This is the one row where you must not blindly compute $\tilde{O} / ℓ$ .

Cell E — single block (edge case)

Worked example E: one block must equal plain softmax

Only one block: $x = [3, 1, 0]$ , values $v = [10, 20, 30]$ .

Forecast: with a single block there is nothing to merge. Does Flash Attention give a different number from ordinary softmax? Guess.

$m = 3$ . Exponents $[e^{0}, e^{- 2}, e^{- 3}] = [1, 0.1353, 0.0498]$ . Why this step? One block means one max-shift, against the single max $3$ — no merge machinery is invoked at all.
$ℓ = 1 + 0.1353 + 0.0498 = 1.1851$ . Why this step? The denominator is complete after one pass because there is no second block to wait for.
$\tilde{O} = 1 \cdot 10 + 0.1353 \cdot 20 + 0.0498 \cdot 30 = 14.2000$ . Why this step? Values weighted by their exponentials, exactly as in standard softmax attention.
$O = 14.2000/1.1851 = 11.9821$ . Why this step? Normalizing once gives the final answer — the merge step is the identity here.

Verify: plain softmax over $[3, 1, 0]$ then weighted values gives the identical $11.9821$ . ✓ The merge step is the identity when there is one block — a sanity anchor for the whole algorithm.

Cell F — limiting behaviour: one score dominates

Worked example F: as one score grows, softmax becomes one-hot

Scores $[x, 0]$ over two keys, values $v = [10, 20]$ . Study $x = 0, 2, 10, 100$ .

Forecast: as $x \to + \infty$ , which single value should the output snap to?

$x = 0$ : weights $[0.5, 0.5]$ → $O = 0.5 \cdot 10 + 0.5 \cdot 20 = 15$ . Why this step? Equal scores ⇒ equal weights ⇒ plain average — the balanced baseline.
$x = 2$ : $m = 2$ , exponents $[1, e^{- 2}] = [1, 0.1353]$ , $ℓ = 1.1353$ , $O = (1 \cdot 10 + 0.1353 \cdot 20) /1.1353 = 11.1917$ . Why this step? A modest gap already tilts most weight onto key 0, pulling the output toward $10$ .
$x = 10$ : exponents $[1, e^{- 10}] \approx [1, 0.0000454]$ , $O \approx (10 + 0.000908) /1.0000454 \approx 10.00036$ . Why this step? The second exponential is now negligible — key 1 is almost switched off.
$x = 100$ : $e^{- 100} \approx 0$ , $O \to 10$ exactly to machine precision. Why this step? This shows the max-shift is what keeps us safe: we compute $e^{0} = 1$ and $e^{- 100}$ , never the overflowing $e^{100}$ .

Verify: the limit $lim_{x \to \infty} O = 10$ (the value of the dominating key). At $x = 2$ the number $11.1917$ checks against direct softmax. Monotone: as $x$ grows, $O$ slides from $15$ down toward $10$ .

Cell G — word problem: the memory payoff in bytes

Worked example G: how much HBM does the

$N \times N$ matrix cost? A model has sequence length $N = 4096$ , hidden per head $d = 64$ , $H = 16$ heads, $L = 24$ layers, using 4-byte floats. Compare the memory to store the full attention probability matrices vs. Flash Attention's per-row statistics.

Forecast: roughly how many gigabytes just for the $P$ matrices?

One head's $P$ matrix (element count). $N^{2} = 409 6^{2} = 16, 777, 216$ floats. Why this step? Standard attention materializes an $N \times N$ probability grid and parks it in HBM.
Bytes per head. $16, 777, 216 \times 4 = 67, 108, 864$ bytes $= 64$ MB. Why this step? Multiply element count by 4 bytes/float to get the true HBM footprint.
All heads, one layer. $64 MB \times 16 = 1024 MB = 1$ GB per layer. Why this step? Each of the $H = 16$ heads holds its own independent $P$ matrix.
All layers (total, standard). $1 GB \times 24 = 24$ GB. Why this step? Every layer repeats the cost during a forward+backward pass.
Flash Attention (total). Per query row keep $m, ℓ$ (2 floats) + output vector ( $d = 64$ floats) $\approx 66$ floats. Across all rows/heads/layers: $66 \times 4096 \times 16 \times 24 \times 4 bytes = 415, 236, 096 bytes \approx 0.39$ GB. Why this step? Memory is now $O (N \cdot d)$ , linear in $N$ , not $O (N^{2})$ — the whole selling point.

Verify: ratio $24 GB /0.39 GB \approx 62 \times$ less attention memory. Units check: floats × 4 bytes = bytes; ÷ $2^{20}$ = MB, ÷ $2^{30}$ = GB. ✓ This is why long-context models are feasible at all — closely related to the memory savings of gradient checkpointing.

Cell H — exam twist: does block order matter?

Worked example H: process blocks in reverse — same answer?

Reuse Cell B's scores $[1, 0 ∣ 3, 2]$ , values $[10, 20, 30, 40]$ , but stream block 2 first, then block 1. Do we get the same $O = 30.3053$ ?

Forecast: the merge multiplies old sums by $e^{m^{old} - m^{new}}$ . Is that operation order-independent?

First process $[3, 2]$ . $m = 3$ , $ℓ = 1 + e^{- 1} = 1.3679$ , $\tilde{O} = 30 + 0.3679 \cdot 40 = 44.7152$ . Why this step? Whichever block arrives first sets the running max; here it happens to be the true global max $3$ already.
Now merge in $[1, 0]$ . Its local $m_{b} = 1$ , $ℓ_{b} = 1.3679$ , $\tilde{O}_{b} = 10 + 0.3679 \cdot 20 = 17.3576$ . Why this step? The incoming block still measures itself against its own max $1$ before the merge corrects it.
Apply the generic merge rule. $m^{new} = max (3, 1) = 3$ . Factors: current running state (max $3$ ) $\to e^{3 - 3} = 1$ ; incoming block (max $1$ ) $\to e^{1 - 3} = 0.1353$ . Why this step? Now it is the incoming block that gets scaled down toward the same global max $3$ — the rescale formula is symmetric, it always shrinks whatever was measured against a smaller max. $ℓ^{new} = 1 \cdot 1.3679 + 0.1353 \cdot 1.3679 = 1.5530$ . $\tilde{O}^{new} = 1 \cdot 44.7152 + 0.1353 \cdot 17.3576 = 47.0641$ .
Normalize. $O = 47.0641/1.5530 = 30.3053$ . Why this step? Normalizing the order-swapped totals recovers exactly Cell B's answer.

Verify: identical to Cell B ( $30.3053$ ). The merge is associative and commutative because everything is ultimately measured against the single true global max $m^{new}$ . Order of tiles is a free choice — which is exactly why the GPU can schedule blocks however it likes.

Connections

Built on ordinary self-attention and used per-head inside multi-head attention.
The recomputation trick is the same idea as gradient checkpointing — trade compute for memory.
Long-context feasibility feeds directly into scaling laws, and powers both BERT and GPT style models (the causal masks in Cells D and D2 are decoder features).
Masking interacts with positional encoding since positions decide which keys are in the future.

Recall Self-test

In Cell B, why is the block-1 rescale factor $0.1353$ ? ::: Because $e^{m^{(1)} - m^{new}} = e^{1 - 3} = e^{- 2} = 0.1353$ ; block 1 was measured against too-small a max. Why does Cell C never divide by zero? ::: After the max-shift, at least one exponential is $e^{0} = 1$ , so $ℓ \geq 1 > 0$ . What must a correct kernel do for a fully masked row (Cell D2)? ::: Detect $ℓ = 0$ and return the zero vector instead of computing $0/0$ . Does streaming blocks in reverse order change the result (Cell H)? ::: No — the merge is commutative/associative because all blocks are rescaled to the same global max.