kernel tutorials / online safe softmax

Online Safe Softmax

Softmax as everyone first writes it overflows; the safe version costs three passes over the row; the online version gets that back down to one.

The op

Softmax is a division by a sum

Softmax turns a row of scores into a row of probabilities: exponentiate every element, then divide each by the total, so the outputs are positive and sum to 1.

{softmax (x)}_{i} = \frac{e^{x_{i}}}{\sum_{j = 1}^{n} e^{x_{j}}}

No single output can be produced until the denominator is known, and the denominator needs the entire row.

Fix one · overflow

Safe softmax: subtract the max

The naive formula dies on real inputs because $e^{x}$ grows absurdly fast: $e^{90}$ is already infinity in fp32, and in fp16 the ceiling is all the way down at $e^{11}$ . One score past the ceiling and the row is inf/inf = NaN.

The fix rests on one observation: shifting every score by the same constant does not change the answer. Multiply the top and bottom of the fraction by $e^{- c}$ and it cancels:

\frac{e^{x_{i} - c}}{\sum_{j} e^{x_{j} - c}} = \frac{e^{- c} \cdot e^{x_{i}}}{e^{- c} \cdot \sum_{j} e^{x_{j}}} = {softmax (x)}_{i} for any c

Pick $c = m$ , the row maximum, and every exponent becomes $\leq 0$ : every exponential lands in $(0, 1]$ , and overflow is impossible by construction.

The price is that safe softmax is a three-pass algorithm, and the passes cannot overlap; each one needs the previous pass to have finished the whole row:

Pass 1 reads the row and finds the max $m$ .
Pass 2 reads the row again and accumulates $d = \sum_{i} e^{x_{i} - m}$ . It cannot start earlier: it needs $m$ .
Pass 3 reads the row a third time and emits $e^{x_{i} - m} / d$ .

The three sweeps, replayed on the row [2, 4, 3, 1]. Pass 2 sits idle until pass 1 has finished the entire row; each readout on the right only ticks while its own pass is moving.

Three passes are harmless when the row is 32 wide and sitting in L1. In attention the row is a row of scores: one entry per key, so as long as the sequence. At 128K context that is a 128K-wide row per query, produced block by block out of $Q K^{T}$ , ideally consumed the moment each block exists and never stored. "First scan everything to find the max" forbids exactly that.

Fix two · the passes

Online softmax: repair the sum as you go

The three-pass structure exists only because pass 2 is scared of using a stale max. Online softmax walks the row once, carrying two running values: $m$ , the max seen so far, and $d$ , the sum of exponentials relative to that max. When the next score $x$ arrives:

m^{'} = \max (m, x)

if d = \sum_{j = 1}^{k} e^{x_{j} - m} with m = \max_{j \leq k} x_{j}, m^{'} = \max (m, x_{k + 1}) then

The highlighted factor is the entire idea. Every term already inside $d$ has the form $e^{x_{j} - m}$ . If the max just rose from $m$ to $m^{'}$ , every one of those terms is now measured against the wrong reference, but they are all wrong by the same factor, because $e^{x_{j} - m^{'}} = e^{x_{j} - m} \cdot e^{m - m^{'}}$ . So one multiply by $e^{m - m^{'}}$ retroactively corrects the whole accumulated sum at once. No stored term is ever revisited. When the max did not rise, $m^{'} = m$ , the factor is $e^{0} = 1$ , and the update is a plain add.

The same walk over [2, 4, 3, 1], starting from $m = - \infty, d = 0$ :

Sees	New max $m^{'}$	Correction $e^{m - m^{'}}$	New sum $d^{'}$
2	2	–	$e^{0} = 1.000$
4	4 ↑ max rose	$e^{- 2} = 0.135$	$1.000 \cdot 0.135 + e^{0} = 1.135$
3	4	1	$1.135 + e^{- 1} = 1.503$
1	4	1	$1.503 + e^{- 3} = 1.553$

One pass, and the result is exactly the safe denominator: computing $\sum_{i} e^{x_{i} - 4}$ directly gives the same $1.553$ . The answer is identical, not an approximation.

The same row in one pass. When the cursor reaches 4 the max rises, and one multiply repairs everything accumulated so far. Nothing about the answer changes: only when the work is allowed to happen.

Where the last division went

For softmax by itself you still owe one more sweep at the end: the outputs need the final $d$ , so one pass becomes two, not one. The full win shows up in attention, which never needs the normalized scores individually: it only needs their weighted sum against $V$ . So keep a third running value, the accumulated output row $o$ , and rescale it by the same $e^{m - m^{'}}$ factor whenever the max rises, then divide by $d$ once at the very end. One pass total, and the score row is never stored anywhere. That is flash attention.

Why it scales across chips

The update is a merge, and merges make collectives

Look at the update rule again and notice it never cared that $x$ was a single score. Absorbing one element is just the special case of merging in a summary $(m_{2}, d_{2}) = (x, 1)$ . Two summaries of two chunks of the row merge the same way:

m = \max (m_{1}, m_{2})

if d_{1} = \sum_{j \in A} e^{x_{j} - m_{1}}, d_{2} = \sum_{j \in B} e^{x_{j} - m_{2}} and m = \max (m_{1}, m_{2}) then

This merge doesn't care what order the chunks arrive in, and merging three chunks gives the same answer however you pair them up. An operation with those two properties is exactly what a reduction collective needs: it is all_reduce with this merge in place of plain addition. So the row can be cut into chunks, each chunk summarized independently, and the summaries combined pairwise, in a tree, or around a ring.

That is the multichip payoff. When context gets long enough that the sequence dimension is fractured across chips, each chip holds a slice of the keys, so for any query, each chip can only score its own slice. The softmax denominator spans all of them. Gathering the raw score rows would move seq_len numbers per query per chip; merging online-softmax summaries moves two scalars per query (plus the running output row):

A 32K-key sequence fractured over four chips. Each chip attends to its own slice, then the tiny summaries merge across chips. The merge is associative, so it can run as a ring: this is ring attention.

Why plain tensor parallelism never needed this

In a tensor-parallel attention block, attention is split by heads, and a head never leaves its chip: each softmax row lives entirely in one place, so sdpa runs with no collective attached. Fracture the sequence instead of the heads and softmax becomes the op that spans chips. Online softmax is what keeps that affordable.

In the code

The same three symbols, as circular buffers

tt-metal's SDPA compute kernel is this page's math written in tiles: the running max and sum each get a circular buffer (cb_max_A/B, cb_sum_A/B), the correction factor $e^{m - m^{'}}$ has one of its own (cb_exp_max_diff), and the rescale step is a function named, fittingly, correction_block. The ring variant runs the same walk across chips, where each step's "next chunk" is another chip's slice of the sequence.

Source: state buffers, ttnn/cpp/ttnn/operations/transformer/sdpa/device/kernels/compute/sdpa.cpp:88 · rescale step, compute_common.hpp:706 (correction_block) · cross-chip walk, ring_joint_sdpa.cpp in the same directory

Recap

The whole page in four lines

Softmax's denominator needs the whole row before any output can be finished.
Safe: shifting all scores by a constant cancels out, so subtract the max: no overflow, but three dependent passes.
Online: carry $(m, d)$ through one pass; when the max rises, one multiply by $e^{m - m^{'}}$ repairs everything accumulated so far. Exact, not approximate.
Multichip: the update is an associative merge of $(m, d)$ summaries, so a row fractured across chips needs only a tiny reduction: flash attention within a chip, ring attention across them.