Softmax numerics architecture — max subtract, fp32 accumulate, fusion, online

What softmax is, and the range where it breaks

Softmax maps a vector of real logits x: [N] to a probability vector: p_i = exp(x_i) / Σ_j exp(x_j), with every p_i > 0 and Σ_i p_i = 1. Mathematically it is well behaved on all of R^N. In floating point it is not, because exp leaves the representable range almost immediately.

The dangerous quantity is the largest logit. exp(x) overflows when x exceeds log(FLT_MAX): about 11.09 in fp16 (max 65504), about 88.7 in fp32 and bf16, which share an 8-bit exponent. Attention scores of 15 or 30 are routine, so a literal fp16 softmax overflows on real data — and inf / inf is NaN, not a large number but a poison value that propagates through the whole batch. Underflow is the mirror image: exp(x) flushes to zero below about -87 in fp32 and -9.7 in fp16. The two failures are not symmetric, and that asymmetry is what the standard fix exploits.

The max-subtraction identity, and why it is exact

Shift every logit by the same constant c and the answer does not change. The proof is one line of cancellation:

softmax(x - c)_i = exp(x_i - c) / Σ_j exp(x_j - c)
                 = [exp(-c) · exp(x_i)] / [exp(-c) · Σ_j exp(x_j)]
                 = exp(x_i) / Σ_j exp(x_j)  =  softmax(x)_i

This is exact in real arithmetic — softmax is invariant to adding any constant to all its inputs, the same statement as ‘logits are defined only up to a shift.’ Choosing c = max(x) is the choice that makes the computation safe. Every shifted logit is then ≤ 0, so every exp lands in (0, 1] and overflow becomes impossible in any format. The denominator is protected too: the argmax term contributes exactly 1.0, so the sum lies in [1, N] and can never be zero. Terms far below the max still underflow — but such a term was, by construction, a probability below the format’s smallest normal, and dropping it sits well under the rounding error you already accept.

A worked example: naive fp16 versus the stable form

Four logits, evaluated both ways in fp16.

x            = [ 20.0,   18.0,   17.0,   -5.0 ]

naive:  exp(20.0) = 4.85e8  →  > 65504  →  inf
        sum = inf           →  inf/inf    →  [NaN, NaN, NaN, NaN]

stable: x - max(x) = [ 0.0,  -2.0,  -3.0,  -25.0 ]
        exp        = [ 1.0,  0.13534,  0.049787,  1.389e-11 → 0 ]
        sum        = 1.185122
        p          = [ 0.84380,  0.11420,  0.04201,  ~0 ]

The naive path does not degrade gracefully; it returns nothing usable, and inside a training step that NaN reaches the optimizer and destroys the weights. The stable path loses exactly one term, the fourth, whose true value is 1.17e-11 — against fp16’s unit roundoff of 2^-11 ≈ 4.9e-4, seven orders of magnitude below the format’s noise floor. Overflow is fatal, underflow is free, which is why we shift down to the max rather than up toward it.

Why the accumulator is fp32 even when the inputs are not

Max subtraction fixes range. It does nothing for precision, and the sum is where precision dies. bf16 carries 8 significand bits, a unit roundoff of u = 2^-8 ≈ 3.9e-3, against fp32’s 2^-24 ≈ 6e-8.

The failure mode is stagnation. Sequential summation of n terms has a worst-case relative error bounded by roughly n · u (and about √n · u under a random-walk model of the rounding). With n = 4096 and bf16, n · u ≈ 16 — the bound is vacuous. Concretely: bf16 values in [256, 512) are spaced 2.0 apart, so once the running sum passes 256, adding a term of size 1 rounds straight back to 256 and the term vanishes. Since every shifted exponential is ≤ 1, a long row’s entire tail can disappear into a stalled accumulator. The fix is free: keep tensors in bf16, widen the accumulator and running max to fp32. Pairwise summation, which vectorized kernels do anyway, cuts the growth from n to log n.

Online softmax: one recurrence, one fewer pass

The textbook implementation makes three passes over the row: find the max, sum the exponentials, divide. Fine for a vector in cache, terrible for an attention row you would rather never store. The online formulation collapses the first two passes into one by carrying a running max m and running sum d together, rescaling d whenever a larger max arrives:

init:   m = -inf,   d = 0

