Part III · GEMM from tiled to SOTA

GEMM II — pipelining with TMA

Lesson 08's tiled GEMM is correct: right tiles, right layouts, right math, fp32 accumulate. It is also slow — and slow in a way that has nothing to do with the math. The whole kernel takes turns. Threads copy a tile, the tensor core chews it, threads copy the next, the tensor core waits. Two engines that could run side by side instead stand in a line, each stalling on the one before it. This lesson closes that gap without touching the data path: same tiles, same layouts, same instructions. What changes is when the work happens and by whom.

Forced move — the tiled baseline leaves the tensor core idle

Lesson 08 built the smallest correct GEMM and proved the thesis "every step changes how a hop happens, never the hops" — the path is always GMEM → SMEM → TMEM → reg → GMEM. But Step 3 left the tensor core idle for most of the clock. The synchronous loop is a relay: load a tile (all threads copy GMEM→SMEM, then cta_sync), compute (the elected thread issues the MMA), wait, load the next, wait. Each stage stalls on the one before it. Yet loading the next tile and computing the current one touch entirely separate hardware — the copy path and the tensor core never contend. The relay is pure self-imposed serialization.

Closing the gap needs no new data path. The forced move is three steps that change scheduling only: Step 4 hands the bulk GMEM↔SMEM transfers to TMA so the math threads never touch the load; Step 5 adds a software pipeline so the next tile has somewhere to land while the current one computes; Step 6 reshapes the launch into a persistent kernel so a fixed pool of CTAs amortizes setup and keeps operands hot in L2.

Step 4 TMA async load. Replace the all-threads copy with one engine-driven bulk transfer; track its completion with an mbarrier and a byte count instead of cta_sync.
Step 5 Software pipeline. Double-buffer SMEM (PIPE_DEPTH=2) so the load for tile k+depth can be in flight while the MMA for tile k runs. One barrier per stage; phase flips follow slot count.
Step 6 Persistent kernel. Launch one CTA per SM and loop over many tiles, with a tile scheduler that groups operand-sharing tiles to keep them hot in L2 — paying off lesson 08's reuse debt.

From the course (mapped to the real toolchain) The step kernels below are written as pseudocode in the MLC course's tile DSL and mapped to the production stack as we go. Step 4 swaps Ampere cp.async / LDGSTS for cp.async.bulk.tensor (TMA). Step 5 is the CUTLASS multi-stage pipeline PipelineTmaAsync with Stages = 2 (production uses 3–8 to hide more latency). Step 6 is the CUTLASS persistent kernel with PersistentTileSchedulerSm90 / Sm100 and a raster/threadblock swizzle for L2 locality. Read the DSL as the concept and the named CUTLASS/PTX as the contract.

Step 4 — hand the load to TMA

The tiled baseline loads a tile with every thread copying its share of GMEM into SMEM, then a cta_sync() so the whole CTA agrees the tile is in place before the MMA reads it. Those copy instructions are issue slots stolen from the math, and the addresses they compute are registers stolen from the accumulator. Lesson 04 gave us the way out: the TMA engine moves a whole rectangular tile on its own, off the issue path. The math threads just say "go" and walk away.

So the load shrinks to a single thread firing two bulk-tensor copies and announcing how many bytes are coming:

// PSEUDOCODE (tile DSL) — Step 4 load, mapped to PTX in the comments
if (tid == 0) {                                  // ONE thread issues the whole tile
    copy_async(Asmem, A, dispatch="tma");        // → cp.async.bulk.tensor … [Asmem], [A_map, {m,k}], [tma_bar]
    copy_async(Bsmem, B, dispatch="tma");        // → cp.async.bulk.tensor … [Bsmem], [B_map, {n,k}], [tma_bar]
    mbarrier.arrive.expect_tx(tma_bar, byte_count);  // announce the byte budget for this round
}
mbarrier.try_wait(tma_bar, phase);               // block until BOTH arrival + bytes hit zero (lesson 06)
//                  ^ → mbarrier.try_wait.parity.shared.b64

