Learning FlashAttention the Hard Way — Part 2

By Dmitry Trifonov•July 2, 2026

TutorialsGPUCUDACompilersAI

This is Part 2/2. Part 1 built the algebra: attention is an associative reduction over a twisted monoid, the state (m, d, o) merging by rescale-and-add. Read it first, or none of the rescales below will mean anything. This part takes that carrier down to the metal — and organizes the descent not as the FlashAttention version history, but as a stack of generic compiler moves, each with a measured price tag.

You survived the tutorial boss. Part 1 was the bonfire-lit cathedral: clean proofs, commutative diagrams, a monoid that associates because a bijection says so. The math promised that FlashAttention is not a one-off GPU hack but one instance of a generic, parallelizable reduction. A promise is not a kernel. Now the fog gate closes behind you, and the real level begins — the one with the registers, the shared-memory banks, and the mma.sync instruction that does not care how elegant your algebra was.

The thesis of Part 2 is the other half of Part 1's claim: associativity is the only property the schedule needs. Every scheduling decision below — tile the KV axis, keep the state in registers, defer the divide, fuse the matmuls into tensor-core atoms, double-buffer the loads — is licensed by the monoid, not by attention. To make that concrete we will not hand-wave the codegen, and this time we will not hand-wave the clock either. We run an actual compiler (deplodock, the one from the from-scratch compiler series), read the CUDA it emits from the carrier, and benchmark every rung on an RTX 5090 against the strongest baseline there is: PyTorch's bundled FlashAttention-2. The kernel you watch being derived ends the post at 206µs. So does FlashAttention-2.

Two shapes carry the whole post. A listing shape small enough that emitted kernels fit on a screen — q, k, v each [1, 4, 128, 64] fp16 — and a perf shape big enough that the clock means something: [1, 8, 4096, 64] fp16, non-causal, ~34 GFLOP per call. Every latency below is the perf shape, deplodock run --bench (100 iterations, CUDA-graph-captured), RTX 5090, and every configuration is a pinned, reproducible knob string. To follow along:

git clone https://github.com/cloudrift-ai/deplodock.git
cd deplodock && make setup && source venv/bin/activate
SDPA='q=torch.randn(1,8,4096,64,dtype=torch.float16); k=torch.randn(1,8,4096,64,dtype=torch.float16); v=torch.randn(1,8,4096,64,dtype=torch.float16); F.scaled_dot_product_attention(q,k,v)'

The Scoreboard

Since this rewrite has the numbers, let's put them where a reader can keep a thumb on them. Each row is one section of this post; each is a pinned config you can re-run; the right column is measured, not projected.

121× end to end, and a tie. Just as instructive is the graveyard: seven optimizations were built, measured, and refused — a warp-scoped barrier, the fast-exp path (twice), deeper rings, warp specialization, wave-quantization reshaping, a zero-barrier warp-private ring, and a register ping-pong. Each gets its paragraph, because each null carries a mechanism, and two of them are exactly why FlashAttention-3 and 4 cannot land on this chip.

The Schedule the Algebra Refused to Pick

Recall the carrier from Part 1. One query row q streams over the key/value rows; the state is a triple

and two states merge by keeping the larger max and rescaling the smaller side toward it:

We proved this associates. Associativity means any bracketing is legal: one thread folding the whole row left-to-right, a warp folding strided slices, a grid of blocks each folding a KV tile and a second kernel folding the partials. The algebra licenses all of them and prefers none. It is a statement about correctness, silent on cost.

Cost is what the GPU adds. The reduction axis here is the key/value sequence, and the thing we never want to pay for is the [S_q, S_k] score matrix S = QKᵀ. At sequence length 4096 that is a 16M-entry matrix per head, and writing it to HBM and reading it back is the entire cost of attention. So the one scheduling decision that matters — the decision FlashAttention is named after — is: tile the reduction axis, carry the state in registers across tiles, and never materialize S. Associativity is what makes that tiling legal. The memory hierarchy is what makes it necessary. Part 1 gave us the first; the rest of this post is about the second.

Anatomy of a Schedule: Two Dials