per block B (local max m_B, local sum d_B = Σ_{i∈B} exp(x_i - m_B)):
    m' = max(m, m_B)
    d' = d · exp(m - m') + d_B · exp(m_B - m')

finish: p_i = exp(x_i - m) / d

Correctness follows from the same shift identity: d is always the sum of exp(x_j - m) over everything seen, and changing m means multiplying by exp(m_old - m_new). Both correction factors are ≤ 1, so rescaling can never overflow. The price is one extra exp per block boundary — arithmetic traded for a memory pass, the right trade on any modern machine.

From online softmax to tiled attention

The recurrence is what makes FlashAttention-style kernels possible. Standard attention computes S = QK^T / √d_k as an explicit [N, N] matrix, softmaxes it, then multiplies by V. At N = 4096 that matrix is 16.8M entries — 33.5 MB in fp16, written and read back once per head, per layer.

Tiled attention never materializes it. It walks blocks of K and V, computes each score tile on chip, and folds it into three running quantities: m, d, and an unnormalized output accumulator O, which takes the same correction factor as the sum:

O' = O · exp(m - m') + exp(S_B - m') · V_B      # then O / d at the end

The result matches a reference softmax to rounding, but memory traffic drops from O(N^2) to O(N · d). Long context stops being a memory problem and becomes a compute problem — entirely because softmax can be written as a rescalable recurrence.

Log-sum-exp: log-probabilities without probabilities

Training and evaluation want log p, not p, and the worst possible route is to compute probabilities and take their logarithm — any underflowed entry becomes log(0) = -inf. Go straight to log space with log-sum-exp:

LSE(x)            = m + log( Σ_j exp(x_j - m) ),   m = max(x)
log_softmax(x)_i  = x_i - LSE(x)
CE(x, target t)   = LSE(x) - x_t

Every step is safe: the exponentials are bounded by 1, the sum is at least 1 so the logarithm is defined and non-negative, and the final subtraction is ordinary arithmetic. A confidently wrong prediction yields a large finite loss instead of inf — the difference between a run that recovers and one that diverges. It is also why cross-entropy should be handed logits, never a softmax output, and why frameworks expose log_softmax and cross_entropy_with_logits as primitives.

The Jacobian, and the backward pass that never divides

Softmax couples every output to every input, so its Jacobian is dense: ∂p_i/∂x_j = p_i(δ_ij - p_j), i.e. J = diag(p) - p p^T. You never build it. For upstream gradient g, the vector-Jacobian product collapses to O(N): ∂L/∂x = p · (g - g·p), elementwise — one dot product, one pass.

Fusing with cross-entropy is better still. Differentiating L = LSE(x) - x_t gives ∂L/∂x_i = p_i - y_i, with y the one-hot target. Every component lies in [-1, 1], no separate softmax backward is needed, and it never divides. The unfused chain does: cross-entropy’s own gradient is -y_i / p_i, so an underflowed p_t gives inf, which the Jacobian multiply turns into NaN. The fused form is algebraically identical and numerically bulletproof — the highest-value fusion in the stack.

Fusion and the memory-traffic account

Softmax does almost no arithmetic per byte: for a row of N elements it is a handful of flops per element against 2N bytes read and 2N written in bf16. Arithmetic intensity near 1 flop/byte puts it far on the memory-bound side of any roofline. That one fact dictates the strategy: minimize passes, not flops.

A three-pass softmax reads the row three times and writes once; online softmax reads it twice. Fusing with the op that produced the logits and the one that consumes them — the matmuls in attention, the loss in the output head — makes it effectively one, because the scores are consumed in registers before they ever reach memory. That is why the recurrence earns its extra exponentials, and why an unfused softmax between two fused matmuls is often the largest stall in an otherwise well-tuned layer.

On CPU, and the pitfalls worth naming

On CPU the bottleneck shifts: there is no hardware transcendental unit, so exp is a polynomial. Vectorized kernels compute exp(z) = 2^(z · log2 e), split the argument into integer and fractional parts, build 2^n by writing the exponent field directly, and evaluate a degree-5 minimax polynomial on the remainder — roughly ten operations across 8 or 16 fp32 lanes. A standard trick folds log2(e) / √d_k into one scale constant so the kernel calls exp2 and saves a multiply per element.

Three pitfalls recur. Taking the max with a horizontal reduction inside the loop rather than lanewise with one reduction at the end. Holding the running max or sum in the compute dtype, which reintroduces stagnation. And masking by multiplying weights by zero instead of adding -inf to logits before the max — that breaks normalization, and a fully masked row becomes 0/0. Clamp those rows explicitly.