byte_count = (BLK_M*BLK_K + BLK_N*BLK_K) * F16_SIZE;   // A tile + B tile, fp16 = 2 B

This is lesson 06's machinery applied to a GEMM tile: arrive.expect_tx records the transaction byte budget and counts as the issuing thread's arrival; as the TMA engine streams the tile in, complete-tx updates drive the byte count to zero; the round completes — and the phase flips — only when arrival count and pending bytes both reach zero. The math path waits on that one barrier, then reads SMEM.

Exam-worthy: tid == 0, never elect_sync()

The byte budget must be announced to the mbarrier exactly once. It is tempting to reach for elect_sync() — it picks "one thread" and reads as the idiomatic way to do single-thread work. But elect_sync() elects one lane per warp. A Hopper/Blackwell warpgroup is 128 threads = 4 warps, so elect_sync() would elect four threads, and arrive.expect_tx would announce the byte count 4×. The barrier would then expect four times the bytes that the TMA engine will ever deliver — the pending-byte count never reaches zero and the try_wait deadlocks. Using tid == 0 picks exactly one thread across the whole warpgroup. (elect_sync() is the right tool when you want one issuer per warp — it is the wrong tool for a once-per-CTA announcement.)

Why cta_sync() no longer suffices — and the byte-count trap

cta_sync() (i.e. __syncthreads()) waits only for the CTA's own threads to reach the line, and it orders only their SMEM writes. It knows nothing about an in-flight TMA transfer — the copy is happening on an engine, not on a thread, so a thread barrier cannot tell whether the bytes have landed. Completion must be tracked by the mbarrier plus the byte count. And that byte count has to be exact: if expect_tx is set too small, the byte half hits zero before the last bytes arrive, the barrier releases early, and the MMA reads an incomplete tile (silent garbage). If it is set too large, the byte half never reaches zero and the wait never releases (deadlock). The byte count is part of the contract, not a hint.

What Step 4 buys, and what it does not: the load is now off the issue path, but the loop is still strictly serial — load → wait → MMA → wait → load …. We replaced who does the copy, not when it happens relative to the math. The tensor core still sits idle through every load. That is the wall Step 5 removes.

Step 5 — a software pipeline so the next tile has somewhere to land

To overlap a load with a compute, the next tile needs a buffer that is not the one the MMA is currently reading. So we double-buffer SMEM: give each operand a small ring of PIPE_DEPTH stages, indexed by a dimension-0 axis that is the pipeline stage.

// PSEUDOCODE — Step 5 allocation. PIPE_DEPTH = 2 (a double buffer).
constexpr int PIPE_DEPTH = 2;
Asmem(PIPE_DEPTH, BLK_M, BLK_K);          // dim-0 index = pipeline stage (the ring slot)
Bsmem(PIPE_DEPTH, BLK_N, BLK_K);

mbarrier  tma_bar[PIPE_DEPTH];            // ONE TMA barrier PER STAGE
mbarrier  mma_bar;                        // ONE MMA-completion barrier (1 slot, revisited every iter)

The barrier counts are the subtle part, and they fall straight out of lesson 06's "one barrier per stage, not per iteration." Each ring slot gets its own TMA barrier — PIPE_DEPTH of them. The MMA has a single barrier because there is only ever one MMA in flight. The K-loop runs in three phases:

// PSEUDOCODE — Step 5 mainloop. phase_tma, phase_mma start at 0.

// ── PROLOGUE: prime the pipe — prefetch the first PIPE_DEPTH stages ──
for (int s = 0; s < PIPE_DEPTH; ++s)
    if (tid == 0) tma_load(stage = s, k_tile = s);     // fire the first `depth` loads, no waiting

