The op that turns into `NaN` if you write it the way it's defined

Lesson 7 · numerical stability and streaming reduction · the prerequisite for FlashAttention

Softmax looks like the easiest op in the model: exponentiate, divide by the sum. Written exactly that way it produces NaN on real attention logits in fp16, and even once fixed it reads its input three times — which, per Lesson 2, is three times the cost that matters. Both problems have the same fix, and that fix is the core of FlashAttention.

Problem one: the definition overflows

Softmax as defined, and the value of x at which it breaks in each float format.

Listen first: softmax of x sub i is e to the x sub i, divided by the sum over j of e to the x sub j. The trouble is the exponential. In float thirty-two the largest representable number is about three-point-four times ten to the thirty-eighth, and the exponential passes that at an input of about eighty-eight-point-seven. In float sixteen the ceiling is only sixty-five thousand five hundred and four, which the exponential passes at an input of just eleven-point-one. Attention logits routinely exceed eleven. Once one term becomes infinity, the sum becomes infinity, and infinity divided by infinity is not a number — so a single large logit poisons the entire row.

softmax(x)_i = exp(x_i) / Σ_j exp(x_j)          # the textbook definition

exp overflows when:
    float32   x >  88.7      (max ≈ 3.40e38)
    float16   x >  11.1      (max = 65504)     ← attention logits pass this easily

one inf in the numerator → inf in the sum → inf/inf → NaN for the WHOLE row

The fix that costs nothing mathematically

Subtract the row maximum from every element before exponentiating. Multiply the top and bottom by exp(−m) and it cancels exactly — the result is identical in exact arithmetic, and in floating point the largest exponent is now exp(0) = 1, so overflow is impossible.

The max-subtraction identity: same value, no overflow, but an extra pass.

Listen first: let m be the maximum of x. Then softmax of x sub i equals e to the quantity x sub i minus m, divided by the sum of e to the x sub j minus m. Every exponent is now at most zero, so every exponential is at most one and underflow to zero is harmless. The catch is the cost: you must find m before you can exponentiate anything, and you must finish the sum before you can divide. That is three separate sweeps over the same data — read for the max, read for the sum, read again to normalise.

m = max(x)
softmax(x)_i = exp(x_i - m) / Σ_j exp(x_j - m)   # identical result, every exponent ≤ 0

but now:   pass 1 → find m
           pass 2 → exp and accumulate the sum
           pass 3 → divide
           THREE reads of x from memory

Problem two: passes are the cost

Softmax does a handful of FLOPs per element against several bytes of traffic — arithmetic intensity well under 1. It is squarely memory-bound, so by Lesson 2's rule its runtime is bytes ÷ bandwidth, and the number of passes is the runtime. Going from 3 passes to 2 is a 33% speedup, for free, forever.

The online trick: fuse the max and the sum

You do not actually need the final max before you start summing — you need to be able to correct the running sum whenever the max changes. If your running max jumps from m_old to m_new, every term already accumulated was scaled by exp(−m_old) and should have been scaled by exp(−m_new). So multiply the running sum by exp(m_old − m_new) and carry on.

Online softmax: one fused pass maintaining a running max and a corrected running sum.

Listen first: start with the running max at negative infinity and the running sum at zero. For each new element, the new max is the larger of the old max and this element. Then the new sum is the old sum times e to the old max minus the new max — that is the correction — plus e to this element minus the new max. When the max does not change, the correction factor is e to the zero, which is one, and nothing happens. When the max does change, the correction shrinks everything already accumulated onto the new scale. After one pass you have both the true max and the correctly scaled sum, and a second pass divides.

m = -inf ;  d = 0
for x_i in x:                                  # ONE fused pass
    m_new = max(m, x_i)
    d     = d * exp(m - m_new) + exp(x_i - m_new)      # rescale, then add
    m     = m_new

y_i = exp(x_i - m) / d                         # second pass: normalise

This is the whole idea behind FlashAttention. If softmax can be computed by streaming — accumulating a running max and a running, rescaled sum — then attention never has to materialise the full sequence-by-sequence score matrix in HBM. It can work tile by tile and correct as it goes. Lesson 7 is the arithmetic; Lesson 8 is the consequence.

Micro-world: break it, then fix it

Eight logits, three algorithms

Fixed input values scaled by the slider, so every run is reproducible. Step through the online version to watch its internal state correct itself.

Algorithm

Accumulate in

Logit scale ×1

Per-element internal state

Start naive, fp16, scale 1. It works. This is why the bug ships — at small logits the textbook formula is completely fine.
Drag the scale to 4 with naive + fp16. The largest logit passes 11.1, one exp returns infinity, and the whole row becomes NaN. Not the one big element — every element. That is the failure this lesson exists to reproduce.
Switch to max-subtracted, same settings. Fixed, and the outputs are the values you would have got. Note the pass count went from 2 to 3 — the fix is not free, it just costs bandwidth instead of correctness.
Switch to fp32 and push the scale to 30. Naive breaks again, just later. Wider floats move the cliff; they do not remove it.
Choose online and press "Step one element" repeatedly. Watch the running max and the correction factor. Most steps have a correction of exactly 1.000 — the rescale only fires when a new maximum arrives, which is what makes this cheap.

Check yourself

Cold-recall defense (one breath):
exp overflows above 11.1 in fp16 and 88.7 in fp32, and one inf makes the whole row NaN. Fix: subtract the row max — algebraically identical, every exponent ≤ 0 — but it costs a third pass (max, sum, normalise). Softmax is memory-bound, so passes are the runtime. The online version fuses max and sum into one pass by rescaling the running sum by exp(m_old − m_new) whenever the max grows. Same answer, 2 passes — and because it streams, attention never has to materialise the score matrix. That is FlashAttention.

Primary source

Read Online normalizer calculation for softmax (Milakov & Gimelshein, 2018) — it is 5 pages and the algorithm above is Algorithm 3. Then look at the Triton fused softmax tutorial and notice what "fused" means in bandwidth terms: the whole row stays in registers, so the passes cost nothing extra.

💬 I'm your teacher for this — ask me followups any time. Worth asking: "if the row doesn't fit in SRAM, what breaks?" — that constraint is exactly why FlashAttention tiles, and it is the bridge into the next lesson.

The op that turns into `NaN` if you write it the way it's defined

Problem one: the definition overflows

The fix that costs nothing mathematically

Problem two: passes are the cost

The online trick: fuse the max and the sum

Micro-world: break it, then fix it

Eight logits, three algorithms

Check yourself

Three questions on stability

Primary source

Read next

The op that turns into NaN if you write it the way it's defined

Problem one: the definition overflows

The fix that costs nothing mathematically

Problem two: passes are the cost

The online trick: fuse the max and the sum

Micro-world: break it, then fix it

Eight logits, three algorithms

Check yourself

Three questions on stability

Primary source

Read next

The op that turns into `NaN` if you write it the way it's defined