Before walking the stack, it is worth being precise about what a "scheduling move" even is, because the moves below are not unrelated tricks — they are settings of exactly two dials, and the algebra of Part 1 tells you which dial each one turns.

Dial 1: where the brackets go. Part 1's shortcut said every value the reduction ever combines is the state of some chunk of KV rows, and merging chunk states is gluing chunks. A schedule, for a fixed carrier, is nothing more than a choice of brackets over the KV sequence plus an assignment of each bracket depth to a piece of hardware. Factor the KV length as $N = G \cdot T \cdot S$ and the fold becomes a tower:

The innermost bracket is the KV tile Move 2 hands to the tensor cores; the middle level is Move 1's streaming loop, the running $(m, d, \mathbf{o})$ folding tile after tile in registers; the outermost is split-KV, partial states written out and folded by a tiny combine kernel. $G = 1$ is the FlashAttention-1/2 forward pass; $G > 1$ is Flash-Decoding. When $N$ doesn't divide the tower, pad the last tile with the identity $(-\infty, 0, \mathbf{0})$ — Part 1's masking story, load-bearing at last. The free axes — query rows, heads, the value dimension — sit outside the fold; they are map, not reduce, so partitioning them needs no algebra at all. And staging or pipelining doesn't relocate a single bracket — those moves schedule the data movement for a tree already chosen.

Dial 2: which algebra to bracket. The same attention function factors into algebraic operations in more than one way, and the brackets live inside a factorization. The materialized form — matmul, softmax, matmul, with the [S_q, S_k] score tensor written between them — is one factorization. Flash is a different factorization of the same function, and every step from one to the other is a trick Part 1 already taught: tuple the max and sum folds into one fold over (m, d) (online softmax), twist the coordinates by $e^{-m}$ (transport of structure), absorb the downstream P·V fold into the carrier as the $\mathbf{o}$ channel (the expectation-semiring move). No bracketing of the materialized form ever deletes the $O(N^2)$ score traffic; no bracketing of the flash form ever frees the carrier's registers. The dials don't mix.

One update from the earlier draft of this post: the dial-2 decision is no longer behind a feature flag. The recognizer fuses by default — greedy compilation ships the flash form — and the materialized escape is a structural placement pin, PLACE@fold=cut. Which means the "watch it do the wrong thing" demo below now requires asking for the wrong thing. Progress has its costs.

The Move Stack

Compiler Part 2 lowered a matmul by walking a fixed stack of small, generic rewrite rules: register-tile the outputs, chunk the K loop, stage the inputs, double-buffer, cp.async, TMA, software-pipeline. This post makes a claim that sounds bolder than it is — FlashAttention is the same stack, walked on a monoid carrier — and then, unlike the earlier draft, proves it on a clock. Every move below is offered by the compiler as a fork: a set of schedule rows the autotuner ranks, each spelled by the same knob codecs a matmul uses (TILE@<axis> for the mma geometry, STAGE@<axis> for the operand pipeline, REDUCE@<axis> for the fold partition). The pinned winner, for the record — every listing below reproduces with these two environment variables:

export DEPLODOCK_TILE='a:mma_m16n8k16_f16/w4x1/f2x8/k4'   # 4 warps × 2 query tiles × 64-key blocks
export DEPLODOCK_STAGE='d2/tma/ring'                       # 2-slot TMA prefetch ring

Move 1 — Stream the Reduce

The IO-aware tiling FlashAttention is named after. Start by watching the compiler do the wrong thing — which, since fusion is now the default, takes a pin:

DEPLODOCK_KNOBS="PLACE@fold=cut" deplodock compile -c "$SDPA" --ir loop

Two kernels. The first computes every score and writes the full [1, 8, 4096, 4096] tensor to HBM; the second reads it straight back for the softmax and the P@V. That trailing store is the villain of the entire field. Everything below is the campaign to delete it.

Drop the pin, and the greedy default emits the fused form — the recognizer builds the structural tree (a TWISTED streaming reduction whose per-step partial holds both contractions) and the loop IR reads like Part 1 wearing a hard hat:

=== 0: scaled_dot_product_attention ===
    Reduction[kv] twisted
        for kv in 0..4096                       # the streaming reduce axis
            for dd in 0..64                     # score: q · kⱼ
                k_e = load k[b0, b1, kv, dd]
                q_e = load q[b0, b1, m, dd]
                sacc <- add(sacc, multiply(k_e, q_e))
            s = multiply(sacc, scale)
            m_i__t0 = maximum(m_i, s)            # m  = max(m_i, sⱼ)
            m_i__t2 = exp(subtract(m_i, m_i__t0))  # α = e^{m_i − m}  — the twist
            m_i__t5 = exp(subtract(s, m_i__t0))    # p = e^{sⱼ − m}
            l_i <- add(multiply(l_i, m_i__t2), m_i__t5)         # d = d·α + p
            O_i <- add(multiply(O_i, m_i__t2), O_i__pv)         # o = o·α + p·vⱼ
            m_i <- maximum(m_i, s)
    O_i__proj = divide(O_i, l_i)                 # the deferred normalize, last moment

Init(m_i, maximum) / Init(l_i, add) / Init(O_i, add) is the carrier with its identity (−∞, 0, 0). The score s is consumed immediately and never stored; there is no [4096, 4096] tensor anywhere. This is not "inspired by" the algebra of Part 1; it is the algebra of Part 1, compiled. And because the schedule rode on the algebra and not on a shape, causal masks, GQA's head // group indexing, and additive biases all fold into the same kernel with no new code path — the coverage suite (tests/compiler/e2e/test_attention_coverage.py, 66 cases) sweeps all of them, static and symbolic-seq_len alike, against torch references.

Correct, fused, and slow: 24943µs on the perf shape, every FMA on the CUDA cores while the tensor cores idle. The rest of the post wakes them up.

Move 2 — Atomize the Contractions (mma.sync)

The scalar kernel leaves ~95% of the FLOPs on the table: QKᵀ and P@V are matmuls, and matmuls belong on the tensor cores. deplodock already owns the move — atomize, the same rule that fuses a Linear's multiply-accumulate cell into an mma.sync.m16n8k16 — and the whole trick of "tensor-core attention" is letting it fire on the two matmul cells buried inside the stream. Structurally, flash is a two-Contraction tree: the score Σ Q·K at the head of the streaming step and the Σ P·V fold as the expect channel's lift. Give both nodes an mma TilePlan and the tree realizes at fragment residence: scores live in mma C-fragments, the softmax runs between the fragments, and the running merge is regenerated from the carrier's channel spec — rowmax as a register reduction plus a __shfl_xor_sync butterfly, the twist α as two expfs, the rescale in place on the output fragments.

The operands want fp16 (or bf16) with fp32 accumulation — and the accumulation costs nothing to ask for, because mma.sync.m16n8k16.f32.f16.f16.f32 accumulates into fp32 C-fragments natively and the carrier (m, l, O) is fp32 by construction. Part 1's numerical story survives the precision drop exactly as promised; the numerics section has the measured table.

deplodock compile -c "$SDPA" --target sm_120 --ir cuda | grep -E 'mma|ldmatrix' | head -4

dpl_mma_load_a_gmem(_sacc_a, &q[...], 64);                        // Q fragment (A)
dpl_mma_load_b_gmem_trans(_sacc_b0, &k[...], 64);                 // K fragment (B, native col-major)
dpl_mma_m16n8k16_f16(sacc_f0, _sacc_a, _sacc_b0, sacc_f0);        // S += Q·Kᵀ, fp32 accumulate
...

612µs — a 41× cliff from one move, and the biggest single number in the post. It is also, note well, still 3× away from FlashAttention-2's 206. The remaining gap is not FLOPs (the mma count is now identical); it is everything around the mmas, and closing it is where the article stops being a derivation and becomes an investigation. From here on, every step follows the same loop: profile, attribute, fix, re-measure — Nsight Compute counters first, hypotheses second.