// ── STEADY STATE ──
for (int k = 0; k < K_TILES; ++k) {
    int stage = k % PIPE_DEPTH;                         // which ring slot holds tile k

    try_wait(tma_bar[stage], phase_tma);                // wait for THIS stage's load to land
    mma(stage, accum = (k != 0));                       // first MMA initializes the accumulator
    try_wait(mma_bar, phase_mma);  phase_mma ^= 1;      // 1 slot, revisited every iter → flip EVERY iter

    int next_k = k + PIPE_DEPTH;                        // the tile this slot will hold next
    if (next_k < K_TILES && tid == 0)
        tma_load(stage, next_k);                        // refill THIS slot for its next trip
    if (stage == PIPE_DEPTH - 1) phase_tma ^= 1;        // ring wrapped → flip the TMA phase ONCE per trip
}
// ── EPILOGUE: TMEM → registers → GMEM (the TMA-store path below) ──

Map to the toolchain: this is exactly CUTLASS's PipelineTmaAsync with Stages = 2. The library hides the spin loop and the phase bookkeeping behind a PipelineState; production code sets Stages to 3–8 to keep more loads in flight and hide more HBM latency. The structure — prologue prefetch, steady-state wait/compute/refill, per-stage barriers — is identical.

The phase-flip rule follows slot count

This is the one place Step 5 is easy to get wrong, and the rule is mechanical once you see it: flip a barrier's phase exactly when you come back around to a slot you have used before.

mma_bar — 1 slot

There is a single MMA barrier, revisited on every iteration. Every iteration is a fresh round on that one slot → flip phase_mma every iteration.

tma_bar[ ] — PIPE_DEPTH slots

The TMA barriers form a PIPE_DEPTH-element array. Each stage is revisited once per trip around the ring → flip phase_tma only when stage == PIPE_DEPTH − 1 (the ring has wrapped and the next iteration reuses stage 0 with a new load).

And the reason each stage needs its own barrier, not one shared barrier: a single shared TMA barrier could not distinguish "tile 0's first load completed" from "tile 0's second load completed" — both arrive on the same object. The consumer would wait on a stale completion (exactly the phase-reuse bug from lesson 06) and read a half-refilled slot → race or deadlock. One barrier per stage, with its own phase register, is what keeps the rounds distinct.

Honest caveat — this is prefetch, not yet true concurrency

Be precise about what Step 5 achieves. The TMA load for tile k+PIPE_DEPTH is queued in the same iteration as the MMA for tile k, so the copy engine is kept busy ahead of the math. But the loop still waits for the current MMA to finish before the next iteration proceeds (try_wait(mma_bar, …) is on the critical path). The load and the compute overlap only to the extent that the next load was already in flight — the math thread is not free to issue the next MMA while the current one runs, because the same threads do both. This is prefetch: the load latency is hidden behind compute, which is most of the win. It is not full load/compute overlap. Getting the two engines to run genuinely at once requires splitting the threads into producer and consumer roles — warp specialization, which is Step 7 in lesson 10.

Interactive · the pipeline Gantt — serial vs prefetch

Two lanes, one per engine: TMA (the copy engine) and MMA (the tensor core). Step 4 (serial) takes turns — load, wait, compute, wait — and the gaps are the tensor core sitting idle. Step 5 (prefetch) lets TMA run ahead so the MMA lane is fed nearly back-to-back, and the idle % drops sharply. The PIPE_DEPTH slider sets the ring size — how many loads can be in flight at once; here a load fits behind a single compute, so PIPE_DEPTH=2 already hides the latency and deeper rings mainly buy headroom for longer or burstier loads. (Either way this is prefetch: the MMA still gates each iteration — full overlap is lesson 10.)

Pipeline Gantt · TMA + MMA lanes · serial vs software-pipelined prefetch

A K-loop of tiles fed through a ring buffer. Serial = Step 4 (each MMA waits for its own load). Prefetch = Step 5 (loads run ahead by PIPE_DEPTH). The model uses a fixed load time and MMA time per tile; the engine-idle KPI is the fraction of the makespan the tensor core is not computing.

mode: PIPE_DEPTH: 2 K tiles: 8 load : MMA cost: 1.2×

makespan

—

tensor-core idle

—

MMA back-to-back?

—

loads in flight (max)

—

Step 6 — a persistent kernel that keeps operands hot

