Online Softmax — the mathematical heart of FlashAttention

If you understand this one derivation, everything about FlashAttention, FlashDecoding, and FlashInfer becomes mechanical. If you don’t, everything else in Phase 3 is spellcasting. Do this on paper first, then read on.

Original paper: Milakov & Gimelshein, Online normalizer calculation for softmax, arXiv <phone_number_or_numberic_id_or_random_id_50>. 2018. Short, brilliant, foundational.

1. Standard (offline) softmax — 3 passes

Given vector x = [x_1, ..., x_N], softmax is:

y_i = exp(x_i) / sum_j exp(x_j)

Numerically-stable version (subtract the max to avoid overflow):

m    = max_j x_j                    # pass 1
d    = sum_j exp(x_j - m)           # pass 2
y_i  = exp(x_i - m) / d             # pass 3

3 passes over x, each touching all N elements. Memory traffic is 3N reads. For attention where x is one row of the score matrix S = QK^T, this means the N×N score matrix must live somewhere addressable — that “somewhere” is HBM, and that HBM traffic is what FlashAttention refuses to pay.

2. The online recurrence (1 pass over the vector, in tile order)

Define, over the first k elements seen so far:

  • m_k = max(x_1, ..., x_k) — running max

  • d_k = sum_{j≤k} exp(x_j - m_k) — running denominator, referenced to the current max

Recurrence when the (k+1)-th element x_{k+1} arrives:

m_{k+1} = max(m_k, x_{k+1})
d_{k+1} = d_k * exp(m_k - m_{k+1})  +  exp(x_{k+1} - m_{k+1})

The key move: when the max grows, rescale the previous denominator by exp(m_old - m_new). That rescale is the entire trick.

At the end, y_i = exp(x_i - m_N) / d_N.

3. Sanity check with tiny numbers

Let x = [1.0, 3.0, 2.0].

Offline:

  • max = 3, exponents = [exp(-2), exp(0), exp(-1)] = [0.1353, 1.0, 0.3679]

  • sum = 1.5032

  • y = [0.0900, 0.6652, 0.2447]

Online:

Init: m_0 = -∞, d_0 = 0.

Step 1 (x=1.0):

  • m_1 = max(-∞, 1.0) = 1.0

  • d_1 = 0 * exp(-∞ - 1.0) + exp(1.0 - 1.0) = 0 + 1 = 1

Step 2 (x=3.0):

  • m_2 = max(1.0, 3.0) = 3.0

  • d_2 = 1 * exp(1.0 - 3.0) + exp(3.0 - 3.0) = exp(-2) + 1 = 0.1353 + 1 = 1.1353

Step 3 (x=2.0):

  • m_3 = max(3.0, 2.0) = 3.0

  • d_3 = 1.1353 * exp(3.0 - 3.0) + exp(2.0 - 3.0) = 1.1353 + 0.3679 = 1.5032

Matches. Now normalize: y_i = exp(x_i - 3.0) / 1.5032 → same answer as offline.

This is worth doing by hand once. Then again with a 4-element sequence and 2 tiles of size 2 to see the tile version.

4. Tile-level (block) recurrence — the FA form

Attention doesn’t stream elements one at a time; it processes them in tiles (block of K columns at a time). So we need the block form.

Let x be split into tiles T_1, T_2, .... For tile T_r with local max m^r_local = max(x in T_r) and local sum d^r_local = sum exp(x - m^r_local) for x in T_r:

Running state before tile r: m_{r-1}, d_{r-1}. After tile r:

m_r = max(m_{r-1}, m^r_local)
d_r = d_{r-1} * exp(m_{r-1} - m_r)  +  d^r_local * exp(m^r_local - m_r)

The formula is identical to the scalar recurrence, promoted to tile granularity.

5. The attention promotion — not just softmax, but softmax(QK^T)V

Attention is:

A_i = sum_j softmax(QK^T)_ij * V_j = (sum_j exp(s_ij - m) V_j) / d

So alongside m and d, we maintain a running output vector o (dimension d_head). When the max updates, we rescale o the same way we rescale d:

Per Q-tile, iterating over K/V-tiles:

State per Q-row: (m, d, o) where o ∈ R^{d_head}.

On arrival of tile r (produces score row s^r and reads V^r):

m_r     = max(m_{r-1}, max(s^r))
alpha   = exp(m_{r-1} - m_r)            # rescale factor for old state
beta    = exp(s^r    - m_r)             # unnormalized weights for new tile (vector)
d_r     = d_{r-1} * alpha  +  sum(beta)
o_r     = o_{r-1} * alpha  +  beta @ V^r

At the very end, output_i = o_final / d_final.

This is FlashAttention forward, distilled. Every FA-family kernel is this recurrence with (a) tile-shape optimizations, (b) masking, (c) precision choices, (d) async pipelining. Master the recurrence and the code becomes readable.

6. Why this saves memory traffic — the arithmetic

Standard attention:

  • Compute S = QK^T, materialize N×N in HBM: N² * 2 B bytes (bf16) written + read.

  • Softmax over S: 3 passes = 3 * N² * 2 B.

  • Multiply by V: N² * 2 B read + N*d * 2 B written.

  • Total HBM traffic: O(N² · d + N²) — dominated by the score matrix.

FlashAttention:

  • Never materializes S. Iterates Q-tiles × KV-tiles. Each tile fits in SMEM (size M).

  • For a single Q-tile: reads all of K and V once. Total HBM reads per Q-tile: N·d·2 B (for K) + N·d·2 B (for V) = 2 N d bytes.

  • Total Q-tiles: N / tile_rows.

  • Total HBM traffic: O(N² · d² / M).

Ratio (FA / naïve): d² / (M · d + M) ≈ d / M. With d=128, M ~ 100 KB / 2 B = 50k elements, ratio ≈ 128/50000 ≈ 0.003 — ~300× less HBM traffic. That’s the source of the speedup.

Read: FA1 paper §3 has this arithmetic laid out with exact constants.

7. What can go wrong (the numerics rabbit hole)

  • exp(m_{r-1} - m_r) when m_{r-1} = -∞: need a sentinel. In practice initialize m = -large_negative (say -1e30).

  • d accumulate in bf16: DON’T. Accumulate d, o, and intermediate exps in fp32 even if inputs and outputs are bf16 / fp16. FA papers and Triton tutorial both do this. See 06_numerics_discipline.md.

  • Fully-masked rows (all -inf scores): d will be 0, giving 0/0 in the final divide. Guard with d = max(d, tiny) or by post-hoc zeroing the row.

  • Score = -inf then + finite: the rescale exp(m_old - m_new) * d_old will vanish correctly, but only if you do the max update before computing the new alpha. Order matters.

8. What to actually do this week

  1. Derive the online softmax recurrence on paper. Twice.

  2. Do the tiny numerical example above by hand.

  3. Do the attention promotion on paper: state = (m, d, o), write the update rule.

  4. Implement online softmax as a pure Python function; verify against torch.softmax.

  5. Extend it to online attention (still pure Python, tiny shapes): verify against torch.nn.functional.scaled_dot_product_attention.

Only then move to 04_writing_fa_in_triton.md.

References