A word on Move 3, since it fired silently just now: partitioning the free axes. The query rows went onto the grid in blocks, the value dimension folded into the output fragments, and the streaming KV axis stayed serial — FA-2's "work partitioning," which in deplodock is not a decision anyone hand-makes but a geometry grid the fork enumerates: warps-per-CTA × key-atoms-per-block × query-tiles-per-warp, every divisibility-legal point a schedule row the tuner can rank. The winner's geometry is where the sweep put it. Hold that thought for the reg_m section, where the third dimension of that grid becomes the closing move of the post.

Move 4 — Stage the Reused Reads, and the First NCU Lesson

Once KV is tiled, each tile's K and V rows are read by every query row in the block — the textbook condition for staging into shared memory, exactly what the stage move does for a matmul's operands. The STAGE@<kv> codec puts a resolved pipeline on the streaming axis: d1/cp is a single-buffered cooperative cp.async fill, and the streaming step becomes the drain of the same staged_kloop skeleton the matmul tier runs.

for (int _fk = threadIdx.x; _fk < 128; _fk += 32) {                       // cooperative K fill
    cp.async.cg.shared.global [_k_smem + ...], [k + ...], 16;
}
cp.async.commit_group;  cp.async.wait_group 0;  __syncthreads();

511µs — staging pays 1.20× immediately. And then the profiler earns its seat at the table. The first NCU pass on the staged kernel showed a number that dwarfed everything else on the card:

One hundred thirty-two million bank-conflict replays. The K/V slabs have power-of-two rows — 128 bytes at head-dim 64 — so every ldmatrix row read in the drain lands on the same bank group and serializes eight ways. This is the classic failure the padding trick exists for, and the fix is comically small: pad each slab row by 16 bytes. One ldmatrix-row of padding shifts consecutive rows across bank groups while keeping every 16-byte alignment intact.

A 1.62× jump from sixteen bytes per row — the single largest post-mma win in the campaign, and the conflicts now sit below the hand-written reference's. (When the transport upgrades to TMA later, the padding retires in favor of the hardware swizzle — the same fix, done by the copy engine.)

Below the Dials: Two Codegen Fixes

The next two rungs turn no dial at all — they change neither the brackets nor the algebra, only how well the chosen schedule is realized. Both were found by reading counters, and both preserve a property worth naming: from the first staged kernel to the end of this post, every optimization is bit-identical — same values, same operation order — and the test suite enforces it (test_staged_warp_flash_bit_identical_to_gmem_direct). One ladder of kernels, one set of output bits; the numerics section shows it in the error table's identical rows.

The C→A register repack (271µs). P is the first mma's output (a C-fragment) and the second mma's A operand — in what look like incompatible layouts. The first version of this kernel bounced P through a shared-memory slab and ldmatrix-ed it back: ~36 shared-memory operations plus a sync, every streaming step. But write out the two PTX-fixed lane maps and a gift appears: the m16n8 C-fragment is element-for-element lane-aligned with the m16k16 A-fragment's k-halves. Two adjacent score fragments convert into one P@V operand fragment with four cvt.rn.f16x2.f32 instructions — no shuffle, no shared memory, no sync. The smem handoff, its slab, and its barrier all deleted. (Foreshadowing: when we later read the FlashAttention-2 source, this exact trick is there under the name convert_layout_acc_Aregs. Both kernels need it; neither could afford the round-trip.)

The Q hoist (233µs). Q's operand index carries the query and head-dim axes and never the stream axis — a structural fact, not an observation about attention — so its fragments are loop-invariant, and the emitter was reloading them from global memory every KV block anyway: sixteen scattered loads per warp per step, 64 steps, for values that never change. Hoisted to resident registers, loaded once. (Second foreshadowing: FlashAttention-2 has this as a compile-time option, Is_Q_in_regs — and ships it disabled, re-reading Q from shared memory every block. On this geometry, our measurement says registers win.)

Move 5 — Pipeline the Loads, or: the Ring That Refused to Win

Here is the best sequence in the campaign, because it is a story about when an optimization is real. The d2/cp/ring knob — double-buffer the K/V slabs, prefetch tile i+1's copies under tile i's math, FlashAttention-3's asynchrony in miniature — was available from the moment staging landed. And it lost, every single time it was measured:

Three consecutive regressions, then a win — with the same knob. The mechanism: software pipelining hides copy latency behind math, but for three rungs the streaming step was so bloated (bank conflicts, the smem handoff, the Q reloads) that the loop body was its own latency blanket — the copies were already hidden behind waste, and the ring's doubled slab footprint just cost occupancy. Only after the repack and the hoist thinned the step enough to expose the copy latency did the prefetch have anything to hide. Software pipelining pays exactly when the loop stops insulating it. If you tune knobs by folklore instead of re-measuring after every structural change, you will write this move off forever on day one.

One more instruction-level cleanup rode in here: a peephole pass that fuses two slab-adjacent ldmatrix.x2 drains into one ldmatrix.x4 (the same fusion the matmul tier's staged drains want, which is why it's a pass and not an emitter tweak). LSU instructions dropped 30%, latency 222 → 219µs — and the muted gain was itself diagnostic: at that point the kernel had stopped being memory-instruction-bound. The counters said so; the next move had to come from somewhere else.

Move 0 Returns — Register-Tile the Queries

Where else: the dependency chain. NCU at 219µs showed the SM as busy as FA-2's but the fma-pipe slightly behind, and the per-step structure explains it — QK mma → rowmax shuffle → exp → rescale → PV mma is a serial chain, hidden only by however many independent chains are in flight. FlashAttention-2's kernel famously runs at 255 registers per thread and ~12% occupancy: few warps, each carrying lots of independent work. That design has a name in the moveset, and it is the very first move of the matmul article: register tiling, f<FM>x<FN> — the same codec field, applied to the flash geometry grid's third dimension. Each warp takes FM=2 query tiles: two independent (m, l, O) chains against shared K/V fragments (every B fragment loaded once, fanned across both tiles' mma chains — the m.reg operand sharing, verbatim from the matmul tier).

212µs. And the NCU profile converged to the reference kernel's own shape:

The register-tiling move that opens the GPU-matmul story (5.2× on a naive matmul) is the move that closes the flash gap. The moveset ends where it began.

New Hardware: the TMA Ring

The last landed rung uses silicon FlashAttention-2 predates. cp.async.bulk.tensor — TMA — copies whole tiles with one instruction from one elected thread, against a descriptor encoding the tensor's shape. The folklore says TMA can't feed attention because its descriptor is 2-D and K/V are batched 4-D tensors; the folklore is wrong — descriptors go to rank 5, so the batched operands encode as rank-4 boxes (1, 1, block, head_dim) with the kernel's own batch/head index expressions riding as box-origin coordinates (GQA's h // group included, zero special cases). The slabs stay dense and take the hardware swizzle instead of padding — the copy engine permutes 16-byte chunks on the way in, and each ldmatrix XORs its address on the way out. Same bank-conflict fix, done by hardware.

cp_async_bulk_tensor_4d(&_k_smem[...], _desc_k, 0, kv0, b1, b0, &_mbar[slot]);   // one thread, whole tile

206µs — parity. And the memory counters now beat the hand-written kernel: 3.31M LSU instructions vs FA-2's 4.36M (the per-thread cp.async fill loops became one thread's box copies), 4K conflicts vs 22K. Both kernels sit at ~167 TFLOPS — about 80% of the RTX 5090's fp16/fp32-accumulate tensor roofline — separated by noise.

What Measurement Refused

Seven optimizations died in this campaign, and the article would be dishonest — and much less useful — without them. Each was built, run, and measured on the same harness as the wins.

1. Warp-scoping the handoff barrier (~0µs). The C→A smem handoff's __syncthreads looked like a classic convoy — four warps re-synchronized every streaming step. Narrowed to __syncwarp: dead flat. The staged loop's transport barriers re-converge the warps anyway. (The whole handoff died to the repack shortly after, but the null stands: the convoy was never the cost.)
2. Fast exp (~0µs, measured twice). libm expf is a ~15-instruction sequence; __expf is one FMUL plus MUFU.EX2 — and swapping ~70 exponentials per lane per step for the cheap form moved nothing, at two different kernel structures. The fma-pipe counter sat at 7% and meant it: the exponentials ride entirely in the mma shadow. This also settles the fancier version (fold scale·log₂e into the score multiply and run the loop in base 2 — which FA-2 does): algebraically free, measurably irrelevant here.
3. Deeper rings (d3, d4): occupancy loss beats extra overlap at every geometry.
4. Wave-quantization reshaping: the winning grid (256 CTAs) is fully co-resident in one wave on 170 SMs; the 512- and 128-CTA variants both lost.
5. Warp specialization (213 vs 206µs). The FA-3 signature move — a producer warp band driving the TMA fills while consumer warps compute — was wired end-to-end (the fills' elected thread rides the wrapped thread index, the stores guard to the compute band) and it lowers, validates, and loses: at a 128–256-thread CTA whose uniform TMA ring already keeps the tensor pipe at reference utilization, a dedicated producer band only spends occupancy.
6. The warp-private ring (291 vs 206µs, reverted). The boldest attempt: give each warp its own ring slots and mbarriers, elected per-warp TMA fills, __syncwarp-only ordering — a streaming loop with zero CTA barriers, something FA-2's cooperative fills structurally cannot have. It lowered correctly, passed every accuracy test, and lost by 40%, for a reason worth framing: flash's K/V are read by every warp in the CTA, so per-warp slabs forfeit exactly the sharing that justified staging (×4 smem replication clamped every ring to a single buffer). The move is only correct where warps stream disjoint data — a within-CTA split-KV, i.e. a decode-shape kernel. The code was deleted; the finding was kept.
7. The register ping-pong (p2, ~0µs). Double-buffering the drain's operand fragments — explicit software pipelining of the ldmatrixs. Flat: ptxas was already scheduling the unrolled straight-line drain as well as the explicit form. Kept only because it makes a documented knob apply uniformly.

Notice the shape of this list. Items 5 and 6 are the asynchrony toolkit — FlashAttention-3's whole thesis — correctly implemented and correctly rejected by this hardware at this scale. That is not a failure of the moves; it is a measurement of the chip, and it sets up the last two sections.

Move 6 — Split-K the Carrier (the Move This Shape Doesn't Want)

The cleanest callback to Part 1 in the whole post, kept short because the honest answer for this shape is "no." Split-KV — Flash-Decoding — cuts the KV axis across CTAs: each block streams its slice into a partial (m, d, o), and a tiny second kernel folds the partials with ⊕, legal in any order because Part 1 proved the carrier commutes. An additive reduce would be wrong here — the partials carry different maxes and must rescale before adding — so this is the twist, load-bearing for the most important inference workload there is. deplodock emits the cross-CTA split for an ordinary matmul today (DEPLODOCK_KNOBS="REDUCE@k=g4k" … — the same carrier-general combine machinery the flash version reuses).

But split-K redistributes work; it never reduces it. It exists to fix grid starvation — one query row against a mile of KV cache, nothing to parallelize. The perf shape is the opposite regime: 256 CTAs, all co-resident, tensor pipe at 80% of peak. Splitting here would add partial-state traffic and a combine launch against a full machine. The right time for this move is the decode chapter — where, pleasingly, the warp-private ring from the graveyard above becomes correct too, because decode's cooperating warps stream disjoint KV. The moveset keeps its rejects filed by precondition.

Face-Off: Reading the FlashAttention-2 Source

We tied a kernel we never looked at. Then we looked. The exact instantiation PyTorch ran (from NCU's demangled name): flash_fwd_kernel<Flash_fwd_kernel_traits<64, 128, 128, 4, Is_Q_in_regs=false, Share_Q_K_smem=false, half_t>> — and here is the generated winner against the Dao-AILab source it corresponds to:

Three things deserve to be said out loud.

The geometry converged. The autotuner's winning point — 128 query rows, 4 warps, 2 m-atoms per warp — is Dao's hand-picked constant set. Better: the FA-2 source comments document the author's manual sweep — "Using 8 warps is 18% slower for seqlen=2k … block size (256 x 64) is 85% slower, because of register spilling" — and they read line-for-line like our fork-row sweep table, register-spill cliff included (our 4-tile point spilled to 341µs the same way). A hand-tuning notebook and a measured schedule fork are the same document, produced by different processes.

The tricks are mutual. The C→A register repack we derived from the PTX lane maps is in their source as convert_layout_acc_Aregs. Their exp2f/log₂e fold is the trick our flat fast-exp measurements had already priced at zero on this chip — hygiene, not speed, for both kernels.

The differences are explainable, and two favor the generated kernel. Q residency: they ship the register variant disabled and re-read Q from smem per block; our measurement says registers win at our block width. Transport: single-thread TMA boxes versus per-thread cp.async loops is precisely the 1M-instruction LSU gap in their disfavor — TMA simply postdates their design. Their 128-key blocks suit cooperative fills; our 64-key blocks suit a TMA ring. And their peeled causal block-skip is the one behavior we haven't built (it matters only for causal shapes, and it restores parity rather than breaking it — FA-2 skips, so would we).

The compiler never knew it was attention. It measured its way to the same kernel the author reasoned his way to — and on the memory path, a little past it.

Why FA-3 and FA-4 Cannot Land on sm_120

The natural next question — we tied FA-2 using post-FA-2 hardware, so why not keep climbing the version numbers? Because the ladder's upper rungs are built from instructions this chip does not have, and — the part measurement adds — because the preconditions for those tricks are absent even where imitations are possible.

FlashAttention-3 is a Hopper kernel. Its three pillars are wgmma (warpgroup MMA — an asynchronous tensor-core instruction a warpgroup issues and walks away from), TMA, and a warp-specialized producer/consumer pipeline deep enough to exploit both. Consumer Blackwell (sm_120) has TMA — we used it — but no wgmma: its tensor core is the synchronous mma.sync, which blocks the issuing warp until the result lands. Warp specialization without async MMA is a body without a heartbeat, and we didn't have to take that on faith: we built the producer/consumer split on our TMA ring, and it measured 213 vs 206 — the machinery is only worth its occupancy cost when consumers can stack asynchronous MMA work behind the producer's copies, which mma.sync forbids. FA-3's FP8 path is the one pillar sm_120 could genuinely take up (the chip has FP8 tensor cores), but that is a precision change, not a schedule move — a different article.

FlashAttention-4 is a B200 kernel. It targets tcgen05 MMA — asynchronous and accumulating into TMEM, the 256 KB tensor memory that exists only on datacenter Blackwell (sm_100) — plus 2-CTA MMA across paired SMs. None of that silicon is present on sm_120. And FA-4's two portable-sounding ideas fail the measured-precondition test here too: the FMA-emulated exp answers a datacenter-Blackwell imbalance (tensor throughput doubled while SFU stood still), whereas our fma-pipe counter shows the exponentials already riding free in the mma shadow at consumer scale — we measured the fast-exp swap flat, twice; and the conditional rescale (skip the α multiply when the max barely moves) attacks a rescale cost that is a rounding error in our profile. Both kernels on this chip are parked at ~80% of the tensor roofline with the softmax hidden; the remaining 20% is intra-chain latency that only asynchronous MMA can overlap. That is FA-4's asymmetric hardware scaling thesis, confirmed from below: on silicon where the asymmetry hasn't opened, the tricks that answer it have nothing to answer.

So the honest frontier statement for consumer Blackwell: FA-2 parity is the fixed point, TMA is the one post-FA-2 gift the chip gives (and it shows up as our LSU lead), and the FA-3/4 rungs stay locked until the silicon under them ships in a consumer part.

Numerics at the Metal

The associativity of Part 1 was free; the numerical safety was not, and the metal is where the bill comes due. Three facts carry over and are now measured rather than asserted:

• fp32 statistics under low-precision matmuls. The QKᵀ and P@V mmas run in fp16, but the carrier (m, d, o) accumulates in fp32 — natively, in the mma C-fragments. Part 1's perturbation bound applies to the fp32 reduction regardless of the matmul's input precision.
• The rescale never amplifies. Every merge multiplies the running state by α = e^{Δm} ≤ 1; streaming over 64 tiles is as stable as one batched softmax.
• The measured table (listing shape, max abs error vs an fp64 torch reference):

The last two rows are identical to the digit — that is the bit-identity chain made visible. Ten optimizations, one set of output bits; the f16 error is input-rounding scale, exactly where Part 1's analysis puts it.

The Payoff: the Compiler Never Knew It Was Attention

Step back to the bonfire. Part 1's whole reason for the abstract-algebra detour was a claim about compilers: a system that can lower one monoid reduction can lower all of them, attention included, with no per-architecture special case. The receipts, in the pipeline's own vocabulary: flash is recognized as a TWISTED streaming reduction whose per-step partial happens to hold two Contraction nodes — the same structural kinds a softmax (one Reduction) and a matmul (one Contraction) lower through. Its schedule is spelled by the same axis-keyed knob codecs as a matmul's (TILE@dd, TILE@pj, STAGE@kv, REDUCE@kv); its fork rows are ranked by the same tuner; its staging rides the same staged_kloop skeleton; even the tensor-core realization is "the contraction nodes carry an mma TilePlan," not an attention emitter. The pass-authoring rule that enforces all this is written into the compiler's own docs: dispatch on the carrier algebra, never on a named shape. A twisted monoid is a monoid.

And the number that grades the claim: the kernel this generic machinery produces, tuned over its own fork, is latency-identical to FlashAttention-2 on an RTX 5090 — with the same geometry the FA-2 author chose by hand, the same register tricks, and a cleaner memory profile. 66 end-to-end tests pin the behavior across causal/GQA/masked/static/symbolic variants; every rung of the ladder above is one pinned knob string away from reproduction.

Wrap-Up

Part 1 ended with four claims. Part 2 is the receipt for all of them, with instrument readings attached:

1. The schedule needs only associativity. Every landed move — stream the KV axis, carry (m, d, o) in registers, atomize the matmuls, stage and pipeline the tiles, register-tile the queries — is licensed by the monoid, not by attention, and together they took one kernel from 24943µs to 206µs.
2. Efficient is a stack of generic moves, plus a profiler. The moves supplied the search space; NCU counters picked the path through it. The two biggest wins (16 bytes of padding; a lane-map coincidence) and all seven refusals came from attribution, not intuition — and the refusals were as load-bearing as the wins, twice telling us which decade of FlashAttention the silicon can actually express.
3. The monoid's split-K is Flash-Decoding — filed, with its preconditions, for the shape that wants it. Long-context decode is Part 1's associativity cashed at scale.
4. Numerical safety is paid separately, and the receipt is printed. fp32 stats, the never-amplifying rescale, one set of output bits across ten optimizations.

You came for a GPU kernel and you leave with a reduction that happens to run at the speed of the hand-written state of the art. Same as it ever was.

Don't you dare go hollow.

References

Series

1. Learning FlashAttention the Hard Way — Part 1. The algebra: the twisted monoid, transport of structure, the numerical analysis.
2. A Principled ML Compiler Stack in 5,000 Lines of Python — Part 2. The Tile IR and the generic move stack (matmul mma.sync, staging, TMA, pipelining) this post reuses.
3. A Principled ML Compiler Stack — Part 3. The autotuner that searches the move space.
4. GPU Matmul Optimization. Register tiling, padding, swizzle, split-K — the same moves, on the kernel that taught them.

FlashAttention

1. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
2. Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023. The kernel this post ties; source for the face-off section (v2.7.4, flash_fwd_kernel.h / kernel_traits.h / softmax.h).
3. Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. 2024. wgmma, TMA, warp specialization, FP8 — the Hopper rungs.
4. Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, Tri Dao. FlashAttention-4. 2026. tcgen05, tensor memory, the FMA-emulated exp, the conditional rescale — the B200 rungs.
5. Tri Dao, et al. Flash-Decoding for long-context inference. 2023. Split-KV — the monoid's split-K.

GPU programming primers

1. Simon Boehm. How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance. The why behind tiling, staging, and pipelining.
2. Vasily Volkov. Better Performance at Lower Occupancy. GTC 2010. The register-tiling argument — the one FA-2's 255-register design and our reg_m rung both descend from.
3. NVIDIA. PTX ISA — ldmatrix, mma.sync, cp.async.bulk.tensor. The lane maps behind the C→A repack and the paired-x4 drains.