Vol. XVI · Kernels and Efficiency Plate 3 of 8online softmax, numerical stability, running maximum, rescaling, one pass, blocks, FlashAttentionAC

The Online Softmax.

Find the maximum and the normaliser in one pass, correcting the running sum each time a bigger value appears.

Softmax divides e^{x_i} by the sum of all the exponentials. Computed naively, large logits overflow — e^{1000} is infinite in any float format — so the safe version first finds the maximum m and computes e^{x_i − m}: one pass for the maximum, one for the sum, one to normalise. When the inputs live in HBM, that is three reads of the whole vector, and in attention the vector is a whole row of the score matrix.

Milakov and Gimelshein (2018) showed the first two passes can be merged. Keep a running maximum m and a running sum l of e^{x − m}. When a new block arrives with a larger maximum m′, multiply the old sum by e^{m − m′} before adding the new block’s terms. At the end l is exactly the safe normaliser. The workshop runs it block by block on random logits and compares it with the two-pass answer and with the naive one.

Below: left, the running maximum as blocks arrive, with the true maximum. Right, the running normaliser on a log scale, with the exact value. Readouts compare the results.

What you will learn
Workshop
n logits drawn from a normal distribution and multiplied by the chosen scale, processed in blocks of the chosen size. The naive sum of e^{x} is computed in double precision and overflows past e^{709}; the two-pass and online results are compared as the normaliser relative to the true maximum.

One pass, exactly

With 64 logits in blocks of 8, the online softmax makes 8 block updates in one pass, rescaling the running sum 2 times as the maximum rises, and its normaliser agrees with the two-pass computation to a relative error of 4.4 × 10⁻¹⁶ — rounding, not approximation. With logits scaled up a thousand-fold, the naive sum of exponentials overflows (finite: 0) while the online version’s relative error is 0. With blocks of one it rescales 2 times, once for each new maximum, and the result is the same.

The point is not the saved pass for a single vector. It is that softmax becomes something that can be computed block by block while the blocks stream through fast memory, carrying only a maximum and a sum per row. In attention the output — a weighted sum of value vectors — can be rescaled the same way. That is the step that lets FlashAttention (plate XVI.4) tile attention without ever storing a row of scores, and lets ring attention (plate XIV.8) combine partial results from different devices.

The same rescaling identity, written for pairs of partial results (m₁, ℓ₁) and (m₂, ℓ₂), is associative — so partial softmaxes can be merged in any order, in a tree, or across devices.

Terms

Safe softmax
Subtracting the maximum before exponentiating, to avoid overflow.
Online softmax
Computing the maximum and normaliser in one pass with rescaling.
Rescaling
Multiplying a running sum by e^{m − m′} when the maximum rises to m′.
Normaliser
The sum of exponentials that softmax divides by.
Keeping a running maximum and a rescaled running sum gives softmax’s exact normaliser in one pass — the step that makes tiled attention possible.

Sources

  1. Milakov and Gimelshein, “Online normalizer calculation for softmax”, arXiv:1805.02867, 2018.
  2. Dao et al., “FlashAttention”, NeurIPS 2022, section 3.1.