Part IV · the second capstone and the craft

Flash Attention 4 — warp-role choreography and the edges

Lesson 11 gave the FA4 algorithm: two MMAs with softmax, masking, and rescaling wedged between, all sharing one TMEM allocation. That structure is what makes the kernel hard — it forces many roles and many handoffs. This lesson is the choreography. Who does what, which barriers connect them, and how the real-world edges — causal masking, grouped-query attention, scheduling — fold in.

The cast: 512 threads, split by function not by data

FA4 launches four warpgroups per CTA — 4 × 128 = 512 threads. The split that matters is not "this warpgroup owns this slice of the output." It is by function: a warpgroup (or even a single warp) is a job title. This is warp specialization from lesson 10, pushed to its limit: GEMM needed a producer, a consumer, and a writeback team; FA4 needs a memory mover, an MMA driver, two softmax teams, a correction team, and a store warp — because there is real CUDA-core work (the softmax) sitting on the critical path between the two MMAs, and it has to overlap with the math.

Why one warp for both MMAs? Because on Blackwell a tcgen05.mma is issued by a single elected thread anyway (lesson 03) — the warp is just the issuing context, not the compute width. Putting both MMAs on the same warp means the score→value ordering and the TMEM accumulator lifetime are decided in one place, by one program counter, with no cross-warpgroup negotiation. The cost of that simplicity is that everything else — softmax, correction, loads, stores — has to be choreographed around that single MMA stream with barriers.

The wiring: about a dozen named barriers

Lesson 06 taught the mbarrier: an SMEM object with an arrival count and a phase bit, where completion flips the phase and a waiter blocks until the phase differs from what it last saw. GEMM used four such barriers (two producer–consumer rings). FA4 needs roughly twelve, because there are more handoffs: load→score, score→softmax, softmax→value, value→correction, correction→store, plus a brand-new channel that GEMM never had.

