Part III · assemble it — GEMM from tiled to SOTA

GEMM I — the tiled baseline

Lessons 02–07 handed us five primitives in isolation, all expressed in the layout algebra of lesson 02: the tcgen05 tensor core, TMA, TMEM, mbarriers, and clusters. Knowing each one is not the same as making them cooperate. This lesson assembles them into the first working GEMM — deliberately the slowest correct one — and lets its limitations name every step that follows.

The one invariant the whole ladder hangs on

Before any code, fix the mental model. A GEMM moves operands through a fixed chain of memories, and that chain never changes. Every step on the ladder changes only how one hop happens — never which hops exist:

The thesis of the ladder: the hops never change; every step only changes how one hop happens. Step 1 fills all five boxes the slow way. Step 4 (lesson 09) changes only the GMEM→SMEM hop (thread-copy → TMA). Step 7 (lesson 10) changes only who drives each hop (one warpgroup → specialized warps). Keep this strip in your head; it is the map for the next three lessons.

The problem we are solving, fixed for the whole ladder:

Operands are fp16; we accumulate in fp32 and cast back to fp16 only at writeback, to minimize rounding error across the K contraction. Storing B as N×K (row-major) rather than K×N is deliberate: it means a CTA loads a row-band of A and a row-band of B — and a row-band of B is a column-band of Bᵀ. Both operands stream in the same row-major shape, which is exactly what the tensor core's descriptor wants. (This is the row.col orientation you saw on the mma instruction back in lesson 03.)

Step 1 — the smallest correct tile

Compute one 128×128 output tile, with K=64, on a 1×1 grid. One CTA = 128 threads = 4 warps = one warpgroup. No K-loop, no spatial tiling, no async. Just: get bytes in, do one tensor-core op, get bytes out, and check the answer. The kernel has exactly four pieces.

1. ALLOCATE the on-chip working set: a TMEM address handle for the accumulator, one mma_bar mbarrier, and two swizzled SMEM tiles — Asmem 128×64 and Bsmem 128×64, both fp16, SWIZZLE_128B. (CuTe: a TiledMMA + a SharedStorage struct of swizzled SMEM layouts; the TMEM handle comes from tcgen05.alloc, see lesson 05.)
2. LOAD GMEM→SMEM with an all-threads cooperative copy — every one of the 128 threads strides over the tile and copies its share — then cta_sync() to publish those SMEM writes before any thread reads them for the MMA. (This is the classic __syncthreads()-guarded SMEM staging from gpu_kernels; in PTX the copies are plain global loads + shared stores.)
3. DISPATCH the MMA from one elected thread: if warp_id==0: if elect_sync(): then a single full-tile tensor-core op, and commit it to the barrier. (PTX: tcgen05.mma then tcgen05.commit.mbarrier::arrive [mma_bar].)
4. WRITE BACK TMEM→reg→GMEM: copy the accumulator out of TMEM into registers, wait for it, cast fp32→fp16, and write. Each of the 128 threads writes one output row — warp i, lane l writes row i·32 + l. (PTX: tcgen05.ld then tcgen05.wait::ld; the cast-and-store epilogue is ordinary register code, identical to the Ampere/Hopper path — which is the whole point of the l/4 fragment layout from lesson 05.)

// PSEUDOCODE — Step 1: one 128×128 tile, K=64, grid 1×1, 1 warpgroup (128 threads)
//   CuTe equiv: a single TiledMMA over one k-tile with a naive cooperative copy.

const BLK_M = 128, BLK_N = 128, K = 64;

// ── (1) ALLOCATE ──────────────────────────────────────────────
tmem    acc   = tmem_alloc(BLK_M, BLK_N);        // fp32 accumulator, lives in TMEM
mbar    mma_bar = mbarrier_init(count=1);        // completion signal for the MMA
smem    Asmem = smem_alloc(BLK_M, K, fp16, SWIZZLE_128B);
smem    Bsmem = smem_alloc(BLK_N, K, fp16, SWIZZLE_128B);   // B is N×K → a row-band

// ── (2) SYNCHRONOUS LOAD  GMEM → SMEM ─────────────────────────
cooperative_copy(Asmem, A_gmem_tile);            // all 128 threads stride-copy
cooperative_copy(Bsmem, B_gmem_tile);
cta_sync();                                      // PUBLISH smem writes before MMA reads

