The Online Softmax Recurrence

TL;DR: Softmax over a row can be computed in one pass by carrying three things, the max so far $m$, the normalizer $\ell$ referenced to that max, and the unnormalized output, and multiplying the old $\ell$ and output by $e^{m_{\text{old}} - m_{\text{new}}}$ whenever the max rises. The result is exact, not approximate, and because the merge is associative it is what lets attention avoid writing the 128 MiB per-head score matrix of an 8K sequence to HBM.

The obstacle

Softmax attention for one query row is

Two things block streaming: the denominator needs every score before any output is final, and the exponentials overflow. In FP16, $e^{x}$ is infinite for $x > \ln 65{,}504 \approx 11.09$, and attention logits pass 11 routinely. The fix for overflow is the "safe softmax" identity: subtracting any constant from every $s_i$ cancels between numerator and denominator, so subtract the row max and every exponent is at most zero. That seems to make streaming worse, since now the max is needed first. Online softmax instead uses a reference point that can move, and a cheap way to re-reference what is already summed.

The recurrence

Hold a state $(m, \ell, a)$ for the keys seen so far: $m$ is their max, $\ell = \sum e^{s_i - m}$ and $a = \sum e^{s_i - m} v_i$. A new block of scores with max $\tilde m$ arrives:

and $a$ updates with the same factor. The proof: every term inside $\ell$ is $e^{s_i - m}$, and $e^{s_i - m} \cdot e^{m - m'} = e^{s_i - m'}$, which is exactly the term it would have been had $m'$ been known from the start. By induction the final state equals the full-row sums referenced to the true max, and $o = a / \ell$.

The more general form is the one that matters for kernels: two partial states for disjoint key sets merge into one.

The merge is associative and commutative in exact arithmetic, so blocks can be combined left to right (FlashAttention's inner loop), in a tree, or in any order (split-KV decoding). In floating point the order changes the rounding, which is the root of one class of batch-dependent nondeterminism.

Worked example

Row of four scores, two blocks of two: [3, 1] then [5, 2]

after block 1   m = 3
                l = e^(3-3) + e^(1-3) = 1 + 0.1353         = 1.1353
block 2 arrives m' = max(3, 5) = 5
                scale = e^(3-5)                             = 0.1353
                l' = 1.1353 x 0.1353 + e^(5-5) + e^(2-5)
                   = 0.1537 + 1 + 0.0498                    = 1.2034
direct check    e^(3-5) + e^(1-5) + e^(5-5) + e^(2-5)       = 1.2034
sanity: no exponent ever exceeded 0, so no intermediate exceeded 1 per term

The merge below was checked against a direct softmax on 1,000 keys (score standard deviation 30, uneven blocks including one of length 1, one block fully masked). Sequential, reversed, shuffled and tree orders all match to $1.3 \times 10^{-15}$.

import numpy as np

def state(scores, values):
    """Partial softmax state (m, l, acc) for one block of keys."""
    if scores.size == 0 or np.all(np.isneginf(scores)):
        return -np.inf, 0.0, np.zeros(values.shape[1])
    m = scores.max()
    p = np.exp(scores - m)
    return m, p.sum(), p @ values

def merge(a, b):
    (ma, la, xa), (mb, lb, xb) = a, b
    m = max(ma, mb)
    if np.isneginf(m):                      # both blocks fully masked
        return a
    sa, sb = np.exp(ma - m), np.exp(mb - m)  # exp(-inf) = 0 for an empty side
    return m, la * sa + lb * sb, xa * sa + xb * sb

What it buys

Without the recurrence, attention materializes $S = QK^\top$. At 8,192 tokens one head's score matrix is $8{,}192^2 \times 2$ bytes, 128 MiB in BF16, written once and read back for the softmax and again for $PV$. With it, a kernel loads a tile of K and V into shared memory, updates the per-row state in registers, and moves on. The FLOPs do not change. The score matrix never leaves the SM, and since the unfused softmax pass does a few FLOPs per byte against an H100 SXM ridge of 989 / 3.35 = 295 FLOP/byte (dense BF16), bytes were the cost.

The per-row state is two FP32 scalars plus a $d$-wide accumulator, and the log-sum-exp $m + \ln \ell$ summarizes the whole row in one number, which is all FlashAttention's backward pass stores.

Where it breaks

Initialize $m$ to $-\infty$, never 0. With $m = 0$, a row whose scores are all negative keeps a reference max that never occurred, and the exponentials can underflow to zero. A row whose every key is masked has $m = -\infty$ and $e^{-\infty - (-\infty)}$ is NaN, so kernels special-case it (the guard in merge above). The statistics must stay in FP32 even when Q, K and V are FP8 or BF16, because $\ell$ is a long sum of terms in $(0, 1]$ and loses low bits fast in a narrow format. Forgetting to rescale the accumulator along with $\ell$ gives an output off by a factor that depends on where the row's max happens to fall, and tests that put the max in the first block still pass.

What interviewers listen for

They want the rescale factor written correctly and the one-line exactness argument, not "it approximates softmax". Strong candidates say the saving is HBM traffic, not FLOPs, and generalize from a running loop to a merge of two states, because that is the step that unlocks split-KV decoding and sequence-parallel attention. A sharp follow-up is the masked row; answering it without prompting reads as someone who has written the kernel.

Say it in 30 seconds

Online softmax keeps a running max, a normalizer referenced to that max and an unnormalized output, and when a new block raises the max it multiplies the old normalizer and output by $e^{m_{\text{old}} - m_{\text{new}}}$. That re-references every old term exactly. Because the merge is associative, attention streams K and V through shared memory without writing the 128 MiB per-head score matrix of an 8K sequence, and long contexts can be split across thread blocks and merged at the end.