The three flavors matter because they complete on different signals. A TMABar flips when the announced byte count lands (TMA's expect_tx); a TCGen05Bar flips when the tensor core says the MMA is done (tcgen05.commit); a plain MBarrier flips on a counted arrival. Mixing them up — waiting on a thread-arrival barrier for a byte transfer that hasn't landed — is one of the classic ways this kernel silently corrupts.

The value MMA waits on three gates at once

The single hardest synchronization point: before WG3·w0 can issue O += P·V, three independent conditions must all hold. It cannot proceed on any two of them.

That third gate is the subtle one and it is why correction and the value MMA are coupled: the value MMA accumulates into O, but the correction warp also writes O (the rescale). If the MMA accumulated before the rescale, it would add P·V onto an O that is still in the old, wrong scale. The barrier serializes them: rescale first (or release if no rescale was needed), then accumulate.

The interleave: why softmax hides behind the tensor core

Here is the move that makes the whole thing fast. The MMA driver does not run all the score MMAs and then all the value MMAs. Once both Q pipeline stages have primed, it interleaves them, ping-ponging between the two Q stages and between the two MMA kinds:

score  Q0 · K[n-1]
score  Q1 · K[n-1]
value  P0 · V[n-1]
score  Q0 · K[n-2]
value  P1 · V[n-1]
score  Q1 · K[n-2]
value  P0 · V[n-2]
…

Read the dependency: the moment S for Q0 is produced by a score MMA, WG0 can start its softmax on it — and while WG0 chews through exp2 and the row reductions, the MMA driver is already running the score MMA for Q1 and the value MMA for the previous block. Softmax on one Q stage overlaps the MMA work on the other. The inter-phase latency that lesson 11 warned about — the tensor core stalling while CUDA cores do softmax — is hidden because the tensor core always has the other stage's MMA to chew on. Pipeline depths: Q = 2, K/V = 3, TMEM = 2 — two Q stages is exactly what makes the ping-pong possible; one stage would serialize softmax behind MMA again.

Widget · the FA4 pipeline timeline

Five role lanes, a few inner-loop iterations. Watch the MMA-driver lane (lane 2) interleave score and value MMAs while the two softmax lanes (3 and 4) light up on opposite Q stages — softmax on one stage running underneath the MMA on the other. The correction lane (5) fires only when a block's max moved enough to need a rescale. Toggle causal masking to see diagonal-straddling blocks marked (the score MMA still runs; softmax masks the invalid columns to −inf).

The correction: the heart of why FA4 is hard

Online softmax means the running max can change when a new K/V block arrives. When it does, every value already accumulated into O was scaled by the old max and is now in the wrong scale. WG2's job is to fix that, in place, in TMEM. The fix is a full round trip — there is no "scale TMEM" instruction — so the accumulator goes TMEM → registers → TMEM in 16-wide chunks (RESCALE_TILE = 16):

// WG2, per 16-column chunk of the O tile (pseudocode → tcgen05.ld / .st)
o_row  = tcgen05.ld(O_region, chunk)        // TMEM → registers
acc_scale = exp2((m_old - m_new) * scale_log2)
o_row *= acc_scale                           // rescale in registers
tcgen05.st(O_region, chunk, o_row)           // registers → TMEM
tcgen05.wait.st()                            // make the store visible

That round trip over the whole O tile is expensive, so it is guarded twice:

• Did anyone need it? WG2 runs the rescale only if the softmax warpgroup signals — via any_sync across the warp — that at least one row saw its max move. If no row's max changed this block, the whole rescale is skipped.
• Did it move enough to matter? Even when some row moved, WG2 skips the rescale when the log-scaled delta stays above −rescale_threshold (= 8.0). If the max barely budged, acc_scale = exp2(small negative) ≈ 1, and multiplying O by ≈1 is pure waste — so the costly TMEM round trip is elided.

The ordering constraint ties back to the three-gate rule above: the value MMA must not accumulate into O until WG2 has rescaled or explicitly released it. Correction and accumulation share the O region, and the barrier (p_o_rescale) is what keeps one from clobbering the other.

Edge 1 · causal masking — three regions, not one

Causal attention forbids a query from attending to keys in its future. Naively you would compute the full score matrix and then mask the upper triangle — but that wastes the entire above-diagonal half. FA4 uses the same three-region split as FA2/FA3, deciding per (Q block, K/V block) which case applies:

Two things follow. First, masking the diagonal block before exp2 (not after) is what keeps the running softmax correct — a future column must never inflate the max or the sum. Second, the diagonal block is the source of work imbalance: an early Q block touches almost no K/V blocks (it can only look a little way back), while a late Q block touches all of them. That imbalance is exactly what motivates a smarter scheduler.

Edge 2 · GQA — free in the compute core

Grouped-query attention shrinks the KV cache by letting several query heads share one K/V head. With num_qo_heads = 32 and num_kv_heads = 8, the ratio is

GQA_RATIO = num_qo_heads / num_kv_heads = 32 / 8 = 4

so four query heads ride on each K/V head. FA4 packs those GQA_RATIO query heads into one Q tile. With BLK_M = 128, each packed head gets SEQ_Q_PER_TILE = BLK_M / GQA_RATIO = 128 / 4 = 32 query positions, and a row index is reinterpreted as

seq_pos = row // GQA_RATIO
q_head  = row %  GQA_RATIO

through a 3D Q-load view. Here is the elegant punchline:

Edge 3 · tile scheduling — which CTA owns which task

FA4 is a persistent kernel (lesson 09): a fixed pool of CTAs stays resident and pulls work. The unit of work is an attention task — a tuple (batch, kv_head, m_block). The scheduler's only job is to decide which task each CTA takes next; it never touches how the tile is computed.

The LPT scheduler is the direct payoff of the diagonal-block imbalance from Edge 1: because late Q blocks are demonstrably heavier, scheduling them first balances the finish times instead of leaving one CTA grinding through a long tail while the others sit idle. The governing rule, the same one from the GEMM ladder:

Does it actually compute attention? The correctness test

The reference is PyTorch's F.scaled_dot_product_attention with GQA enabled. The test config:

• Shapes: batch B = 1, sequence S = 1024, query heads Hq = 32, KV heads Hkv = 8 (so GQA_RATIO = 4), head dim D = 128.
• Dtype: fp16 in, with fp32 accumulation in TMEM.
• Tolerance: rtol = atol = 1e-2 — deliberately loose.

The tolerance is loose on purpose: it has to absorb fp16 rounding, the exp2 reformulation (lesson 11 folded 1/√d and log₂(e) into scale_log2 so the hardware exp2 replaces a software exp), the online-softmax reordering (a streaming sum lands in a different order than a batched one), and the final cast back to fp16. A tight tolerance would flag these as bugs when they are exactly the approximations the kernel is built on.

The whole point, in one table: what FA4 adds over GEMM

Lessons 08–10 built GEMM; lessons 11–12 built FA4. Every difference between them traces to a single structural change — a second MMA with softmax inserted between the phases. The hardware is the same Blackwell tensor core, TMA, TMEM, and mbarriers; FA4 just has more tiles and more handoffs.