Standard attention computes a whole N×N score matrix, runs a softmax over it, then multiplies by the values — three passes over data that will not fit in fast memory. FlashAttention gets the identical answer in a single streaming pass that never stores that matrix. The magic is not a hardware trick and not an approximation; it is a small, exact algebraic recurrence — the online softmax. This article derives that recurrence from scratch: the running maximum m_i that keeps exponentials finite, the running denominator l_i, the rescaling factor that retroactively corrects everything seen so far when a larger score appears, and the matching update to a running output accumulator. We then prove the two claims that make it trustworthy: that no exponential ever overflows, and that the final vector equals ordinary attention exactly.
What we must compute, and why storage is the problem
Fix a single query vector q and let the keys and values be k_1…k_N and v_1…v_N, each of dimension d. Attention produces one output vector:
s_j = (q · k_j) / √d (scalar score, j = 1..N)
a_j = exp(s_j) / Σ_t exp(s_t) (softmax weights, sum to 1)
o = Σ_j a_j · v_j (output, a d-vector)The textbook route computes the whole score vector s (for all queries at once, the whole matrix S = QK^T, shape [N, N]), then the softmax, then the weighted sum. Each stage reads and writes an O(N^2) array. That array is the villain: at long context it dwarfs the inputs and cannot stay in a chip’s small fast memory, so every stage pays to shuttle it to and from slow memory. FlashAttention’s goal is to produce o while touching each key and value once and never holding all of s at once.
Safe softmax: the max-subtraction we cannot give up
Softmax is written with a bare exp, but no real implementation uses it that way. Scores can be large; exp(89) already overflows a 32-bit float. The fix exploits an exact invariance: subtracting any constant m from every score leaves the weights unchanged, because the constant cancels between numerator and denominator.
a_j = exp(s_j - m) / Σ_t exp(s_t - m) for ANY constant mChoosing m = max_j s_j makes every exponent s_j - m ≤ 0, so every exp lands in (0, 1] — no overflow, ever. This is ‘safe softmax.’ The catch for streaming is that it needs the global maximum before any term, forcing one pass to find m, a second for the denominator, and a third for the weighted sum. Three passes over s is exactly what we are trying to avoid; the online softmax removes that ordering constraint without giving up the safety.
The running state: a max, a denominator, and an output
Process the keys in blocks, in any order, and carry three pieces of state that summarize everything seen so far. After absorbing some prefix of the scores:
m = max of the scores seen so far (scalar)
l = Σ exp(s_j - m) over scores seen (scalar denominator)
o~ = Σ exp(s_j - m) · v_j over scores seen (unnormalized output, a d-vector)Note that l and o~ are both defined relative to the current m. That is the whole subtlety: m is provisional. If a later block contains a score larger than the current m, every exponential we already summed was computed against the wrong reference and must be corrected. Crucially, we keep the output as the unnormalized sum o~, not the normalized o~/l — deferring the division lets numerator and denominator be rescaled by the same factor, which makes the correction a single cheap multiply rather than a re-derivation.
The rescaling identity, derived
Everything hinges on one line of algebra. Suppose our running max is m_old and a new block pushes it up to m_new ≥ m_old. A term we stored earlier holds exp(s_j - m_old), but it should now read exp(s_j - m_new). Relate the two:
exp(s_j - m_new) = exp(s_j - m_old) · exp(m_old - m_new)
= exp(s_j - m_old) · α, where α = exp(m_old - m_new)The correction is the same scalar α for every stored term, so it factors straight out of the sums: rescaling the accumulated l and o~ is just l ← α·l and o~ ← α·o~. Because m_new ≥ m_old, α ≤ 1 — the correction only ever shrinks old contributions, so it too cannot overflow. This single identity is the engine of the entire algorithm.
The block-wise update, and the full algorithm
Now the per-block step writes itself: compute the block’s local max, fold it into the running max, rescale the old state by α, then add the block’s contribution measured against the new reference. Assembled for a single query, iterating over blocks of keys and values:
init: m = -∞, l = 0, o~ = 0 (0 is the d-vector)
for each block B of (k, v):
s_j = (q · k_j)/√d for j ∈ B
m_blk = max_{j∈B} s_j
m_new = max(m, m_blk)
α = exp(m - m_new) (rescale old state, ≤ 1)
p_j = exp(s_j - m_new) (all ≤ 1)
l = α·l + Σ_j p_j
o~ = α·o~ + Σ_j p_j·v_j
m = m_new
return o = o~ / lInitializing m = -∞ makes the first block behave correctly: α = exp(-∞) = 0 wipes the empty accumulators. Each block contributes a matmul P·V and a reduction for l, both streaming operations that touch the block once; a single final division normalizes. In a real kernel this loop runs for a whole tile of queries at once, so the score block is a small matrix and P·V is a dense matmul.
A fully worked numeric example
Take four keys with scores s = [1, 3, 2, 5] and scalar values v = [10, 20, 30, 40], processed in two blocks of two. The true answer (reference max 5) has weights exp(s-5) = [0.0183, 0.1353, 0.0498, 1], denominator l = 1.2034, and output o = 44.383 / 1.2034 = 36.88. Running the recurrence:
Block 1: s=[1,3] m_blk=3 m_new=3 α=exp(-∞)=0
p=[exp(-2),exp(0)]=[0.1353,1]
l = 0·0 + 1.1353 = 1.1353
o~ = 0·0 + (0.1353·10 + 1·20) = 21.353
Block 2: s=[2,5] m_blk=5 m_new=5 α=exp(3-5)=0.1353
p=[exp(-3),exp(0)]=[0.0498,1]
l = 0.1353·1.1353 + (0.0498+1) = 1.2034
o~ = 0.1353·21.353 + (0.0498·30+40) = 44.383
o = 44.383 / 1.2034 = 36.88 ✓ identicalThe rescale factor α = 0.1353 is precisely what demotes the block-1 contributions once the larger score 5 appears, and both paths agree to every digit.
Proof of numerical stability
The stability claim is that no exponentiation in the algorithm can overflow, at any point, regardless of how large the raw scores are. It follows from two observations already established. First, the block probabilities: for every j ∈ B, m_new = max(m_old, m_blk) ≥ s_j, hence s_j - m_new ≤ 0 and p_j = exp(s_j - m_new) ∈ (0, 1]. Second, the rescale factor: m_new ≥ m_old gives m_old - m_new ≤ 0, so α = exp(m_old - m_new) ∈ (0, 1].
Every exponential argument in the whole run is therefore non-positive, so every exp is bounded by 1 — overflow is impossible by construction. The only floating-point loss is benign underflow: a score far below the running max yields a tiny p_j that rounds toward zero, exactly as in single-pass safe softmax, and contributes negligibly anyway. The recurrence inherits safe softmax’s numerical safety while dropping its three-pass requirement.
Proof of exactness
The second claim is stronger: the output is not close to standard attention, it is equal to it in exact arithmetic. Argue by loop invariant. Assume that after any number of blocks, m is the maximum of all scores seen, l = Σ exp(s_j - m), and o~ = Σ exp(s_j - m) v_j, both sums over exactly the scores seen. The update preserves this: multiplying by α = exp(m_old - m_new) rewrites every stored term from base m_old to base m_new (the rescaling identity), and the new block adds its terms in base m_new.
After the final block, m is the global maximum, so l = Σ_j exp(s_j - m) and o~ = Σ_j exp(s_j - m) v_j over all N keys, and o~/l = Σ_j a_j v_j — the definition of attention. Block size and traversal order never enter the result; they change only the intermediate values, not the fixed point.
O(N) memory, and the same FLOPs
The state per query is one scalar m, one scalar l, and one d-vector o~ — size O(d), independent of sequence length. Across a tile of queries we hold the query tile, the current key/value block, and the accumulators, so the working set never includes the O(N^2) score matrix.
That is the headline: memory drops from O(N^2) to O(N) in sequence length. The arithmetic is unchanged — still O(N^2 d) multiply-adds, because we still form every score and every weighted value. FlashAttention is not a FLOP reduction; it is a memory-traffic reduction. By keeping the score block resident in fast memory and streaming keys and values through once, it turns a bandwidth-bound operation into a compute-bound one — where the wall-clock speedups come from.
Why the same math matters on a CPU
The recurrence is usually explained for GPUs shuttling tiles between HBM and SRAM, but the algebra is hardware-agnostic and the benefit reappears on a CPU running a small language model. A CPU has the same shape of problem — a fast, tiny cache in front of comparatively slow main memory — and the N×N score matrix blows past the cache just as it blows past SRAM.
Applying the recurrence keeps the active score block small enough to sit in cache, so the CPU streams keys and values through once instead of re-reading a materialized matrix from RAM. The update is a tidy sequence of vectorizable reductions — a max and two scaled sums — that maps cleanly onto SIMD lanes. For CPU-hosted SLMs, where there is no HBM to hide behind and memory bandwidth binds at long context, not materializing S is often the single most effective attention optimization available.
Pitfalls in implementing the recurrence
Four mistakes recur when coding this by hand. Rescaling only l, not o~. Both accumulators are stated relative to the running max; forgetting the α·o~ multiply corrupts the output while the denominator still looks fine. Initializing m to a finite value. Use -∞ so the first block’s α is exactly 0; any finite value leaks a spurious weight into the empty state.
Masking after the block max. For causal or padded attention, set masked scores to -∞ before taking m_blk; masking afterward corrupts the reference. Normalizing too early. Divide by l only once, after the final block — normalizing per block breaks the rescaling algebra, because a normalized partial output can no longer be corrected by a single scalar multiply.
m, a running denominator l, and an unnormalized output o~; when a block raises the max, rescale the old state by α = exp(m_old - m_new) ≤ 1 and add the block’s terms against the new reference. Every exponent is non-positive, so nothing overflows; the loop invariant holds, so the final o~/l equals ordinary attention exactly — block size and order change only the intermediate numbers. The win is memory: the N×N score matrix is never materialized, dropping memory from O(N^2) to O(N) while the FLOP count stays the same. That is why it turns a bandwidth-bound kernel into a compute-bound one — and why the very same recurrence is the highest-leverage attention optimization for a small model running against a CPU cache.