Part II · the five hardware primitives

TMEM — the accumulator's new home

Lesson 04 fixed the operand-supply problem: TMA streams A and B into shared memory so the math warps never touch the load. But there is a second starvation pressure, and it sits on the output side of the tensor core. On Blackwell, the cure is a third kind of on-chip memory that exists for one job: holding the accumulator. This lesson is about TMEM — what it is, why the kernel must name it explicitly, and the three instructions that move data in and out of it.

What TMEM is

TMEM (Tensor Memory) is a Blackwell-only, 2D scratchpad that sits on each SM and is scoped per-CTA. It is the dedicated home for the tcgen05.mma accumulator (and, for block-scaled MMA, the scale factors). It is not the register file, not shared memory, and not L2 — it is a fourth address space that exists specifically to decouple the accumulator's size from each thread's register budget.

Its physical shape is rigid and worth memorizing, because every layout you write against it is phrased in these two axes:

• 128 hardware Lane rows (indexed 0–127), and
• up to 512 Col columns (indexed 0–511), each Col 32 bits wide.

The 128×512 shape is not arbitrary: it mirrors how tcgen05.mma writes its accumulator. The 128 Lane rows line up with the M dimension of the output tile, and the Col columns hold the N dimension, 32 bits per fp32 element. So an accumulator of width N occupies a 128 × N block of the grid.

TMEM is part of the layout interface, not a hidden cache

This is the conceptual hinge. Unlike a CPU cache — which the program never names — TMEM lives inside the tile-layout semantics. The CuTe-style layout notation from lesson 02 grew two new named axes for it: TLane and TCol. An accumulator that is 128 rows tall and N columns wide is written

S[(128, N) : (1@TLane, 1@TCol)]

Read that the same way as any other layout: the shape is (128, N), and the strides say "step one position along the TLane axis to move down a row, step one position along the TCol axis to move right a column." Because TMEM coordinates appear directly in the layout algebra, the kernel must name it, allocate it, and layout-match it exactly like it does for SMEM. There is no silent backing store doing the right thing on your behalf. If the layout you declare against TMEM disagrees with what the MMA wrote there, you read scrambled accumulator values — the same "layout is part of the instruction interface" rule from lessons 02 and 03, now applied to the output side.

Allocation — a budgeted CTA resource

Registers are assigned by the compiler; TMEM is not. The kernel explicitly allocates and frees it, and the mental model is "a budgeted CTA resource, much like shared memory." The protocol:

1. One warp in the CTA requests a range of TMEM columns. The request is in units of 32 columns — you ask for a multiple of 32, and the hardware rounds your request up to the next multiple per its allocation rules.
2. The CTA receives a base TMEM address, which every subsequent tcgen05 instruction uses to locate its rows and columns.
3. When the work is done, the kernel frees the range explicitly.

The toolchain instructions are vendor-real PTX:

// one warp in the CTA does the allocation; ncols rounded up to a multiple of 32.
// the base TMEM address is written back to an SMEM slot the CTA can read.
tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32   [tmem_base_ptr], ncols;

// ... run the K-loop: tcgen05.mma accumulates into [tmem_base + col offsets] ...

// give back the allocation permit, then free the columns.
tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;
tcgen05.dealloc.cta_group::1.sync.aligned.b32             tmem_base, ncols;

That relinquish_alloc_permit step is a real correctness gotcha you will meet again in the debugging lesson: allocation is gated by a permit, and a CTA must surrender the permit before it (or a successor) can cleanly tear down. For now the takeaway is simply that TMEM is requested and released like a heap allocation, in 32-column granules, by exactly one warp on behalf of the whole CTA.

Access — three dedicated instructions

Here is the rule that surprises people: ordinary ld.shared / st.shared cannot touch TMEM. It is a separate address space, so the shared-memory load/store instructions simply do not address it. Blackwell adds three dedicated instructions, one per direction of movement:

tcgen05.ld — and why the epilogue code is reusable

The read-back is the one to understand in detail, because it is what lets Blackwell reuse the epilogue code written for Ampere and Hopper. A single tcgen05.ld lowers to four warp-level operations — one per warp, and each warp handles 32 of the 128 Lane rows. (Four warps × 32 rows = the 128 Lane rows; this is the same four-way split a warpgroup uses.) The instruction has a load-shape family and a repeat factor:

// load-shape family:   .16x64b  .16x128b  .16x256b  .32x32b  .16x32bx2
// repeat factor:       .x1 .x2 .x4 ... .x128   (how many shapes per issue)
tcgen05.ld.sync.aligned.16x256b.x4.b32   {d0, d1, ... },  [tmem_addr];
tcgen05.wait::ld.sync.aligned;            // accumulator now in registers

The payoff is the lane mapping. For the common epilogue, lane l receives values from TMEM row l/4 and two columns — which reproduces exactly the m8n8 fragment that lesson 03 traced through every generation. Because the accumulator arrives back in registers in that same shape, the Blackwell epilogue can reuse the identical register-level cast-and-store code as an Ampere mma.sync or a Hopper wgmma — even though the accumulator spent the entire K-loop living in TMEM instead of registers. The home changed; the fragment shape the epilogue sees did not.

Interactive · the TMEM map and the tcgen05.ld read-back

The grid below is one SM's TMEM: 128 Lane rows × 512 Col columns. Pick an accumulator width N and watch it occupy a 128×N block, allocated in units of 32 columns (the hardware rounds your request up). Toggle a second MMA with M = 64 to see a strided Lane allocation that leaves gaps. Then press animate tcgen05.ld to watch the four-warp read-back sweep the rows, each warp owning 32 lanes, with lane l mapping to row l/4.

Async completion — you cannot read TMEM early

All three instructions are asynchronous. They return immediately and signal completion separately:

• tcgen05.ld → wait with tcgen05.wait::ld.
• tcgen05.st → wait with tcgen05.wait::st.
• tcgen05.cp → completes via a commit group + mbarrier, exactly like tcgen05.mma.

Block-scaled MMA — why scale factors take an extra hop

Lesson 03 introduced Blackwell's block-scaled formats (MXFP8, NVFP4), where a block of elements shares one scale factor. Those scale operands have a specific home: they are read from TMEM, while the data operands A and B stay in SMEM. That asymmetry forces an extra step, because TMA only lands data in SMEM:

So tcgen05.cp exists largely to perform that SMEM→TMEM hop for the scales, repacking them into the layout the tensor core reads. The packing is compact and a little clever. A 128-row scale vector folds into just 32 Lane rows — lane = r mod 32, column = r/32 — and is then broadcast ×4 so all four warps of the reading warpgroup can find a copy:

S[(32, sf) : (1@TLane, 1@TCol)]  +  R[4 : 32@TLane]

The R[4 : 32@TLane] part places four replicas at Lane offsets 0, 32, 64, 96. Each of the four warps reads from its own 32-lane window, and each window finds a full copy of the scales — no warp has to reach across the warpgroup for them. It is the same "lay it out so the consumer finds it locally" discipline the whole track keeps returning to, applied to scale factors.

The full data path

Now assemble every hop a block-scaled tcgen05 GEMM tile makes. This is the capstone picture for the lesson, and it is the skeleton of the GEMM mainloop you will build in Part III:

Three layout contracts have to stay matched all the way down this path, and a mismatch in any one reads scrambled bytes rather than crashing: