Raphael Friedmann
← The log

Online softmax, and the register that can't move

The deep version of FlashAttention on consumer Blackwell: how online softmax avoids the giant score matrix, and why keeping the output in registers forces raw mma.sync over WMMA.

17 June 2026 Expert

The kernels post called register-resident FlashAttention “the defining property” and kept moving. This is the slow version of that sentence: what online softmax actually does, why the output accumulator has to live in registers, and why that one constraint rules out the convenient high-level matrix API and forces raw mma.sync.

The matrix you don’t want to build

Naive attention forms the full score matrix, every query against every key, softmaxes it row by row, and multiplies by the values. The matrix is the problem. For a long sequence it is enormous, and materialising it means writing and reading a quadratic amount of data the kernel does not actually need to keep.

Online softmax

FlashAttention never builds it. It streams the keys and values past in tiles and keeps two running numbers per row: the maximum score seen so far, and the running sum of exponentials.

Online softmax: one tile at a time, no giant matrix
KV tile 1KV tile 2KV tile 3KV tile 4running max + sumrescale O by exp(m_old - m_new)output Othe full N x N score matrix is never built
Instead of forming the whole score matrix and softmaxing it, the kernel streams the keys and values in tiles. It keeps a running maximum and sum, and rescales the output each time a bigger value appears. The quadratic matrix never exists.

The trick is the rescale. When a new tile contains a larger score than anything before it, the running maximum jumps, and everything accumulated so far has to be corrected for the new maximum. So per tile the kernel updates all three running quantities with the same factor exp(m_old - m_new):

m_new = max(m_old, rowmax(S))
l_new = l_old * exp(m_old - m_new) + rowsum(exp(S - m_new))
O_new = O_old * exp(m_old - m_new) + exp(S - m_new) @ V

The running normaliser l (the sum of exponentials) is rescaled by exactly the same factor as the output accumulator O, every tile. The identity that makes it exact is simple: exp(S - m_old) * exp(m_old - m_new) = exp(S - m_new). Apply it every tile, and after the last tile divide the accumulator by the normaliser, O <- O / l, to get the final result. Without that final division the accumulator is unnormalised. Once l is rescaled each tile and O is divided by l at the end, the output is identical to the full-matrix softmax, without the matrix ever existing.

Why the output must live in registers

Here is the part that dictates everything else. That rescale, O *= exp(m_old - m_new), is per row, and it happens on every tile. To apply it, the kernel has to know precisely which register holds which output row.

The convenient high-level matrix-multiply API (WMMA) hides that. It hands you the accumulator as an opaque fragment and keeps the register-to-row mapping to itself. So there is no way to reach in and rescale row 7 of the output, you would have to spill the whole accumulator to shared memory, rescale it there, and read it back, destroying the entire point.

Raw mma.sync.m16n8k16 exposes the registers directly. You know exactly which one is which row, so the rescale is a handful of register multiplies. That is why the defining property, the output accumulator never touches shared memory across the whole loop, requires an API that exposes the mma register-to-thread layout (raw PTX, or CuTe/CUTLASS), not the opaque WMMA fragment API. WMMA is ruled out not for stylistic reasons but for arithmetic ones.

The budget that follows

Because the output costs no shared memory, it lives in registers, the shared-memory budget can go entirely to the data that has to be staged. It is not free, though: it costs registers instead, and that register pressure is exactly what caps occupancy at the ~17% ceiling reported below.

One attention block, 97 KB of a 99 KB budget
Q tile32 KBK/V ring32 KBscores (S/P)32 KBfree2 KBO accumulator: 0 KB - it lives in registers
It fits, barely, which is why a deeper pipeline (which would need 102 KB) is impossible. The output accumulator costs nothing because it never enters shared memory; it stays in registers the whole time. That one fact is what forces raw mma.sync instead of the friendlier WMMA.

It fits in 97 KB of a 99 KB cap (estimates), which is exactly why you cannot pipeline deeper: a three-stage version would need about 102 KB and there is nowhere to put it.

The kernel claws back a little more with a no-round-trip trick, the column-to-lane mapping of the score output happens to match what the value multiply needs, so the scores are repacked straight from registers instead of bouncing through shared memory. Accumulating the value product in f16 halves the register footprint; doing the same for the score product was tried and reverted, the unpacking cost more than it saved.

The ceiling

After all of that, the kernel is latency-bound at about 17% occupancy, stalled on the async-copy wait.

You cannot pipeline deeper (no shared memory) and you cannot hide the matrix multiply itself, because consumer Blackwell has no asynchronous matrix engine, no tcgen05, no tensor memory. The instruction blocks the warp that issued it, full stop. Trying to raise occupancy by packing two blocks per SM measured 24% slower. The remaining gap to a data centre FlashAttention-4 kernel is partly silicon: the FlashAttention-4 warp-specialised TMEM design cannot be ported to sm_120 (no TMEM/tcgen05 on the chip). Other software improvements are not foreclosed.