Will Ahmad

Linear algebra for LLMsPart 9: Numbers on real hardware

Stable and online softmax

7 min

exp overflows fast

ex>65,504 (FP16 max) once x>11.1\displaystyle e^{x} > 65{,}504 \text{ (FP16 max) once } x > 11.1
ex>3.4×1038 (FP32 max) once x>88.7\displaystyle e^{x} > 3.4 \times 10^{38} \text{ (FP32 max) once } x > 88.7

Attention scores and logits easily pass 11. Computed naively, the softmax numerator becomes infinity and the result becomes NaN.

Subtract the max (again)

m=max⁡jzj,softmax⁡(z)i=ezi−m∑jezj−m\displaystyle m = \max_j z_j, \qquad \operatorname{softmax}(z)_i = \frac{e^{z_i - m}}{\sum_j e^{z_j - m}}

Part 5’s shift invariance, now as a rule for hardware: after subtracting mm, every exponent is ≤0\le 0, so every ezj−me^{z_j - m} is in (0,1](0, 1] and the sum is between 1 and the number of entries. Nothing overflows.

Write

Stable softmax: subtract m=max⁡zm = \max z before exp⁡\exp.

Log-sum-exp

LSE⁡(z)=log⁡∑jezj=m+log⁡∑jezj−m\displaystyle \operatorname{LSE}(z) = \log \textstyle\sum_j e^{z_j} = m + \log \sum_j e^{z_j - m}
log⁡softmax⁡(z)i=zi−LSE⁡(z)\displaystyle \log \operatorname{softmax}(z)_i = z_i - \operatorname{LSE}(z)

Cross-entropy needs the log of a softmax probability. Computing it as zi−LSE⁡(z)z_i - \operatorname{LSE}(z) avoids taking the log of a tiny number that rounded to 0. That’s what F.cross_entropy does inside: pass it logits, not probabilities.

Write

log⁡softmax⁡(z)i=zi−LSE⁡(z)\log \operatorname{softmax}(z)_i = z_i - \operatorname{LSE}(z), with LSE⁡=m+log⁡∑ezj−m\operatorname{LSE} = m + \log \sum e^{z_j - m}.

The problem: rows don't fit

S=QK⊤ ⁣:  T×T per head\displaystyle S = QK^\top\!: \; T \times T \text{ per head}

At T=8192T = 8192 that’s 67 million scores per head, too big for fast on-chip memory. We’d like to process KK and VV in blocks, but softmax needs the max and the sum over the whole row before it can normalize anything.

Online softmax: keep a running max and sum

m′=max⁡(m,max⁡j∈blockzj)\displaystyle m' = \max(m, \max_{j \in \text{block}} z_j)
ℓ′=ℓ em−m′+∑j∈blockezj−m′\displaystyle \ell' = \ell \, e^{m - m'} + \textstyle\sum_{j \in \text{block}} e^{z_j - m'}

Read the row one block at a time. When a new block raises the max, the old sum was computed against a smaller mm, so multiply it by em−m′e^{m - m'} to bring it to the new reference. At the end, ℓ\ell is exactly ∑jezj−m\sum_j e^{z_j - m} over the whole row.

Write

m′=max⁡(m,block max)m' = \max(m, \text{block max}); ℓ′=ℓem−m′+∑blockezj−m′\ell' = \ell e^{m - m'} + \sum_{\text{block}} e^{z_j - m'}.

On numbers: [1, 3 | 2, 5]

block 1: m=3,ℓ=e−2+e0=1.1353\displaystyle \text{block 1: } m = 3, \quad \ell = e^{-2} + e^{0} = 1.1353
block 2: m′=5,ℓ′=1.1353 e−2+e−3+e0=1.2034\displaystyle \text{block 2: } m' = 5, \quad \ell' = 1.1353\, e^{-2} + e^{-3} + e^{0} = 1.2034

Check against the whole row at once: e−4+e−2+e−3+e0=0.0183+0.1353+0.0498+1=1.2034e^{-4} + e^{-2} + e^{-3} + e^{0} = 0.0183 + 0.1353 + 0.0498 + 1 = 1.2034. Same answer, without ever holding all four scores.

✏️ Your turn

Run it yourself

z=[0,2∣4]\displaystyle z = [0, 2 \mid 4]

Process the two blocks in order. Write mm and ℓ\ell after each, then the three softmax weights.

Work it out on paper, then go to the next card to check.

Answer

block 1: m=2,ℓ=e−2+1=1.1353\displaystyle \text{block 1: } m = 2, \quad \ell = e^{-2} + 1 = 1.1353
block 2: m′=4,ℓ′=1.1353 e−2+1=1.1537\displaystyle \text{block 2: } m' = 4, \quad \ell' = 1.1353\, e^{-2} + 1 = 1.1537
softmax⁡=[e−4,e−2,e0]/1.1537≈[0.016,0.117,0.867]\displaystyle \operatorname{softmax} = [e^{-4}, e^{-2}, e^{0}] / 1.1537 \approx [0.016, 0.117, 0.867]

The output gets the same rescale

o′=o em−m′+∑j∈blockezj−m′ vj,output=o/ℓ\displaystyle o' = o\, e^{m - m'} + \textstyle\sum_{j \in \text{block}} e^{z_j - m'}\, v_j, \qquad \text{output} = o / \ell

Attention also needs the weighted sum of values. Keep an unnormalized running output oo, rescale it whenever the max moves, and divide by ℓ\ell once at the end. This is the core of FlashAttention (Dao et al., 2022): exact attention, computed tile by tile in fast memory, never writing the T×TT \times T matrix to GPU memory.

Write

Running output: rescale by em−m′e^{m - m'}, add the block’s weighted values, divide by ℓ\ell at the end.

📌 Recap

Stable and online softmax in one card

  • Subtract the max before exp⁡\exp; use log-sum-exp for log-probabilities.
  • Online softmax: running mm and ℓ\ell, rescaling by em−m′e^{m - m'} when the max grows.
  • Rescale the running output the same way: FlashAttention.

Write

The three lines above.

Sources

  • Milakov and Gimelshein, “Online normalizer calculation for softmax” (2018), arXiv:1805.02867.
  • Dao, Fu, Ermon, Rudra and Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” (2022), arXiv:2205.14135.