Standard attention allocates an N×N score matrix in GPU memory, softmaxes it, then throws it away. For a 32K-token sequence that intermediate is a billion floats per head, and it never needs to exist. Online softmax is the recurrence that lets FlashAttention produce the exact same output while only ever holding a few blocks in fast memory at a time. Not an approximation. Bit-for-bit the same result, minus the floating-point reordering.
If you have ever wondered how FlashAttention gets a "no approximation" asterisk while sparse and low-rank attention variants do not, the answer is this one numerical trick.
TL;DR
- Online softmax computes softmax incrementally over blocks of the input, keeping a running max and running normalizer instead of scanning the whole row twice.
- It lets FlashAttention fuse the QKᵀ, softmax, and ·V steps into a single kernel that never writes the N×N attention matrix to HBM — memory drops from O(N²) to O(N).
- The result is numerically exact (up to float reassociation), unlike sparse or linear attention, because the recurrence rescales earlier partial sums when a new maximum appears.
- The win is memory bandwidth, not FLOPs: attention is memory-bound, so removing the N×N read/write is what produces the wall-clock speedup.
- The same running-normalizer state is what makes streaming and KV-cache decoding work — you can extend attention one token at a time without recomputing the row.
What problem does online softmax actually solve?
Softmax over a row of scores s is usually taught as three passes: find the max m, compute exp(s_i - m) and sum them into l, then divide each term by l. The max-subtraction is mandatory for numerical stability — exp(30) overflows fp16 instantly, so you shift by the max to keep exponents in range.
The catch in attention: you need the whole row of scores before you can start, because the max depends on every element. So the naive kernel materializes the full row (and, batched across queries, the full N×N matrix), reads it back, and only then normalizes. That matrix is the single largest tensor in a transformer forward pass, and it lives in HBM (high-bandwidth memory), the slow pool. Every byte written there and read back costs bandwidth you cannot spare.
Online softmax removes the "need the whole row first" constraint. It computes the correct normalized result while streaming the row in blocks, carrying just enough state to retroactively fix its earlier work when a larger score shows up later.
How does the online softmax recurrence work?
You maintain three running values as you consume blocks of keys/values: the running max m, the running sum of exponentials l, and the running weighted output O. When a new block arrives with its own local max, you rescale the accumulated state by exp(m_old - m_new) so everything is expressed relative to the new, larger maximum, then fold in the new block.
Here is the core recurrence in NumPy, deliberately unfused so you can see the state:
import numpy as np
def online_attention(Q, K, V, block=64):
# Q: (d,) single query for clarity
# K, V: (N, d)
N, d = K.shape
m = -np.inf # running max of scores
l = 0.0 # running sum of exp(scores - m)
O = np.zeros(d) # running unnormalized output
for start in range(0, N, block):
Kb = K[start:start+block]
Vb = V[start:start+block]
s = Kb @ Q # block scores, shape (block,)
m_new = max(m, s.max()) # updated running max
alpha = np.exp(m - m_new) # rescale factor for old state
p = np.exp(s - m_new) # new block weights
l = alpha * l + p.sum() # correct the normalizer
O = alpha * O + p @ Vb # correct the output, add new block
m = m_new
return O / l # normalize once, at the end
Two lines carry the whole idea. alpha = exp(m - m_new) is the correction factor: if a block introduces a bigger score, every previously accumulated exponential was computed against a smaller max and is now too large by exactly exp(m_old - m_new). Multiplying l and O by alpha retroactively re-bases them. The final division by l happens exactly once.
Compare it to the textbook version and the outputs match to floating-point tolerance:
def naive_attention(Q, K, V):
s = K @ Q
p = np.exp(s - s.max())
return (p @ V) / p.sum()
Q = np.random.randn(64)
K = np.random.randn(4096, 64)
V = np.random.randn(4096, 64)
print(np.max(np.abs(online_attention(Q, K, V) - naive_attention(Q, K, V))))
# ~1e-15
That 1e-15 is the entire claim. Online softmax is not "close enough" — it is the same computation with the summation reassociated, so it differs only by the rounding you would get from adding the same numbers in a different order.
Why is online softmax numerically stable?
Because every exponential is always evaluated against the current running maximum, never against a stale one. The instant a block raises the max, alpha shrinks the old accumulators before any new large term is added, so no partial sum is ever computed with a positive exponent. exp(s - m_new) has s ≤ m_new by construction, meaning every exponent is ≤ 0 and every term is in (0, 1]. You never overflow, and you never divide by an underflowed normalizer, because l accumulates the same bounded terms.
This is the property sparse and linear-attention methods sacrifice for speed. They change the math — dropping score entries or replacing softmax with a kernel feature map — so their outputs diverge from full attention and need retraining or accuracy caveats. Online softmax changes only the evaluation order, so it slots under any existing trained model with no quality loss.
Does online softmax approximate attention?
No. This is the most common misconception. FlashAttention is exact attention; online softmax is why. The kernel produces the identical tensor a naive softmax-then-matmul kernel would, differing only by floating-point non-associativity (the same reason (a + b) + c can differ from a + (b + c) in the last bit).
The confusion comes from lumping FlashAttention in with the "efficient attention" literature — Performer, Linformer, Longformer, BigBird — which do approximate. Those trade exactness for sub-quadratic FLOPs. FlashAttention keeps the quadratic FLOP count and attacks a different bottleneck: memory traffic. Different axis, no accuracy trade.
Why does removing the N×N matrix matter if the FLOPs are the same?
Because attention is memory-bound on modern GPUs, not compute-bound. The QKᵀ and ·V matmuls are relatively cheap; the expensive part is shoving the N×N intermediate out to HBM and reading it back for the softmax. HBM bandwidth is roughly an order of magnitude slower than on-chip SRAM, and the arithmetic intensity of attention is low enough that the hardware sits idle waiting on those transfers.
Online softmax lets you keep the running m, l, and O in registers/SRAM and process each block of K and V without ever writing scores to HBM. FlashAttention tiles Q, K, and V into blocks sized to fit SRAM, runs this recurrence in the inner loop, and writes only the final O back to global memory. Memory for the intermediate goes from O(N²) to O(N) — you hold a few blocks and the running state, not the full matrix. That is what turns a bandwidth-bound kernel into one that actually uses the GPU's FLOPs, and it is the source of the wall-clock speedups the FlashAttention papers report, plus the ability to train on much longer sequences without running out of memory.
The follow-up work (FlashAttention-2 and 3) is mostly about better work partitioning across warps and overlapping the matmul with the softmax rescaling — but the online-softmax core is unchanged. The recurrence above is still the beating heart of the kernel.
How does this connect to KV-cache decoding?
The same running m and l are exactly the state you need to extend attention by one token during autoregressive generation. When you append a new key/value to the cache, you are running one more iteration of the online-softmax loop: compute the new score, update the max, rescale, fold in the new value. This is why streaming attention and incremental decoding fall out naturally — the recurrence is inherently online in the sequence dimension. If you have read about "attention sinks" or long-context streaming, the mechanism that lets you drop or add tokens without recomputing the whole row is this normalizer bookkeeping.
It also explains a subtle serving detail: because the normalizer l is computed on the fly, you cannot cache softmax probabilities across requests — only the raw K and V. The probabilities depend on the full set of keys attended to, which changes every step. The KV cache stores pre-softmax state precisely because online softmax reconstructs the normalization cheaply each time.
When does online softmax not help you?
When N is small. For short sequences the N×N matrix fits comfortably in cache and the memory-traffic savings vanish, so a fused naive kernel can match or beat FlashAttention because it skips the rescaling overhead. The alpha multiply on every block is real work; it only pays off once the sequence is long enough that HBM traffic dominates. It also does nothing for the MLP blocks, which are often the larger share of total FLOPs — online softmax is an attention-kernel optimization, not a whole-model one. And if you are running a genuinely sub-quadratic model (state-space, linear attention), you are in a different regime entirely; there is no N×N matrix to avoid.
The bottom line
Online softmax is a two-pass-into-one-pass reformulation of softmax that carries a running max and running normalizer, rescaling earlier partial results by exp(m_old - m_new) whenever a larger score appears. It is what lets FlashAttention compute exact attention in a single fused kernel without ever materializing the N×N score matrix in HBM, cutting attention memory from O(N²) to O(N) and turning a bandwidth-bound operation into one the GPU can actually saturate. The output is numerically identical to naive attention up to floating-point reassociation — so unlike sparse or linear attention, it is a pure systems win with zero accuracy cost, and the same recurrence is what makes KV-cache decoding and long-context streaming work one token at a time.
Top comments (0)