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.

Advertisement

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 m

Choosing 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.

Advertisement

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~ / l

Initializing 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   ✓ identical

The 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.