Steps 4–5 fixed one tile's schedule. Step 6 fixes how tiles are handed out across the whole problem. The tiled baseline used a 2D launch: one CTA per output tile. For a 4096² GEMM with 128×128 tiles that is 32×32 = 1024 ephemeral CTAs, each paying kernel setup (allocate SMEM, allocate TMEM, prime barriers) once and then vanishing after a single tile. The setup cost is paid 1024 times, and nothing the first CTA loaded survives for the next — every CTA re-reads its operands cold from HBM.

The persistent kernel inverts that: launch a fixed pool of CTAs sized to the machine — one per SM — and have each one loop over many output tiles. Setup is paid once per SM, not once per tile, and operands a CTA already touched can stay resident in L2 for the next tile it draws.

// PSEUDOCODE — Step 6 persistent kernel. SM_COUNT = 148 on B200.
constexpr int SM_COUNT = 148;                  // launch ONE CTA per SM (the whole pool)

tile_scheduler = ClusterPersistentScheduler2D(
        num_m_tiles, num_n_tiles,
        l2_group_size = 8,                     // group nearby tiles so they share operands in L2
        num_clusters  = SM_COUNT);

allocate_tmem(accum);                          // allocated ONCE — persists across every tile

while (tile_scheduler.valid()) {
    (m_st, n_st) = tile_scheduler.current();   // this CTA's current output tile
    phase_tma = 0;  phase_mma = 0;             // RESET phases — each tile is a fresh K-loop

    prologue_prefetch();                       // prime the pipe (Step 5)
    steady_state_K_loop();                     // the Step-5 pipelined mainloop
    epilogue_tma_store(m_st, n_st);            // TMA-store the result tile (below)

    tile_scheduler.next();                     // advance to the next tile in this CTA's list
}

Two sizing facts and two state-reset rules carry the whole step:

CTA count = SM count. Over-launch (more CTAs than SMs) and the extra CTAs queue behind resident ones, idling until a slot frees — the launch overhead you were trying to amortize comes back. Under-launch and some SMs never get a CTA — wasted hardware. One CTA per SM is the sweet spot; SM_COUNT = 148 on a B200 (a device-queried number — read it, don't hardcode it in production).
Phases reset to 0 at the start of every tile. Each output tile runs its own fresh K-loop, so phase_tma and phase_mma start over from 0. Forgetting this reset would carry a stale phase into the new tile's first try_wait → the lesson-06 reuse bug, one tile later.
TMEM allocation is NOT reset. The accumulator's TMEM is allocated once and persists across every tile the CTA processes. Re-allocating per tile would pay tcgen05.alloc repeatedly for no reason; the persistent kernel reuses the same TMEM region, zeroing it via the accum=false flag on each tile's first MMA instead.
l2_group_size = 8 pays off lesson 08's reuse debt. Step 3 of the tiled baseline planted a debt: neighbors recompute over the same A (row-band) and same B (column-band) operands. Grouping nearby tiles so a CTA runs them back-to-back keeps those shared operands hot in L2 — the second tile in a row-band finds its A already cached. This is the same operand-reuse idea TMA multicast attacks spatially (lesson 07); the scheduler attacks it temporally.

The arithmetic: 4096² → 1024 output tiles spread over 148 SMs ≈ 7 tiles per CTA. Map to the toolchain: this is the CUTLASS persistent kernel with PersistentTileSchedulerSm90 / Sm100 and a raster (threadblock-swizzle) ordering that is exactly the L2-locality grouping. And the static "next tile = the scheduler's formula" here is precisely what lesson 07's Cluster Launch Control upgrades: CLC replaces the precomputed assignment with hardware work-stealing, killing the tail when tile costs are uneven.

The TMA-store epilogue (used in Steps 4–6)

Once the K-loop finishes, the accumulator sits in TMEM and has to reach GMEM. The path is the reverse of the load, and it ends with TMA again — but the completion signal is different, and that difference is worth pinning down.

// PSEUDOCODE — TMA-store epilogue, mapped to PTX
copy_async(regs, accum, dispatch="tmem");        // TMEM → registers  (→ tcgen05.ld)
tcgen05.wait.ld;                                 // wait for the TMEM read to complete
regs_fp16 = cast_fp32_to_fp16(regs);             // narrow for the output dtype
write(Dsmem, regs_fp16);                         // registers → SMEM (staging buffer)

fence.proxy_async("shared::cta");                // make thread SMEM writes visible to the TMA proxy

if (tid == 0) {
    copy_async(D, Dsmem, dispatch="tma");         // SMEM → GMEM  (→ cp.async.bulk.tensor store)
    cp.async.bulk.commit_group();                 // close the group of outstanding bulk stores
    cp.async.bulk.wait_group(0);                  // wait until 0 groups remain → Dsmem reusable
}

Two things are load-bearing here. The fence.proxy_async("shared::cta") exists because the thread writes to Dsmem and the TMA store reads it through a different memory proxy; without the fence, the TMA proxy may not see the thread's writes yet and would copy stale bytes. (This is the "thread → engine handoff needs a fence" rule from lesson 06, on the store side.) And then the completion mechanism flips:

Load completion — mbarrier + byte count

On a load you know exactly how many bytes are coming in, so you announce that budget with arrive.expect_tx(bytes) and the barrier completes when the byte count hits zero. The consumer waits on a phase.

Store completion — commit_group + wait_group

On a store there is no expected byte count that makes sense — the bytes are going out, and the kernel doesn't need to know when they arrive, only when its stores are done so it can reuse Dsmem. So stores are tracked with commit_group (close a batch) + wait_group(0) (block until no batches remain).

This load/store asymmetry first appeared in lesson 04; here it is the concrete reason the epilogue looks nothing like the mainloop's wait. Loads track "how many bytes have landed"; stores track "are my outgoing transfers retired."

Where Step 6 leaves us — and the bottleneck for lesson 10

After three scheduling changes and zero data-path changes, the kernel is dramatically faster: the tensor core is fed off an engine, its loads are prefetched a ring ahead, and the launch is amortized over a persistent pool with operands kept hot in L2. But there is a ceiling baked into the structure. Step 6 still runs the whole pipeline with all threads marching in lockstep. One team of threads issues the TMA, waits, issues the MMA, waits, runs the epilogue — it is still the case that one set of threads is where all three engines meet, and the try_wait(mma_bar, …) on the critical path is the proof: the loop cannot issue the next MMA while the current one runs, because the same threads must do both.

Takeaway

Lesson 08's tiled GEMM was correct but took turns — the tensor core idled through every load. Three scheduling-only steps close the gap with no new data path. Step 4 hands the bulk copies to TMA, tracked by an mbarrier + an exact expect_tx byte count (issued by tid == 0, never elect_sync(), so the budget is announced once not four times; cta_sync() can't see an in-flight engine transfer). Step 5 double-buffers SMEM into a PIPE_DEPTH-stage ring with one barrier per stage; the phase-flip rule follows slot count — phase_mma flips every iteration (1 slot), phase_tma flips once per ring trip. This is prefetch, not full overlap. Step 6 launches one CTA per SM looping over many tiles, resetting phases per tile but persisting the TMEM accumulator, with l2_group_size grouping operand-sharing tiles to pay off lesson 08's reuse debt. The store epilogue uses commit_group/wait_group, not a byte count — the load/store asymmetry. The wall that remains: all threads still march in lockstep, so the two engines can't truly run at once.

Where this connects

The load engine and its expect_tx byte budget are lesson 04 · TMA; the phase bit, the try_wait semantics, "one barrier per stage," and the reuse bug are all lesson 06 · mbarriers — Step 5's ring is the PIPE_DEPTH=2 pipeline that lesson promised. The persistent kernel's static tile scheduler is the thing lesson 07 · clusters & CLC upgrades to hardware work-stealing, and l2_group_size is the temporal cousin of that lesson's TMA multicast.

Step 6 still runs the whole pipeline with all threads marching in lockstep — one team of threads is still where all three engines meet. Split them into producer and consumer roles so the copy engine and the tensor core finally run at the same time. That is lesson 10 · GEMM III: warp specialization & clusters — the climax of the ladder.