Part III · GEMM from tiled to SOTA

GEMM III — warp specialization & clusters to SOTA

This is the climax of the GEMM ladder. We have a persistent kernel with a TMA-fed software pipeline — correct, prefetching, and still leaving the tensor core idle half the time. The reason is structural, not a tuning knob: one team of threads is being asked to load and compute and write, so it can only ever do one at a time. The forced move is to stop sharing the work and give each engine its own warps. Then we pool two SMs into a cluster, then stack a second math consumer on top — and the kernel lands at cuBLAS parity.

The whole ladder, in one table

Before the climax, here is the entire journey on one screen. Every row is a forced move that removed one more thing starving the tensor core. The target throughout: a B200 (148 SMs, 228 KB SMEM/SM), M = N = K = 4096, fp16 in, fp32 accumulate.

Read the shape of this table, not just the numbers. One step dominates everything: TMA async is ≈142× — the inflection where the kernel stopped paying the thread-copy tax (lesson 09). Everything in this lesson is the long tail that closes the last gap: the software-pipeline-plus-warp-specialization combination is worth roughly 2.2×, clusters another roughly 2.2×, and the final multi-consumer step about 10%. By Step 8 the kernel is compute-bound — the clusters doubled arithmetic intensity and removed the last big memory stalls, so there is simply very little memory latency left to hide. The closing 10% being small is not a disappointment; it is exactly what "we hit the compute ceiling" looks like.

Step 7 — warp specialization: give each engine its own warps

The single-warpgroup kernel from lesson 09 multiplexes three jobs onto one team of 128 threads. Warp specialization assigns those jobs to distinct warps by role, so the memory engine and the tensor-core engine run concurrently on independent warps. We use two warpgroups, WG_NUMBER = 2 (256 threads, 8 warps), and split the roles:

Now the producer warp can be filling stage k+1 of the SMEM ring while the consumer warp is multiplying stage k — the TMA engine and the tensor core are both busy in the same wall-clock window. That is the overlap the software pipeline alone could not deliver, because previously the same threads had to finish loading before they could start computing.

Four barriers = two producer–consumer rings

Concurrency means the warps no longer march together, so we cannot coordinate them with a single block-wide rendezvous. Each handoff between two roles needs its own mbarrier, and each handoff is bidirectional — the consumer must also tell the producer "I'm done with that buffer, you may refill it." So the kernel runs two producer–consumer rings, four named barriers total:

A small state object, PipelineState(PIPE_DEPTH, phase), carries the per-ring bookkeeping: a .stage index (which physical buffer in the ring) and a .phase bit (lesson 06's parity), advanced together by .advance(). This is the same stage-and-phase machinery from lesson 06, now instantiated twice — once per ring.

// PSEUDOCODE (TIRx-style) — the warp-specialized mainloop skeleton.
// Maps to CUTLASS warp-specialized GEMM: producer/consumer warps + PipelineTmaAsync.

WG_NUMBER = 2                       // two warpgroups = 256 threads = 8 warps

if   warpgroup == 1 and warp == 3:  // ── TMA PRODUCER ───────────────────────
    state = PipelineState(PIPE_DEPTH, phase=1)        // NOTE: starts phase 1
    for k in range(num_k_tiles):
        wait(mma2tma, state.phase)                    // buffer free to refill?
        tma_load(A[k], B[k]  ->  Asmem[state.stage], Bsmem[state.stage],
                 arrive=tma2mma[state.stage])         // expect_tx self-completes it
        state.advance()

elif warpgroup == 1 and warp == 0:  // ── MMA CONSUMER ──────────────────────
    state = PipelineState(PIPE_DEPTH, phase=0)        // NOTE: starts phase 0
    for k in range(num_k_tiles):
        wait(tma2mma, state.phase)                    // SMEM tile ready?
        tcgen05.mma(Asmem[state.stage], Bsmem[state.stage], -> TMEM, accum=(k>0),
                    arrive=mma2tma[state.stage])       // free the SMEM buffer
        state.advance()
    tcgen05.commit(arrive=mma2ld)                      // TMEM results ready

elif warpgroup == 0:                // ── WRITEBACK (all 4 warps) ────────────
    wait(mma2ld, phase=0)                              // TMEM ready?
    regs = tcgen05.ld(TMEM)                            // 4-warp readback, lesson 05
    tma_store(regs -> C_gmem)                          // commit_group / wait_group
    arrive(ld2mma)                                     // TMEM free for next tile

The tcgen05.mma here is the Blackwell tensor-core instruction from lesson 03, reading SMEM operands and accumulating into TMEM; tcgen05.ld is the 4-warp TMEM→register readback from lesson 05. The barrier waits lower to mbarrier.try_wait.parity and the loads to cp.async.bulk.tensor — exactly the PTX from lessons 04 and 06.

The subtlest bug: opposite initial phases

Look again at the two PipelineState initializations above. The producer starts phase=1; the consumer starts phase=0. This is not a typo, and getting it wrong is the most insidious failure in the whole kernel. It follows directly from lesson 06's rule that try_wait(bar, phase) passes when the barrier's phase differs from the argument:

• The barrier is freshly initialized at phase 0, and the SMEM buffers start empty.
• The producer's first wait(mma2tma, phase=1) sees the barrier at phase 0 → 0 ≠ 1 → it passes immediately. Correct: the buffers are empty, so the producer is free to start filling them right away.
• The consumer's first wait(tma2mma, phase=0) sees phase 0 → 0 = 0 → it blocks. Correct: no data has been produced yet, so the consumer must wait.

So the two ends are deliberately one phase apart: one must run, one must wait, and the opposite initial phases encode exactly that asymmetry on the same fresh barrier.

warpgroup_sync(10), not cta_sync()

Once the warpgroups run different code paths, the block-wide barrier becomes a trap. cta_sync() uses hardware barrier #0 and insists that every thread in the CTA arrive at it. But inside a warpgroup branch — if warpgroup == 1 and warp == 3 — only some of the CTA's threads are even executing that code. A cta_sync() placed there waits for threads that will never arrive, because they took a different branch entirely.

A real-toolchain technique the course kernel omits: setmaxnreg

Why PIPE_DEPTH=2

The pipeline depth is the number of SMEM buffers in Ring 1. PIPE_DEPTH=2 is the smallest depth that lets load and compute overlap at all: with two buffers, the producer can fill buffer B while the consumer drains buffer A, then they swap. With a depth of one there is only a single buffer, so the producer and consumer would contend for it and serialize again — no overlap. Going deeper (3, 4, …) hides more memory latency, because the producer can run further ahead of the consumer, right up until the SMEM budget runs out (228 KB/SM caps how many tiles you can stage). At depth 2 with warp specialization, the kernel hits 0.23 ms (~309×) — more than double the serial-pipeline Step 4, exactly because the two engines now run concurrently.

Interactive · warp-role choreography vs lockstep

This is the centerpiece. The left timeline is lesson 09's lockstep kernel: one team of threads does load → compute → write, so only one engine is ever lit. The right timeline is this lesson's warp-specialized kernel: the TMA producer (WG1.w3), the MMA consumer (WG1.w0), and the writeback warpgroup (WG0) run as independent lanes, and you can watch the producer fill stage k+1 while the consumer multiplies stage k. The four barriers are drawn as the two rings; step through the startup to see the producer's wait(phase=1) pass immediately while the consumer's wait(phase=0) blocks. The KPIs track the tensor-core busy fraction: lockstep wastes two-thirds of the tensor core (~33%), while specialization keeps it fed back-to-back in steady state (~100%) — the only idle left is the one-tile fill and drain, which amortize away as the K-loop grows. Driving that fraction to 100% is the whole point of the track.

Step 8 — the 2-CTA cluster: one cooperative MMA over a 256×256 tile

Warp specialization keeps a single SM's tensor core busy, but lesson 07 showed a bigger lever: pool two SMs. In Step 8 the two CTAs join a thread-block cluster and issue one cooperative tcgen05.mma with cta_group=2 that reads operands from both CTAs' SMEM and produces a single 256×256 output tile (versus 128×128 from one CTA).

The split mirrors lesson 07's cooperation pattern, by rows:

• A is split by M rows: CTA-0 owns output rows 0–127, CTA-1 owns rows 128–255. Each CTA holds its own half of A.
• Each CTA loads a different 128-row B slice, and the cooperative MMA reads across the pair via distributed shared memory.

// PSEUDOCODE → CuTe/PTX mapping for the cluster step
B_remote     = map_shared_rank(Bsmem, cta_id=1)   // peer's SMEM  ↔  mapa.shared::cluster
tma2mma_cta0 = tma2mma.remote_view(0)             // raise CTA-0's barrier from CTA-1

// TMA arrive byte count is multiplied by CTA_GROUP — each CTA loads its OWN bytes:
arrive.expect_tx(bytes_per_tile * CTA_GROUP)

// the SINGLE cooperative MMA is issued only by CTA-0 (cbx == 0) and wakes BOTH CTAs:
if cbx == 0:
    tcgen05.mma.cta_group::2(A_pair, B_pair -> TMEM_pair,
                             cta_mask = 0b11)      // 3 = both CTAs' barriers fire

ld2mma.init(128 * CTA_GROUP)                       // = 256: both warpgroups must drain
cluster_sync()                                     // replaces cta_sync() before TMEM dealloc

Three details are worth pinning down, because they are the ones easy to get wrong (and they are this lesson's exercises):

• Why multiply the arrive byte count by CTA_GROUP? Because each CTA issues its own TMA and delivers its own bytes into its own SMEM. The completion barrier for the pair must therefore expect the bytes from both CTAs' loads, so the expect_tx budget scales by the number of CTAs in the group. (This is Exercise 2: the byte count is per-delivered-data, and each CTA delivers a share.)
• Why is the MMA issued only by CTA-0? A cooperative cta_group::2 MMA is a single instruction spanning both SMs; if both CTAs issued it you would launch the math twice. So only CTA-0 (cbx == 0) elects to issue it, and it wakes both CTAs' completion barriers via cta_mask = 3 (binary 11) — one signal for the whole pair.
• Why ld2mma.init(128 * CTA_GROUP) = 256 and cluster_sync()? The writeback now spans both CTAs' warpgroups, so 256 threads must finish reading TMEM before the next tile may overwrite it; and the TMEM deallocation fence is now cluster-wide, so cluster_sync() replaces cta_sync().

The payoff: each CTA's B slice is now reused against the other CTA's A slice as well as its own, so the same loaded B bytes feed twice as much math — arithmetic intensity roughly doubles. That is the move lesson 07 promised, and it drops the kernel to 0.104 ms (~676×), already within ~10% of cuBLAS and close to the compute roof. On a B200 whose roofline ridge sits near 250 FLOP/byte (lesson 00), pushing AI up is pushing straight toward the ceiling.

Step 9 — multi-consumer: two math warps share one B tile

Step 8 doubled B's reuse by sharing it across two CTAs. Step 9 squeezes B again, this time within the cluster, by adding a second MMA consumer so two consumers share the same staged B tile. We grow the warpgroup count to WG_NUMBER=3 with NUM_CONSUMER=2, and the cluster output grows from 256×256 to 512×256 (split along M):

Two implementation facts come with the second consumer:

• PIPE_DEPTH=4. With two consumers draining each staged tile and a producer loading two A blocks per stage, a deeper ring keeps the producer far enough ahead to feed both consumers without stalling.
• EPI_N=64 epilogue chunking. Each writeback thread would otherwise hold the full 256 fp32 accumulator values = 1024 bytes per thread, which risks spilling to local memory (slow, off-chip). So the writeback is broken into 64-column chunks, keeping only 64 fp32 live at a time, and the TMEM→register→GMEM drain runs once per chunk.

The result: 0.094 ms (~744×) — cuBLAS parity. The final gain is only ~10%, and that is precisely what to expect this close to the compute ceiling: by Step 8 the kernel was already compute-bound, so there is very little memory stall left to hide. Halving B's per-FLOP load shaves the last sliver of memory traffic, and the kernel sits on the roofline.

Where the roofline ended up

Trace the regime across the ladder: the kernel was memory-bound until roughly Step 7 — the tensor core kept waiting on operands. The cluster steps (8 and 9) doubled arithmetic intensity twice over, removing the last big memory stalls, so by Step 8 the kernel is compute-bound: it is limited by how fast the tensor core can multiply, not by how fast HBM can deliver. That crossover is the whole point of the track's engine — feeding the tensor core was the entire kernel, and once it is fed, the math sets the wall-clock. cuBLAS lives at the same ceiling because it is solving the same feeding problem with the same hardware.

The bridge to Flash Attention 4

Everything assembled in this lesson — TMA loads into an SMEM ring, the tcgen05 cooperative MMA, TMEM readback, and warp-specialized barriers with opposite initial phases — carries straight into the next capstone. Flash Attention 4 reuses all of it. What changes is the schedule: a GEMM is one MMA repeated down the K-loop, but attention wedges an online-softmax step between two MMA phases (score S = QKᵀ, then value O += P·V) instead of repeating a single one. The machinery is identical; the dependency graph is fundamentally harder.