Online softmax

Softmax looks like it needs the whole row at once: each entry is exp(x_i) divided by the sum of exp(x_j) over the row. If that were the only algorithm, you would have to store every score before normalizing. Online softmax keeps a running max and a running sum so you can eat the row in tiles.

The big idea

You are totaling a stack of receipts in batches. Whenever a new batch has a larger tax rate than you assumed, you rescale the old total before adding the new one. You never needed every receipt on the table.

The analogy breaks on floating point. Different addition order means results are not bit-identical to a naive kernel. They are mathematically the same softmax. Production stacks accept the tiny difference.

Walking through it

For each tile of scores:
  m_new  = max(m_old, max(tile))
  rescale the running numerator by exp(m_old - m_new)
  add exp(tile - m_new)
  update the running sum of exp
At the end, divide the numerator by the running sum

The only resident state is the tile plus (max, sum, partial output). You never needed the full row.

import math

xs = [1.0, 3.0, 2.0, 0.0]
m = max(xs)
naive = [math.exp(x - m) / sum(math.exp(y - m) for y in xs) for x in xs]

run_m, run_s, num = float("-inf"), 0.0, []
for tile in ([1.0, 3.0], [2.0, 0.0]):
    m_new = max([run_m, *tile])
    scale = math.exp(run_m - m_new) if run_m != float("-inf") else 0.0
    run_s = run_s * scale + sum(math.exp(x - m_new) for x in tile)
    num = [n * scale for n in num]
    num.extend(math.exp(x - m_new) for x in tile)
    run_m = m_new

online = [n / run_s for n in num]
print([round(v, 6) for v in naive])
print([round(v, 6) for v in online])

Output:

[0.087144, 0.643914, 0.236883, 0.032059]
[0.087144, 0.643914, 0.236883, 0.032059]

That identity is what lets FlashAttention stream K and V tiles without writing S to HBM.

Masks belong inside the tiled loop. Causal scores for future positions, and padded keys, must not enter the running sum, or they enter as -inf. Getting this wrong corrupts the distribution. It is worse than a slow kernel.

Common mistakes

• Thinking "online" means approximate. The algebra is exact softmax. Only the addition order changes.
• Forgetting to rescale when the max grows. Old numerators are then in the wrong units.
• Letting padded keys into the running sum. The model then attends to padding.

Check yourself

1. What three pieces of state does online softmax keep between tiles? Answer: The running max, the running sum of exp, and the partial weighted output.
2. Why rescales when a later tile has a bigger max? Answer: Earlier exp(x - m_old) values must become exp(x - m_new), which is a multiply by exp(m_old - m_new).

Key takeaways

• Softmax needs a row max and a sum, not the whole row resident at once.
• When the running max changes, rescale what you already have.
• The result matches naive softmax up to floating-point associativity.
• This is the math that makes tiled, IO-aware attention possible.

Further learning

Online normalizer calculation for softmaxThe Softmax function is ubiquitous in machine learning, multiple previous works suggested faster alternatives for it. In this paper we propose a way to compute classical Softmax with fewer memory accesses and hypothes…arXiv.org

Next, FlashAttention uses this inside an SRAM tile loop.