Part IV · the second capstone & the craft

Flash Attention 4 — the algorithm and the two-MMA problem

Lesson 10 reached cuBLAS parity on GEMM: one MMA, repeated, fed by a perfectly choreographed supply chain. Attention reuses all of that machinery — TMA loads, tcgen05 MMAs, the TMEM accumulator, warp-specialized mbarriers — and yet it is fundamentally harder. Not because the hardware is different. Because attention is not one MMA repeated; it is two MMAs with real work wedged between them. This lesson derives the algorithm and shows exactly where that wedge breaks the clean GEMM pipeline. Lesson 12 is the warp choreography that tames it.

The memory-bound trap, restated from the roofline

Start with the obvious algorithm and watch it fall off the roofline from lesson 01. Attention over a sequence of length N with head dimension d wants the full score matrix:

S = Q·Kᵀ  ∈  ℝ^N×N,  then  P = softmax(S / √d),  then  O = P·V

The trap is the size of S. At a sequence length of 4096, the score matrix is 4096 × 4096 fp32 values ≈ 64 MB. That is orders of magnitude larger than the SM's shared memory, and far larger than the single 128×512 TMEM region you just learned to allocate in lesson 05. There is simply nowhere on-chip to put it.

So the naive kernel does the only thing it can: it writes S out to HBM, reads it back to compute the row-wise softmax, writes P back to HBM, then reads it again for the second matmul. The arithmetic is cheap; the traffic is enormous. This is the memory-bound trap — an N×N intermediate that the tensor core can produce far faster than HBM can absorb. The whole point of FlashAttention is to never materialize S.

Online softmax — three running states instead of one big matrix

The escape is to stream K and V in blocks and never hold more than one block's worth of scores at a time. The difficulty is softmax: a normal softmax needs the max and the sum over the entire row before it can normalize anything, and we are refusing to see the whole row at once. Online softmax is the trick that lets a row's normalization be built up incrementally.

For one query row we keep exactly three running states:

• row_max — the maximum raw Q·Kᵀ score seen so far (for numerical stability of the exponential).
• row_sum — the running softmax denominator (the sum of exponentials so far).
• O — the running output accumulator, a 1×d vector. This lives in TMEM, exactly like the GEMM accumulator.

For each K/V block we run the update below. It is the entire algorithm; read it slowly, because every later complication is a consequence of these six lines:

// pseudocode — one K/V block, for one query row.   maps to: Score MMA, CUDA-core softmax, Value MMA
S        = Q_block · K_blockᵀ                 // [score MMA]  raw scores for this block
m_new    = max(row_max, rowmax(S))            // extend the running max
scale    = exp( (row_max - m_new) / √d )      // correction factor for everything accumulated so far
P        = exp( (S - m_new) / √d )            // this block's UN-normalized probabilities
row_sum  = row_sum · scale + rowsum(P)        // rescale old denominator, add this block's mass
O        = O · scale + P · V_block            // [value MMA] rescale old output, add this block's contribution
row_max  = m_new                              // commit the new max

The key insight: one scale fixes everything at once

Here is why this works, and it is the single most important idea in the algorithm. The exponentials are computed relative to the running max. When a new block contains a larger score, row_max increases — and the instant it does, every quantity we accumulated under the old max is suddenly in the wrong scale. The old row_sum was a sum of exp(score − old_max); the old O was a weighted sum using those same too-large exponentials.

A single correction factor repairs both:

scale = exp( (old_max − new_max) / √d )   (always ≤ 1, since new_max ≥ old_max)

Multiplying both the denominator row_sum and the output accumulator O by this one scale pulls everything accumulated so far down into the new common scale, so the earlier blocks and the block we are about to add now share the same reference max. One scalar, two corrections — that is what makes "see the row one block at a time" arithmetically identical to "see the whole row at once."

Deferred normalization

Notice what P = exp((S − m_new)/√d) is not: it is not the final, normalized attention weight. We never divide by row_sum inside the loop. The output accumulator O is a sum of un-normalized contributions, kept in a consistent scale by the rescaling. Normalization is deferred to after the last K/V block, where the kernel writes:

O_final = O / row_sum

Deferring the divide keeps the inner loop free of a per-block reciprocal on every element of O — one division per row at the very end instead of one per row per block. The inner loop stays lean; the normalization happens once.

The exp2 trick — fold the scale into the exponent

