Part II · the five hardware primitives

TMA — the async bulk copy engine

The tensor core reads its operands from shared memory through a descriptor. But getting the tile into shared memory the old way means the math threads do the copy themselves — computing addresses, issuing loads, bookkeeping the swizzle — burning the exact instruction slots the tensor core needs. The forced move is to take the copy off the threads entirely and hand it to a dedicated hardware unit. That unit is the Tensor Memory Accelerator.

What the issuing thread actually does

A TMA copy is issued by a single thread. That thread does not loop over elements and it does not compute per-element addresses. It provides exactly two things:

1. a tensor-map descriptor — a small object that records the tensor's shape, its strides, the element size, the tile shape to copy, and the swizzle mode; and
2. the SMEM destination address — where the tile should land.

Then it issues the copy and keeps going. The rest of the CTA is free to compute on a tile that already arrived. The engine walks the rectangular region described by the descriptor, reads it from global memory, and writes it into shared memory — asynchronously, off to the side of the warps.

Contrast this with the Ampere world. Ampere's cp.async (the LDGSTS "load-global-store-shared" instruction) is still thread-issued: every lane issues its own async copy of its own slice, so the warp still spends instructions generating 32 addresses even though the store-to-SMEM now bypasses the register file. TMA collapses all of that into one issue from one thread describing the whole tile. That is the difference between "async per element, still on the threads" and "one engine, off the threads."

The descriptor records the layout, including the swizzle

The swizzle field is the load-bearing one. Swizzle is applied automatically as the bytes land in SMEM — there is no separate "swizzle pass" the kernel runs afterward. The TMA engine permutes bytes on the write so the tile arrives already in the layout the tensor core expects to read. This is the other half of lesson 03's iron rule made concrete: the swizzle mode in the TMA descriptor must match the swizzle mode named in the MMA descriptor (per lessons 02–03), or the tensor core reads scrambled bytes — a finite-garbage bug, not a crash.

Completion is where loads and stores split apart

Issuing the copy is the easy half. The hard half — and the central teaching point of this lesson — is knowing when it finished, because the consumer cannot touch the tile just because the instruction was issued. The data is in flight. And here loads and stores complete through two different mechanisms, for a simple reason: a load has an in-kernel consumer that must wait; a store usually does not.

The rule, in one line: loads wait via an mbarrier and a byte count; stores wait via a commit/wait group. Keep these straight — using the wrong one is a classic source of silent corruption (lesson 13 catalogs it).

Load completion — the mbarrier and the byte budget

The asymmetry is worth slowing down on, because it is not "wait until the instruction retires." A TMA load completes through an mbarrier (a shared-memory async barrier; lesson 06 covers it in full — this is a preview) in a precise sequence:

1. init / reuse an mbarrier object for this pipeline stage.
2. mbarrier.arrive.expect_tx(bytes) — this does two things at once: it records how many bytes the barrier should expect (the tile's byte count), and it counts as the issuing thread's own arrival.
3. issue the TMA load, pointed at that same mbarrier.
4. the engine performs complete-tx updates as bytes arrive, decrementing the pending-byte counter.
5. the barrier's phase flips only when both conditions hold: the arrival count is satisfied and the pending byte count has reached zero.
6. the consumer waits on that phase; when it flips, the tile is guaranteed fully present and safe to read.

So the completion condition is a conjunction — arrivals done AND bytes == 0. The byte budget is why a plain thread-only barrier cannot do this job: __syncthreads() can wait for threads, but it has no notion of "200 KB are still in flight from a copy engine." The expect_tx count is the hook that lets a barrier wait on a data-movement event rather than a thread-arrival event. (This is exactly why TMA forces lesson 06 into existence.)

Store completion — commit group and wait group

A TMA store typically writes the kernel's final output back to GMEM, where nothing inside the kernel reads it again. There is no consumer to release — so an mbarrier would be ceremony with no reader. What the kernel does need to know is when the SMEM region it stored from is free to be overwritten by the next tile. That is a producer-side question, answered by a lighter mechanism:

// issue one or more bulk-tensor stores SMEM → GMEM …
cp.async.bulk.tensor.2d.global.shared::cta  [tmap, {x,y}], [smem_src];
// … then bundle everything issued so far into a group:
cp.async.bulk.commit_group;
// … later, block until at most N groups are still outstanding:
cp.async.bulk.wait_group 0;          // 0 ⇒ all prior stores have drained
// after the wait, the SMEM source region is safe to reuse.

No byte count, no mbarrier, no phase. You commit a batch and later wait for the batch to drain. That is the whole store protocol.

Interactive · how a TMA load actually completes

The model below runs the load-completion sequence step by step. In TMA mode, one elected thread announces expect_tx(bytes) and issues the copy, then is free; the engine fills the byte budget on its own lane; the mbarrier phase bit flips only when bytes reach 0 and the arrival is in; then the MMA consumer is released to read. Flip to thread-copy mode to watch the alternative: the math warp itself drives the copy, lane by lane, with the tensor core idle the entire time. Step through, or hit play.

Two things the animation is built to make obvious. First, in TMA mode the issuing thread's lane goes idle the instant it issues — the work moves to the engine lane — whereas in thread-copy mode the math lane stays pinned to the copy. Second, the phase bit does not flip early: it waits for the byte counter to hit exactly zero. If you (wrongly) released the consumer on "instruction issued," it would read a half-arrived tile. That single guarantee — flip only at bytes == 0 ∧ arrived — is the entire reason a byte-counting barrier exists.

Why this only pays off inside a pipeline

Issuing one async copy and then immediately waiting for it buys nothing — you've just renamed a synchronous load. TMA earns its keep when the copy for a future tile is in flight while the tensor core computes the current one. That is a pipeline, and it needs at least two SMEM buffers ("stages"):

• one SMEM stage holds the tile currently being consumed by the MMA;
• another stage is being filled by a TMA load for an upcoming iteration;
• as the loop advances, the roles rotate (stage k mod num_stages);
• before the MMA reads a stage, it waits on that stage's load barrier (the mbarrier from above);
• before a TMA load overwrites a stage, the kernel ensures the previous consumer is done with it.

That choreography is precisely why TMA and mbarriers always appear together in Hopper and Blackwell kernels: TMA moves the bytes off the threads, and the mbarrier is the only thing that can tell the consumer "your stage is ready" and the producer "your stage is free." Neither is useful without the other. We build this double-buffered mainloop explicitly in the GEMM ladder; here, just hold the shape of it.

The real toolchain — PTX and the descriptor object

None of this is a course abstraction; it is the actual CUDA/PTX surface. The device-side instruction family is:

// TMA load: GMEM → SMEM, multidim, descriptor-driven, swizzled,
// completion tracked by an mbarrier via complete-tx byte updates.
cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes
    [smem_dst], [tensor_map, {coord_x, coord_y}], [mbar];

// the matching arrival that announces the byte budget (see lesson 06):
mbarrier.arrive.expect_tx.shared::cta.b64  _, [mbar], ;

// TMA store: SMEM → GMEM, drained via commit/wait group:
cp.async.bulk.tensor.2d.global.shared::cta  [tensor_map, {x,y}], [smem_src];
cp.async.bulk.commit_group;
cp.async.bulk.wait_group  0;

A few real distinctions worth nailing down:

• cp.async.bulk (raw, 1-D) vs cp.async.bulk.tensor (multi-dimensional, descriptor-driven, swizzle-aware). The plain cp.async.bulk moves a contiguous byte range; the .tensor form is the one that reads a rectangular tile through a tensor map and applies the swizzle. Dimensionality runs .1d … .5d.
• The descriptor is a CUtensorMap object. It is typically built on the host with cuTensorMapEncodeTiled(...) and passed to the kernel as a const __grid_constant__ CUtensorMap parameter — encoding the shape/strides/tile/swizzle once, on the CPU, rather than recomputing addresses per thread on the GPU.
• The multicast variant .multicast::cluster lets one load feed several CTAs at once — a preview of lesson 07's clusters and distributed shared memory.
• Contrast the whole thing with Ampere's cp.async (LDGSTS), which is thread-issued and per-slice. TMA is the generational step from "async copy, still on the threads" to "a copy engine, off the threads."

// host side: encode the tensor map once, pass it as a grid-constant kernel arg
CUtensorMap tmap;
cuTensorMapEncodeTiled(&tmap, CU_TENSOR_MAP_DATA_TYPE_FLOAT16,
    /*rank=*/2, gmem_ptr, gmem_extents, gmem_strides,
    tile_extents, /*elem_strides=*/{1,1},
    CU_TENSOR_MAP_INTERLEAVE_NONE,
    CU_TENSOR_MAP_SWIZZLE_128B,            // MUST match the MMA descriptor
    CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
    CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
// kernel signature:
__global__ void gemm(const __grid_constant__ CUtensorMap tmapA, ...);