Online stable softmax
Derive the maximum, normalization-sum, and weighted-numerator updates for tiled softmax
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Max and sum reductions
- Ordinary stable softmax
- Exponent rules and elementary algebra
A transformer attention kernel receives one row of scores such as [1000, 999, 997]. Directly computing exp(1000) overflows ordinary floating-point formats, while subtracting the row maximum produces exp(0), exp(-1), and exp(-3), which remain finite. A small PyTorch or CUDA test can record the direct and stabilized outputs, and a correctness test can compare the result with torch.softmax in FP32.
A long attention row may be processed in tiles that fit in on-chip memory. After the first tile, the kernel retains the largest score seen and the sum of exponentials measured relative to that largest score. When a later tile contains a larger maximum, the kernel rescales the earlier sum before combining it. Engineers call this one-pass update online softmax. The running values can be logged in a CPU reference implementation and checked after every tile before they are embedded in a GPU kernel.
Stable softmax normally uses three logical phases: find the maximum, compute exponentials and their sum, then divide each exponential by that sum. For data larger than on-chip storage, a naive implementation rereads the row or materializes exponentials. Online softmax instead retains a running maximum and a normalization sum that are sufficient to merge the next block.
Online softmax maintains a running maximum and a normalization sum relative to that maximum. When a new block has a larger maximum, rescale the previous sum to the new maximum. The recurrence is numerically stable and supports tiled attention that processes one score block at a time.
The recurrence is easier to understand when the maximum is treated as a reference used to express all exponentials. When the reference changes, previous terms are still valid, but they must be converted to the new reference before they can be added to terms from the new block.
Why stable softmax subtracts the maximum
Softmax for a row x is exp(xi) divided by the sum of exp(xj). Direct exponentiation can overflow for a large positive input. Subtracting the row maximum does not change the result because the same factor exp(−m) appears in every numerator and in the denominator. The largest shifted input is zero, so its exponential is one and all other exponentials are at most one.
A conventional stable implementation first finds m, then computes l = Σ exp(xj−m), then writes exp(xi−m)/l. If the row does not fit in registers or shared memory, this requires another read or an intermediate. Online softmax reduces the stored state needed while tiles arrive.
Derive the update for two blocks
Initialize an empty prefix with m = −∞ and l = 0. Assume a nonempty processed prefix then has maximum m and stabilized sum l = Σprefix exp(x−m). A new nonempty block has maximum mb and local stabilized sum lb = Σblock exp(x−mb). The combined maximum is mnew = max(m,mb). The two sums cannot be added directly if m and mb differ because their exponential terms use different references. A fully masked row must use the operator's declared empty-row behavior instead of evaluating −∞−(−∞).
Convert the prefix terms by multiplying l by exp(m−mnew). Convert the block terms by multiplying lb by exp(mb−mnew). The combined sum is lnew = exp(m−mnew)l + exp(mb−mnew)lb. Every original exponential is now represented relative to mnew. Only m and l from the prefix are required; the prefix scores are not required to update the normalization statistics.
Carry a weighted value numerator for attention
Attention needs Σ exp(scorej−m)Vj in addition to the normalization sum. Maintain a vector numerator o for the processed key positions. When the maximum changes, multiply the old numerator by exp(m−mnew), exactly as for l. Add the new block's weighted V contribution after expressing it relative to mnew.
After all blocks, divide o by l. The output vector can remain on chip while key and value blocks arrive. Retaining the output vector is stronger than online standalone softmax: attention can consume each unnormalized score block immediately, so the kernel never needs to emit the full vector of softmax probabilities.
Masks, empty rows, and precision
A masked score acts like negative infinity because its exponential contribution should be zero. A partially masked tile can load or compute a finite placeholder and then force masked lanes to negative infinity before the maximum and sum. A fully masked row requires a specified result; otherwise m can remain negative infinity and the expression m−m produces an undefined value.
Use sufficient precision for m, l, and the output numerator. Lower-precision inputs often accumulate these states in FP32. Exponential approximations introduce additional error. Test extreme score ranges, all-masked and single-valid-element rows, odd tile boundaries, and both forward outputs and gradients against a stable reference. Numerical tolerance should be tied to datatype and downstream requirements.
Worked example: Merging two score blocks
The first block has maximum 10 and stabilized exponential sum 1.8. The second has maximum 12 and local stabilized sum 1.3.
- The combined maximum is 12. The first block's reference frame must move upward by two.
- Rescale its sum: exp(10−12) × 1.8, approximately 0.244. The second block is already expressed relative to 12, so it contributes 1.3.
- Add the compatible sums: 0.244 + 1.3 = 1.544. If the old and new sums had been added directly, the result 3.1 would incorrectly treat both blocks as though they used the same maximum.
- The combined normalization sum is about 1.544. Apply the same rescaling to any running weighted-value numerator from the first block before adding the second.
- Check the special case mb < m. The existing maximum remains m, so the old sum is unchanged and only the new block sum is scaled downward.
- Compare with a direct stable softmax over concatenated scores to verify the recurrence and inspect error in the target accumulator precision.
Conclusion: A two-number summary preserves exactly the information needed for normalization across tiles. The large score row need not exist as one stored object.
Common errors
- Keeping a running sum without rescaling when the maximum increases. Terms are then expressed in incompatible reference frames.
- Treating a fully masked row as an ordinary maximum. Negative infinity arithmetic requires an explicit semantic contract.
- Assuming the online recurrence eliminates all passes for standalone softmax. Output storage and row size determine whether inputs remain available on chip.
- Updating the normalization sum but not the attention numerator when the maximum changes. The final value then combines weights from different references.
Reference summary
Standard stable softmax finds a row maximum, computes exponentials relative to it, sums them, then normalizes. Initialize an empty running state with m = −∞, normalization sum l = 0, and weighted numerator o = 0. For a new tile, m₂ is its maximum, l₂ is its exponential sum relative to m₂, and o₂ is its weighted-value sum relative to m₂.
m′ = max(m,m₂); l′ = exp(m−m′)l + exp(m₂−m′)l₂o′ = exp(m−m′)o + exp(m₂−m′)o₂The same rescaling updates an accumulated weighted value. The running state lets attention process key and value tiles without storing the full score matrix. Edge cases—masked rows, negative infinity, empty tiles, and accumulator precision—must be specified rather than patched after profiling.
Knowledge check
Why must the previous running sum be multiplied by exp(m−m′) when a larger maximum appears?
Show the answer guide
- The running softmax state represents the sum of exponentials relative to the old maximum m: l = sum exp(x_i - m).
- When a new value raises the maximum to m-prime, the old terms must be expressed in the new coordinate system. Each old term changes from exp(x_i - m) to exp(x_i - m-prime) = exp(x_i - m) times exp(m - m-prime).
- The previous running sum must therefore be multiplied by exp(m - m-prime) before new exponentials are added. Without the factor, old and new terms use different normalization references and cannot be summed correctly.
- Because m-prime is at least m, the factor is at most one. The rescaling preserves the stable-softmax property and avoids reconstructing or storing the earlier logits.
Trace an online softmax maximum change
Implement scalar and vectorized online softmax for the row [1.0, 2.0, -3.0, 12.0, 11.0, 4.0]. Log the running maximum and running exponential sum after every element. Add a deliberately incorrect version that omits rescaling when 12.0 becomes the new maximum, then test all versions on 1,000 random rows containing occasional logits between 50 and 100.
Evidence to keep: Submit code, the complete state trace for the declared row, comparison against a stable two-pass softmax, maximum absolute and relative error across random tests, and a numeric explanation of the error introduced at the maximum change.
FlashAttention combines this online state with tiled Q, K, and V movement so exact attention avoids materializing its quadratic score and probability matrices.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.