// ── (3) MMA DISPATCH  (one elected thread) ────────────────────
if (warp_id == 0) {
    if (elect_sync()) {                          // exactly one lane of warp 0
        gemm_async(acc, Asmem, Bsmem,            //  PTX: tcgen05.mma
                   accum=False,                  //  False ⇒ initialize/overwrite TMEM
                   dispatch=tcgen05, cta_group=1);
        tcgen05.commit(mma_bar);                 //  arms the barrier for this MMA
    }
}
mbarrier_wait(mma_bar, phase=0);                 // whole CTA waits for the math

// ── (4) WRITEBACK  TMEM → reg → GMEM ──────────────────────────
reg f32_acc[…] = copy_async(acc);                //  PTX: tcgen05.ld   (TMEM → registers)
tcgen05.wait.ld;                                 //  accumulator now in registers
reg f16_out[…] = cast_f32_to_f16(f32_acc);       //  round once, at the end
int row = warp_id*32 + lane_id;                  //  each of 128 threads owns one row
store_row(D_gmem, row, f16_out);

Two design choices in here are not arbitrary; they are forced, and worth pausing on.

Why accum=False. The accumulator in TMEM holds whatever was there before. accum=False tells the tensor core to overwrite rather than add — it initializes the accumulator with this MMA's result. With a single K-tile that is exactly what we want. (In Step 2 the very first iteration will use False and every later one True.)

Why one elected thread issues the MMA. This is the most counter-intuitive line for anyone coming from the SIMT model, where "all threads run the instruction" is the rule. A tcgen05.mma is itself a cooperative, full-tile operation — one instruction commands the whole 128×128×64 multiply. If all 128 threads issued it, you would launch the same full-tile work 128 times over. So exactly one lane issues it (warp_id==0 picks warp 0; elect_sync() picks one lane within it). This is the Blackwell tensor-core programming model from lesson 03, now load-bearing. (Subtle but it bites in lesson 09: elect_sync elects one lane per warp. Inside a single warp that is one thread; across a whole warpgroup it is four. Step 1 is safe because we already narrowed to warp_id==0 first.)

Step 1's four limitations — the list that drives the whole ladder

Step 1 is correct and useless. Name its limitations precisely, because each one is a later step. This list is the syllabus for lessons 08–10:

1. Only a single K-tile. K is pinned to 64; we cannot contract a realistically large K. → fixed in Step 2 (this lesson)
2. Only a single output tile. M and N are pinned to 128 — the grid is 1×1. → fixed in Step 3 (this lesson)
3. Synchronous copies, not TMA. Math warps spend their issue slots copying bytes. → Step 4 (lesson 09)
4. No overlap of data movement with compute. Load, then compute, then write — strictly serial; the tensor core idles during every copy. → Steps 5–7 (lessons 09–10)

This lesson closes (1) and (2). Lesson 09 takes (3) and (4). Everything below is just walking down this list.

Step 2 — the K-loop: accumulate over a large K

To contract a large K, tile it: process K_TILES chunks of 64, and accumulate into the same TMEM slot. The A/B SMEM buffers are fixed-size and get overwritten each iteration (load tile i, consume it, load tile i+1 on top). The TMEM accumulator, by contrast, persists across all iterations — that is the whole reason the accumulator lives in TMEM and not in the SMEM we keep clobbering.

The only new ingredient is the accum flag: accum = (i != 0). Iteration 0 initializes the accumulator (overwrite); every later iteration adds to it.

// PSEUDOCODE — Step 2: K-loop accumulation into a persistent TMEM accumulator.
//   CuTe equiv: the CUTLASS mainloop, in its fully-SYNCHRONOUS (un-pipelined) form.

tmem acc = tmem_alloc(BLK_M, BLK_N);
mbar mma_bar = mbarrier_init(count=1);
int  phase_mma = 0;                              // ← expected phase, kept in a register

for (int i = 0; i < K_TILES; ++i) {
    cooperative_copy(Asmem, A_tile[i]);          // OVERWRITE the same buffer each iter
    cooperative_copy(Bsmem, B_tile[i]);
    cta_sync();

    if (warp_id == 0 && elect_sync()) {
        gemm_async(acc, Asmem, Bsmem,
                   accum=(i != 0),               // i==0: overwrite · i>0: add
                   dispatch=tcgen05, cta_group=1);
        tcgen05.commit(mma_bar);
    }
    mbarrier_wait(mma_bar, phase_mma);           // block until THIS MMA completes
    phase_mma ^= 1;                              // ← TOGGLE — see the table below
}
// ... writeback identical to Step 1 ...

The checkpoint: the phase-flip table

Here is where Step 2 reuses one object — the single mma_bar — across every iteration. That reuse is exactly the situation lesson 06 warned about, now made concrete. The barrier has an internal phase bit that flips each time the MMA completes; mbarrier_wait(bar, p) blocks until the barrier's bit differs from p. So the register phase_mma must always name "the phase I expect to leave behind this round," and we toggle it after every wait. Trace it:

The wait in iteration i passes the bit the barrier was sitting in before iteration i's MMA. The MMA completes, the bit flips away from that value, the wait returns, and we toggle phase_mma to match — ready to do it again next round. Step through the widget below to watch the two values march in lockstep.

Interactive · the phase-flip table, and the bug that skips the wait

Step through the K-loop one iteration at a time. The left column is the phase_mma register; the right is the barrier's internal bit. Watch the MMA arm the barrier, the bit flip on completion, the wait return, and the register toggle to track it. Then flip "forget phase ^= 1": the register freezes at 0, and from iteration 1 on the wait returns before the MMA finishes — the panel flags every early return in red. That is the silent corruption, made visible.

Step 3 — spatial tiling: one CTA per output tile

Step 2 still computes a single 128×128 tile of D. To cover an M×N output, launch a 2D grid — one CTA per output tile — and let each CTA compute its own tile by offsetting where it reads A and B and where it writes D. The grid grows from [1, 1] to [M / BLK_M, N / BLK_N]; nothing about the inner kernel (Step 2's K-loop) changes except the base offsets.

// PSEUDOCODE — Step 3: 2D grid, one CTA per output tile.
//   CuTe equiv: the classic tiled grid launch, BEFORE persistent scheduling.

grid = [ M / BLK_M, N / BLK_N ];                 // e.g. 4096/128 = 32  →  32 × 32 = 1024 CTAs
//      ── each CTA: ──
int bx, by = cta_id();                           // this CTA's tile coordinates
int m_st = bx * BLK_M;                           // row offset into A and D
int n_st = by * BLK_N;                           // row offset into B (= col offset into Bᵀ)

for (int i = 0; i < K_TILES; ++i) {
    cooperative_copy(Asmem, A[m_st .. m_st+BLK_M, i]);   // THIS CTA's A row-band
    cooperative_copy(Bsmem, B[n_st .. n_st+BLK_N, i]);   // THIS CTA's B row-band
    cta_sync();
    if (warp_id == 0 && elect_sync()) {
        gemm_async(acc, Asmem, Bsmem, accum=(i != 0), dispatch=tcgen05, cta_group=1);
        tcgen05.commit(mma_bar);
    }
    mbarrier_wait(mma_bar, phase_mma);
    phase_mma ^= 1;
}
store_tile(D[m_st .. , n_st .. ], acc);          // writeback to THIS CTA's tile

For M = N = 4096 with 128-tiles, that is a 32×32 grid = 1024 CTAs, each independently doing a Step-2 K-loop. It works, it covers the whole output, and it is still entirely synchronous.

Where Steps 1–3 leave us

We have a correct, general GEMM: any M, N, K, covered by a 2D grid of CTAs, each accumulating across the full K with a correctly phase-tracked barrier. It would pass a numerical check against cuBLAS. And it is slow — it sits near the bottom of the performance ladder lesson 00 previewed, because we have only addressed limitations (1) and (2). Look back at the memory strip: every hop is still done the slow way, and they happen strictly in series.

Hand lesson 09 the bottleneck. Our loads are synchronous, and nothing overlaps: the cooperative copy uses the math warps' own issue slots, and the tensor core sits idle for the entire duration of every cooperative_copy before it gets to run. We are paying for the copy and the compute end-to-end when they could happen at the same time. The forced move is to hand the copy to a dedicated engine — TMA — so the warps stop touching the load, and then to pipeline it so the next tile streams in while the current one computes.