Online softmax

The recurrence

Online softmax keeps two running values per row, the maximum m and the sum l, and folds each new block into them:

m_new = tl.maximum(m_i, tl.max(x, axis=1))
alpha = tl.exp2(m_i - m_new)
p = tl.exp2(x - m_new[:, None])
l_i = l_i * alpha + tl.sum(p, axis=1)
acc = acc * alpha[:, None] + tl.dot(p, v)
m_i = m_new

When a block raises the maximum, every term already accumulated was weighted against the old one. Multiplying l and acc by exp(m_old - m_new) puts them on the new scale, and the new block contributes on that same scale. One pass, constant memory per row.

Why the values are preserved

Each rescale multiplies numerator and denominator by the same factor, so the ratio is unchanged, and in exact arithmetic the result equals the two-pass softmax. In floating point it is not bitwise equal: terms are added in a different order and rescaled by a rounded factor, so the correctness check uses a tolerance justified by the dtype.

Initialise honestly

Start from m = -inf and l = 0. On the first non-empty block alpha is exp(-inf - m_new), which is zero, so the stale sum and accumulator vanish and the block stands alone. The attention tutorial starts l at 1.0 instead, which works for the same reason. What does not work is an empty first block: -inf - -inf is NaN, and the row is poisoned before it starts. The attention tutorial prunes causal configurations that can produce an entirely masked diagonal tile, which is the same failure arriving through the mask sentinel rather than through -inf.

Things that break

Rescaling l but not acc is the classic bug: the output is right for a single tile and drifts as the context grows, so a short-sequence test passes. Mixing bases is the second, since a log2(e) factor folded into the logits only makes sense if every exponential in the loop is base two, alpha included. And recomputing the maximum from the block instead of carrying it forward gives two passes with extra steps. These snippets follow Triton 3.9; confirm tl.exp2 against triton.__version__, because older code spells it tl.math.exp2.