# 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.

Source: https://rfriedmann.de/blog/online-softmax-and-flashattention/
Published: 2026-06-17 · Track: log · Level: Expert

The [kernels post](/blog/optimizing-kernels-consumer-blackwell/) 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.

<figure class="chart">
<div class="chart-title">Online softmax: one tile at a time, no giant matrix</div>
[diagram omitted — see the page for the chart]
<figcaption class="chart-caption">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.</figcaption>
</figure>

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.

<figure class="chart">
<div class="chart-title">One attention block, 97 KB of a 99 KB budget</div>
[diagram omitted — see the page for the chart]
<figcaption class="chart-caption">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.</figcaption>
</figure>

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](/blog/what-the-5090-lacks-vs-datacenter/)). Other software improvements are
not foreclosed.