Online Softmax

Definition

Online (streaming) softmax computes a numerically-stable softmax in a single pass over the inputs, maintaining a running max and running denominator and applying a scalar rescale correction whenever a new element raises the max — instead of the standard two/three passes (find max → sum exponentials → normalize). Introduced by Milakov & Gimelshein (arXiv:1805.02867); it is the mathematical heart of FlashAttention, enabling attention to be tiled so the $N\times N$ score matrix never touches HBM.

Key math

Safe softmax over scores $s_1,\dots ,s_N$ subtracts the row max $m={\max }_j{s}_j$ for stability:
$o={\sum }_j\frac{e^{s_j-m}}{{\sum }_l{e}^{s_l-m}}{v}_j.$
Both $m$ and $\ell ={\sum }_l{e}^{s_l-m}$ depend on the whole row — the obstacle to streaming. Maintain, after processing blocks $\mathrm{1..}t$: running max $m^{(t)}$, running denominator ${\ell }^{(t)}$ (referred to $m^{(t)}$), and unnormalized output accumulator ${\mathbf{o}}^{(t)}$. Init $m^{(0)}=-\infty ,{\ell }^{(0)}=0,{\mathbf{o}}^{(0)}=0$. For block $B_t$ with local max ${\tilde{m}}^{(t)}$:
$m^{(t)}=\max (m^{(t-1)},{\tilde{m}}^{(t)}),{\alpha }^{(t)}=e^{m^{(t-1)}-m^{(t)}}\text{(rescale)},$
$\boxed{{\ell }^{(t)}={\alpha }^{(t)}{\ell }^{(t-1)}+\sum _{j\in B_t}{e}^{s_j-m^{(t)}},{\mathbf{o}}^{(t)}={\alpha }^{(t)}{\mathbf{o}}^{(t-1)}+\sum _{j\in B_t}{e}^{s_j-m^{(t)}}{v}_j}$
Final: $o={\mathbf{o}}^{(T)}/{\ell }^{(T)}$.

Why the correction ${\alpha }^{(t)}=e^{m^{(t-1)}-m^{(t)}}$? Old terms were referred to $m^{(t-1)}$; when the max jumps, each must be re-referenced: $e^{s_j-m^{(t-1)}}\cdot e^{m^{(t-1)}-m^{(t)}}=e^{s_j-m^{(t)}}$. Multiplying the whole accumulator by the scalar ${\alpha }^{(t)}$ does this to every stored term at once. If the max is unchanged, ${\alpha }^{(t)}=1$ (no rescale). Since $m^{(t)}\ge m^{(t-1)}$, always ${\alpha }^{(t)}\in (0,1]$ and $e^{s_j-m^{(t)}}\le 1$ — no overflow. The invariant ${\ell }^{(t)}={\sum }_{j\le B_t}{e}^{s_j-m^{(t)}}$ makes the result exact (bit-for-bit up to FP associativity), not an approximation.

Why it matters

Because $\ell$ and $\mathbf{o}$ are just running scalars/vectors, each block’s scores can be discarded immediately after updating state — which is why FlashAttention never writes the $N^2$ matrix to HBM. The same trick unlocks streaming/tiled softmax for very wide rows that exceed SMEM. FA-2 defers the rescale to a single end-of-loop normalization; FA-4 skips rescaling unless the max shift threatens stability (~10× fewer rescale ops).

Taught in

• 06-memory-wall-and-flashattention — §5.1 the full derivation, the invariant, and a worked-by-hand example.
• 07-triton-and-modern-kernels — §5 the fused-softmax kernel and its streaming ceiling.

See also

• flash-attention — the algorithm built on it
• io-aware-algorithms — why streaming enables IO-awareness
• kernel-fusion — softmax fused into attention
• shared-memory-tiling — the tiling online softmax makes exact
• memory-wall · gpu-memory-hierarchy
• gpu-systems-for-llms