Every line above calls exp(x / √d), and exponentials are expensive scalar ops sitting right on the critical path. Two micro-optimizations collapse into one. First, the hardware has a fast base-2 exponential (the SASS MUFU.EX2 path), not a natural one. Second, the 1/√d scale is a constant. So fold both into a single precomputed constant and evaluate every exponential as a base-2 exp2 on the raw score:

scale_log2 = log2(e) / √d   ⟹   exp(x / √d) = exp2( x · scale_log2 )

Now there is no separate 1/√d multiply and no natural-exp conversion on the critical path — just x · scale_log2 feeding the hardware exp2. This is the standard FA2/FA3 softmax optimization. It looks like a triviality, but it matters precisely because softmax is on the critical path: shaving the exponential shaves the wedge between the two MMAs.

Why FA4 is harder than GEMM — the crux

We can now state the thesis precisely:

GEMM is one MMA repeated; FA4 is two MMAs with real work wedged between them: online softmax, causal masking, and the rescaling that keeps earlier and later blocks in a common scale.

Trace where each piece runs. The two matrix multiplies — the Score MMA (S = Q·Kᵀ) and the Value MMA (O += P·V) — run on the tensor core. The softmax between them (exponentials plus row-wise reductions for rowmax and rowsum) runs on the CUDA cores. And because the Value MMA cannot start until P exists, and P cannot exist until the softmax finishes, the softmax sits directly on the critical path between the two tensor-core operations. In GEMM, nothing is ever on the critical path between two MMAs except the next MMA. That is the whole difference.

The consequence is the line worth memorizing:

So much of attention optimization is really softmax optimization — reformulating exp and overlapping softmax with the MMAs instead of stalling on it.

The four computational stages (vs GEMM's one)

Where the GEMM mainloop had a single repeated stage (MMA into TMEM), FA4's inner loop has four, and three of them move data across address spaces:

Stages ② and ③ are the wedge. Both add TMEM → register → TMEM round-trips that GEMM never had, and both happen between the two MMAs. The conclusion the course draws is exactly right:

FA4 is harder than GEMM not because it relies on different hardware, but because there are more tile values and more handoffs between them.

The tile-primitive graph

Draw the data flow as tiles and the hops between them. The GEMM version of this graph was a straight pipe: GMEM → SMEM → TMEM (accumulate) → GMEM. FA4 inserts two new rows — softmax and correction — and each one adds TMEM↔register traffic:

The two red rows are the entire cost of attention over GEMM. They are why the next lesson needs so many warp roles: somebody has to do the softmax and the correction while the tensor core is busy, or the pipeline stalls.

Interactive · online softmax, block by block

Watch the three running states evolve as K/V blocks stream in left-to-right for a single query row. Each block computes its own scores; whenever a block raises row_max, a rescale event fires — the old O and row_sum are multiplied by scale (highlighted) so earlier and later blocks share a common scale. Compare the streaming states (three tiny numbers, constant size) against the materialize-full-scores mode (a 64 MB bar that does not fit on-chip). At the end, the deferred normalize divides O by row_sum.

The contrast is the whole motivation in one picture: the streaming footprint is three scalars plus a 1×d accumulator — it never grows with N. The materialized score matrix grows as N² and is already ~64 MB at N = 4096, with nowhere on-chip to live. Online softmax is what makes attention fit.

TMEM layout — S, P, and O share one allocation via fp32/fp16 aliasing

The four stages produce three tile values that all want to live in TMEM: the scores S, the probabilities P, and the output O. TMEM is scarce (lesson 05: 128 Lane × ≤512 Col per CTA). FA4's answer is to make all three share one 128×512 TMEM allocation (N_COLS_TMEM = 512) by overlaying two views on the identical physical bytes:

• an fp32 view (call it tmem) — used for S (raw scores) and O (the fp32 accumulator), and
• an fp16 view (tmem_as_f16) over the same bytes — exposing twice as many indexable columns, used for P (which the Value MMA wants in fp16).

With MMA_N = 128 and a Q pipeline depth of 2, the layout per stage is:

• S_region = fp32 columns 0–127 (per Q stage).
• O_region = fp32 columns 128–255 (per Q stage). S + O together fill the 512-column allocation across the two pipeline stages.
• P_region = an fp16 alias over the same bytes as S_region. P reuses S's storage — but only after S has been consumed.

The constants worth carrying into lesson 12

The FA4 kernel is parameterized by a handful of tile and pipeline constants. They are worth stating once here, because the choreography in the next lesson is built around them: