CUTLASS Deep Dive: From CuTe Layouts to Blackwell tcgen05

Introduction

CUTLASS is NVIDIA’s C++ template library for writing fast GEMM-like kernels. But “CUTLASS” today is really two libraries stacked on top of each other:

• CuTe — a lower-level abstraction built around layouts and tensors. It is what FlashAttention 3, FlexAttention, and every modern CUTLASS kernel is written in. CuTe is the algebra.
• CUTLASS 3.x — collective builders, warp-specialized kernels, tile schedulers, and fusible epilogues. CUTLASS is the framework that composes CuTe pieces into shippable kernels.

This post is a long-form, code-first tour of both layers. We start with the tiled_copy.cu tutorial to understand what a layout actually is, walk through SM80 sgemm to see pipelined shared memory, then Hopper’s wgmma + TMA, then Blackwell’s tcgen05 + TMEM + 2SM MMA, and finish with a survey of how the higher-level CUTLASS API assembles all of this into production-grade GEMM variants (FP8, mixed dtype, StreamK, grouped, fused epilogues).

All code snippets are pulled verbatim from cutlass/examples/cute/tutorial/ and cutlass/examples/.

🧩 Try it yourself — I packaged every concept in this post as a 9-day puzzle repo: github.com/supercharleszhu/cutlass-puzzle. Each day is a buildable .cu with critical lines replaced by // TODO: and a reference solution.cu next to it. Clone it, fill in the blanks, and you’ll actually write CuTe code — not just read about it.

Part 1 — Layouts: the Algebra Under CuTe

CuTe has exactly one core idea: a layout is a function from coordinates to offsets.

A layout is a pair of a shape and strides of matching rank. You can evaluate a layout like a function:

Layout(shape=(4,2), stride=(1,4))
= function mapping (m,n) -> m*1 + n*4
layout(3,1) = 3*1 + 1*4 = 7

Everything else in CuTe — tensors, partitioning, tiled copies, MMA fragments, swizzles — is built out of layouts and a few operations on them (composition, divide, product, coalesce). If you are comfortable with function composition, CuTe’s type signatures stop looking scary.

A Tensor is a pair (Engine, Layout) where the engine is the underlying pointer (gmem/smem/rmem). Creating one looks like this:

// From tiled_copy.cu:154
Tensor tensor_S = make_tensor(
    make_gmem_ptr(thrust::raw_pointer_cast(d_S.data())),
    make_layout(tensor_shape));

↪ blog/01_tiled_copy.cu#L162

⚠️ make_tensor does not copy data. It’s a zero-cost view — a lightweight wrapper around the existing pointer plus the layout metadata. No allocation happens, no bytes move, and the resulting Tensor aliases d_S. Mutating tensor_S(i, j) mutates the original buffer. Think of it as std::span, not std::vector. The same is true for make_gmem_ptr / make_smem_ptr / make_rmem_ptr — they just tag a pointer with a memory-space type so the compiler picks the right load/store instructions later.

tensor_shape is make_shape(256, 512) — a dynamic 2D shape. Because no stride was given, make_layout picks the default compact column-major stride (1, 256). You can query it, slice it, tile it, compose it with other layouts — all at compile time as long as the shapes and strides are static.

Static vs Dynamic Extents

Convention throughout CUTLASS:

• Lowercase names (m, n, k) are dynamic — runtime integers.
• Uppercase names (M, N, K) or Int<128>{} are static — compile-time constants propagated through types.

Static extents let the compiler constant-fold strides, inline unroll counts, and emit LDGSTS / LDSM / MMA instructions with the right immediate operands. Most of the CuTe speed comes from making as much static as possible.

tiled_divide — the Workhorse Partitioning Op

The central partitioning operation is tiled_divide, which takes an (m,n) tensor and a (M,N) static tile and returns a ((M,N), m', n') tensor:

// From tiled_copy.cu:163
auto block_shape = make_shape(Int<128>{}, Int<64>{});
// Tile the tensor (m, n) ==> ((M, N), m', n') where (M, N) is the static
// tile shape, and modes (m', n') correspond to the number of tiles.
Tensor tiled_tensor_S = tiled_divide(tensor_S, block_shape);  // ((M, N), m', n')

↪ blog/01_tiled_copy.cu#L171

This is purely a layout transformation — no data is moved, no copy happens. The same pointer is now viewed as a 3D (rank-3, with mode-0 grouped) tensor where each (_, blockIdx.x, blockIdx.y) slice is the (M,N) static tile that one thread block will process.

CuTe’s local_tile and local_partition are other flavors of the same idea: express a nested hierarchy (CTA tile → thread partition → MMA fragment) as layout composition.

Part 2 — Tiled Copies: Coalesced Loads, Vectorized

Once you have a static tile, the next question is: which thread loads which element? That’s what TiledCopy answers.

tiled_copy.cu shows two copy strategies side-by-side:

// From tiled_copy.cu:188
// Thread arrangement: 32x8 = 256 threads
Layout thr_layout = make_layout(make_shape(Int<32>{}, Int<8>{}));

// Each thread loads a 4x1 vector of Elements
Layout val_layout = make_layout(make_shape(Int<4>{}, Int<1>{}));

using CopyOp = UniversalCopy<uint_byte_t<sizeof(Element) * size(val_layout)>>;
using Atom   = Copy_Atom<CopyOp, Element>;

TiledCopy tiled_copy = make_tiled_copy(Atom{}, thr_layout, val_layout);

↪ blog/01_tiled_copy.cu#L196

What this expresses, in one declarative statement:

1. 256 threads participate (the size(thr_layout)).
2. Threads are logically arranged 32 × 8 over the tile.
3. Each thread loads a 4-element vector of floats — one 128-bit load, the widest a single thread can issue.
4. The data tile is 128 × 64 (from block_shape above) — this is (32 * 4) × 8, which tiles perfectly.

Under the hood, make_tiled_copy composes thr_layout and val_layout to figure out which logical (m,n) coordinate each thread’s vector covers. Inside the kernel, all we write is:

// From tiled_copy.cu:107–119
ThrCopy thr_copy = tiled_copy.get_thread_slice(threadIdx.x);
Tensor thr_tile_S = thr_copy.partition_S(tile_S);   // (CopyOp, CopyM, CopyN)
Tensor thr_tile_D = thr_copy.partition_D(tile_D);   // (CopyOp, CopyM, CopyN)
Tensor fragment   = make_fragment_like(thr_tile_D); // register-backed

copy(tiled_copy, thr_tile_S, fragment);
copy(tiled_copy, fragment, thr_tile_D);

↪ blog/01_tiled_copy.cu#L107-L119

That’s the whole kernel. partition_S and partition_D give this thread its slice of the tile. copy issues a vectorized load.

Warp layout inside the (32, 8) thread layout

A warp is 32 consecutive threadIdx.x values — a CUDA hardware concept, fixed in silicon, not configurable. The thr_layout = (32, 8) with default stride (1, 32) gives us a layout function (m, n) → m*1 + n*32, which means threads group into warps like this:

threadIdx.x  →  (m,  n)         which warp
──────────────────────────────────────────
  0..31       (0..31, 0)         warp 0
 32..63       (0..31, 1)         warp 1
 64..95       (0..31, 2)         warp 2
 96..127      (0..31, 3)         warp 3
128..159      (0..31, 4)         warp 4
160..191      (0..31, 5)         warp 5
192..223      (0..31, 6)         warp 6
224..255      (0..31, 7)         warp 7

Each warp spans 32 M-positions at exactly one N-position. That isn’t accidental — it’s the entire reason this layout coalesces. With val_layout = (4, 1), thread 0 of warp 0 reads M=0..3 at N=0, thread 1 reads M=4..7 at N=0, …, thread 31 reads M=124..127 at N=0. In column-major memory (M is the contiguous dim), those 32 lanes hit 32 contiguous 16-byte chunks = 512 bytes contiguous per warp → exactly 4 × 128-byte sectors, zero wasted bandwidth.

The 8 warps together cover a (128, 8) slab of the block in one pass; the TiledCopy walks 8 such passes along N to cover the full (128, 64) block. Each pass is 4 memory sectors per warp × 8 warps = 32 sectors, all useful bytes.

The stride on the M-mode is load-bearing here. If you’d written make_layout((32, 8), (8, 1)) (stride 8 on the M-mode), the layout function becomes (m, n) → m*8 + n, so lane 0 → (0, 0), lane 1 → (0, 1), …, lane 8 → (1, 0). Warp 0 now spans 4 M-positions × 8 N-positions. In column-major memory each lane reads an address block_M = 128 floats away → 32 separate cache lines per warp, fully uncoalesced. Same shape, different stride, 30× worse DRAM bandwidth.

That’s the CuTe discipline in a nutshell: shape says how threads logically partition the block; stride says which threads end up on contiguous bytes. Get the shape wrong and you can’t compile. Get the stride wrong and you compile but burn bandwidth. For column-major memory you want stride-1 on the M-mode; for row-major memory you’d flip to (8, 32) default-stride. For MMA instructions, the acceptable thread layouts are dictated by the hardware and CuTe’s MMA atoms bake them in — one more reason atoms exist.

You can confirm the warp-to-(m, n) mapping at runtime by adding a one-liner inside the kernel:

if (blockIdx.x == 0 && blockIdx.y == 0) {
  printf("tid=%3d warp=%d lane=%2d\n",
         threadIdx.x, threadIdx.x/32, threadIdx.x%32);
}

…and cross-check that lane 0 of every warp hits M=0..3, and warp_id == n in the thr_layout.

Why vectorize

A GPU global memory transaction is 32 or 128 bytes. If 32 threads each issue a 4-byte load to adjacent addresses, that’s one 128-byte coalesced transaction — peak bandwidth. If the same threads each issue a 16-byte ld.global.v4.f32, that’s still one 128-byte transaction but the arithmetic intensity of the kernel (ops / load) quadruples. On bandwidth-bound kernels vectorization is often worth 2–4× in wall time.

The point is not that CuTe makes vectorization possible — it’s that it makes the thread/value layout an explicit, type-level artifact you can read, tune, and reuse.

Testing tips: how does each thread know its piece?

A natural question once you’ve stared at copy(tiled_copy, thr_tile_S, fragment) for a while: where does “thread 65 owns M=4..7 at N=2, 10, 18, …, 58” actually come from? The kernel never computes a thread-dependent index. There is no threadIdx.x * something in user code.

The answer is that the thr_layout is that function. get_thread_slice(threadIdx.x) inverts the layout to recover this thread’s (m, n); partition_S(tile_S) composes that with the tile’s layout to produce a new tensor whose base pointer and strides already encode this thread’s slice. It’s all layout algebra — the compiler folds it away at compile time.

The simplest way to see it is to print the partition from inside the kernel:

if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 65) {
  print("tile_S:     "); print(tile_S);     print("\n");
  print("thr_tile_S: "); print(thr_tile_S); print("\n");
}

For thread 65 in block (0, 0) you’ll see something like:

tile_S:     gmem_ptr[32b](…)       o (_128, _64):(_1, _128)
thr_tile_S: gmem_ptr[32b](…+260)   o (_4, _1, _8):(_1, _0, _1024)

Read thr_tile_S as:

• …+260 — base pointer pre-offset for thread 65 at (m=1, n=2) with val_M=4, giving 1*4 + 2*128 = 260.
• (_4, _1, _8) — shape: 4-float vector, 1 M-step, 8 N-steps.
• (_1, _0, _1024) — strides: contiguous within the vector, no M-repeats (mode is size-1), next N-slab is 1024 floats = 8 columns × 128 rows away.

Every thread’s thr_tile_S has the same shape and same strides — only the base pointer changes. That pointer is the one and only runtime-dependent value; everything else is static type information. That’s why the generated SASS is tight: a few address-add immediates, 8 vectorized loads, 8 vectorized stores. No per-thread index arithmetic at runtime.

Mental model cheat-sheet:

Question	Answer
Which bytes does thread t touch?	Composed from thr_layout, val_layout, and tile_S’s layout.
When is it resolved?	Compile time — except for the base pointer, which depends on threadIdx.x.
What do I iterate?	The rank-3 shape (CopyOp, CopyM, CopyN) — copy() unrolls it.
Do I write index math by hand?	Never. partition_S / partition_D produce the right tensor view.

CuTe’s promise in one line: you define the layouts; CuTe emits the indexing code.

One more debugging tip: print_latex(tiled_copy) dumps a LaTeX figure showing the (thread, value) → (m, n) mapping visually. Render it with any LaTeX engine and you get a nicely colored grid that makes the partition obvious at a glance — especially useful for debugging MMA atoms in later days where the thread layouts aren’t as uniform.

Part 3 — A Real SGEMM: Shared Memory, cp.async, and Software Pipelining

The next step up from tiled copy is the SM80 sgemm in sgemm_sm80.cu. Structurally it’s a standard GEMM:

for each k-tile:
  load A-tile and B-tile from gmem -> smem
  load A/B subtiles from smem -> registers
  MMA into register accumulator
axpby(alpha, acc, beta, C)  // epilogue

But the mechanics differ from a textbook CUDA kernel in several important ways.

Software-Pipelined Mainloop

Naively, each iteration waits for its own load to finish before computing. That wastes the tensor cores during load latency. The SM80 mainloop uses multiple shared memory pipes (bP = Int<3>{}) and Ampere’s cp.async to overlap:

// From sgemm_sm80.cu:143–150 (PREFETCH)
CUTE_UNROLL
for (int k_pipe = 0; k_pipe < K_PIPE_MAX-1; ++k_pipe) {
  copy(copy_a, tAgA(_,_,_,k_tile_next), tAsA(_,_,_,k_pipe));
  copy(copy_b, tBgB(_,_,_,k_tile_next), tBsB(_,_,_,k_pipe));
  cp_async_fence();
  --k_tile_count;
  if (k_tile_count > 0) { ++k_tile_next; }
}

↪ blog/02_sgemm_sm80.cu#L143-L150

cp.async is an async copy instruction: gmem -> smem without staging through a register. The thread issues it, then keeps doing work. cp_async_fence() marks a barrier, and cp_async_wait<N>() waits until only N outstanding fence groups remain.

The core of the mainloop then looks like this, laid out so every iteration is doing three things at once:

// From sgemm_sm80.cu:261–301 (simplified)
while (k_tile_count > -(K_PIPE_MAX-1)) {
  for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
    if (k_block == K_BLOCK_MAX - 1) {
      // Slice next pipe + wait for its cp.async to land in smem
      tXsA_p = tXsA(_,_,_,smem_pipe_read);
      tXsB_p = tXsB(_,_,_,smem_pipe_read);
      cp_async_wait<K_PIPE_MAX-2>();
      __syncthreads();
    }
    auto k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX;
    // smem -> rmem for k_block+1
    copy(s2r_atom_a, tXsA_p(_,_,k_block_next), tXrA(_,_,k_block_next));
    copy(s2r_atom_b, tXsB_p(_,_,k_block_next), tXrB(_,_,k_block_next));
    if (k_block == 0) {
      // gmem -> smem into the *next* pipe
      copy(copy_a, tAgA(_,_,_,k_tile_next), tAsA(_,_,_,smem_pipe_write));
      copy(copy_b, tBgB(_,_,_,k_tile_next), tBsB(_,_,_,smem_pipe_write));
      cp_async_fence();
      --k_tile_count; if (k_tile_count > 0) { ++k_tile_next; }
      smem_pipe_write = smem_pipe_read;
      smem_pipe_read = (smem_pipe_read == K_PIPE_MAX-1) ? 0 : smem_pipe_read+1;
    }
    gemm(mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
  }
}

↪ blog/02_sgemm_sm80.cu#L261-L301

Three overlapping streams of work:

Stream	Instruction	What it does
G→S	cp.async	Prefetch next K-tile of A,B into smem_pipe_write
S→R	ldmatrix (LDSM)	Feed the next k_block of A,B into registers
R→R	mma.sync	Compute on the current k_block

Each pass of the inner loop advances all three by one step. The tensor cores are never waiting on memory.

Swizzled Shared Memory Layouts

Notice this in the TN (row-major × col-major) path:

// From sgemm_sm80.cu:358–363
auto swizzle_atom = composition(Swizzle<3,3,3>{},
                                Layout<Shape <_8,Shape <_8, _8>>,
                                       Stride<_8,Stride<_1,_64>>>{});
auto sA = tile_to_shape(swizzle_atom, make_shape(bM,bK,bP));
auto sB = tile_to_shape(swizzle_atom, make_shape(bN,bK,bP));

↪ blog/02_sgemm_sm80.cu#L358-L363

Swizzle<3,3,3> is a permutation that remaps smem addresses so that 32 threads in a warp hitting 32 consecutive rows never collide on the same smem bank. Without it, ldmatrix (the warp-wide smem → register tensor-core load) incurs 8-way bank conflicts and drops to ~1/8 of peak. The swizzle is purely a layout transform — the data isn’t moved, the layout function is rewired.

Understanding CuTe swizzles is usually the single biggest unlock for writing high-performance shared-memory tile code.

SM75 LDSM: Warp-Cooperative smem → register

The MMA needs data in a specific register layout. Rather than have each thread load scalars one at a time, SM75+ has ldmatrix.sync — a warp-collaborative instruction where 32 threads collectively load 4×8 f16 matrices into registers in one shot. CuTe exposes this as a Copy_Atom:

// From sgemm_sm80.cu:383–389
Copy_Atom<SM75_U32x4_LDSM_N, half_t> s2r_atom_A;
Copy_Atom<SM75_U32x4_LDSM_N, half_t> s2r_atom_B;

↪ blog/02_sgemm_sm80.cu#L383-L389

Again — you declare the atom, CuTe handles the stride algebra and emits the right PTX.

The MMA Atom

The compute atom is the m16n8k16 tensor-core MMA, composed into a 32×32×16 TiledMMA over 2×2×1 = 4 atoms:

// From sgemm_sm80.cu:375–377
TiledMMA mmaC = make_tiled_mma(SM80_16x8x16_F16F16F16F16_TN{},
                               Layout<Shape<_2,_2>>{},   // 2x2x1 MMA Atoms
                               Tile<_32,_32,_16>{});     // 32x32x16 Tiled MMA for LDSM

↪ blog/02_sgemm_sm80.cu#L375-L377

Copy_Atom and MMA_Atom are the two kinds of atoms CuTe knows about. Every CUTLASS kernel is a dance between these two.

Part 4 — Hopper: WGMMA and TMA

On SM90 (H100), two things change the shape of a GEMM kernel fundamentally.

WGMMA — the Warp-Group MMA

Ampere’s mma.sync is a 32-thread (one warp) instruction. Hopper’s wgmma.mma_async.sync is a warp-group instruction — 128 threads (4 warps) collectively issue one MMA, and critically, the MMA executes asynchronously from shared memory. The accumulator is in registers, but the A and B operands stream from smem via hardware-managed descriptors.

In CuTe:

// From wgmma_sm90.cu:332
TiledMMA tiled_mma = make_tiled_mma(
  SM90_64x64x16_F16F16F16_SS<GMMA::Major::MN, GMMA::Major::MN>{});

↪ blog/03_wgmma_sm90.cu#L332

The _SS suffix means “both sources from Shared memory.” The MMA reads A and B directly from smem — no copy(sA, tCrA) step like SM80. The “fragment” on Hopper is just a descriptor pointing into smem:

// From wgmma_sm90.cu:163–164
Tensor tCrA = thr_mma.make_fragment_A(tCsA);  // (MMA,MMA_M,MMA_K,PIPE) — descriptor
Tensor tCrB = thr_mma.make_fragment_B(tCsB);  // (MMA,MMA_N,MMA_K,PIPE) — descriptor

↪ blog/03_wgmma_sm90.cu#L163-L164

And the mainloop compute step uses a three-instruction dance — arrive, gemm, commit, wait — because the WGMMA is asynchronous:

// From wgmma_sm90.cu:267–274
warpgroup_fence_operand(tCrC);
warpgroup_arrive();
cute::gemm(mma, tCrA(_,_,_,k_pipe_read), tCrB(_,_,_,k_pipe_read), tCrC);
warpgroup_commit_batch();
warpgroup_wait<0>();
warpgroup_fence_operand(tCrC);

↪ blog/03_wgmma_sm90.cu#L267-L274

warpgroup_arrive() tells the hardware “I am about to start issuing WGMMAs using this accumulator,” warpgroup_commit_batch() closes the batch, warpgroup_wait<0>() waits for the batch to complete so the accumulator is safe to read.

TMA — Tensor Memory Accelerator

Ampere’s cp.async is per-thread async copy. Hopper introduces TMA, a hardware engine dedicated to moving tiles. A single TMA instruction, issued by a single thread, moves an entire multidimensional tile from gmem → smem (or smem → gmem), handles stride arithmetic, and signals a barrier when done.

wgmma_tma_sm90.cu shows the typical setup. On the host you build a TMA descriptor:

// From wgmma_tma_sm90.cu:301–306
Tensor mA = make_tensor(A, make_shape(M,K), dA);
Tensor mB = make_tensor(B, make_shape(N,K), dB);

// Create TMA Atoms with the desired copy operation
Copy_Atom tmaA = make_tma_atom(SM90_TMA_LOAD{}, mA, sA(_,_,0), make_shape(bM,bK));
Copy_Atom tmaB = make_tma_atom(SM90_TMA_LOAD{}, mB, sB(_,_,0), make_shape(bN,bK));

↪ blog/04_wgmma_tma_sm90.cu#L301-L306

The TMA descriptor is constructed on the host and passed to the kernel as a __grid_constant__ argument. Inside the kernel, the load reduces to one thread per warp-group issuing one instruction per tile:

// From wgmma_tma_sm90.cu:174–179
// Set expected Tx Bytes after each reset / init
ProducerBarType::arrive_and_expect_tx(&producer_mbar[pipe], tma_transaction_bytes);
copy(tma_a.with(producer_mbar[pipe]), tAgA(_,k_tile), tAsA(_,pipe));
copy(tma_b.with(producer_mbar[pipe]), tBgB(_,k_tile), tBsB(_,pipe));

↪ blog/04_wgmma_tma_sm90.cu#L174-L179

This is arrive_and_expect_tx: the thread tells the mbarrier how many bytes will land, then issues the TMA. When the TMA completes, the hardware auto-arrives the barrier with the byte count. Consumers wait on the barrier and are released only after the full transaction arrives.

Cluster Launch

SM90 also introduces thread block clusters — a group of CTAs that co-schedule on the same GPC and can share smem via distributed shared memory (DSMEM). A TMA descriptor can be set up to multicast the same tile to all CTAs in a cluster at once — one DRAM read, one wire, delivered to multiple SMs. Set up with:

dim3 dimCluster(2, 1, 1);
cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smemBytes};
cutlass::launch_kernel_on_cluster(params, kernel_ptr, ...);

Multicast is the primary reason an H100 with large cluster sizes can sustain >80% of its advertised FLOPS on GEMM.

Producer-Consumer Pipelines

With WGMMA async and TMA async, the SM80-style “stream three things at once” mainloop is rephrased as a formal producer-consumer pipeline:

// From wgmma_tma_sm90.cu:217–218
auto write_state = cutlass::PipelineState<K_PIPE_MAX>();  // TMA writes
auto read_state  = cutlass::PipelineState<K_PIPE_MAX>();  // MMA reads

↪ blog/04_wgmma_tma_sm90.cu#L217-L218

Each pipeline stage has two barriers: a producer barrier (signals TMA completion, waited on by MMA) and a consumer barrier (signals MMA completion, waited on by the next TMA writing into that smem buffer). The SM80 mainloop interleaves load and compute in a single warp’s instruction stream; the SM90 mainloop decouples them via barriers, which opens the door to warp specialization (some warps do only loads, some warps do only MMA — see CUTLASS example 48).

Part 5 — Blackwell: tcgen05, TMEM, and 2SM MMA

SM100 (B200/GB200) makes three big architectural changes:

1. A new tcgen05 MMA family, ~2× the throughput of WGMMA per clock.
2. A new memory space called Tensor Memory (TMEM) that holds the MMA accumulator — freeing up registers for other work.
3. 2SM MMA, where two CTAs in a cluster collaborate on a single 256×256 MMA, halving the per-CTA register pressure and amortizing smem loads.

The 5 Blackwell tutorials walk through these one step at a time.

Tutorial 01 — Baseline tcgen05 MMA

On SM100 the MMA reads A from smem (or TMEM), reads B from smem, and writes the accumulator to TMEM, not registers. You have to explicitly allocate a TMEM region, point the accumulator tensor at it, and explicitly copy the result back to registers in the epilogue:

// From 01_mma_sm100.cu:229–241
Tensor tCtAcc = cta_mma.make_fragment_C(tCgC);  // TMEM-backed accumulator

uint32_t elect_one_warp = (threadIdx.x / 32 == 0);
using TmemAllocator = cute::TMEM::Allocator1Sm;
TmemAllocator tmem_allocator{};
if (elect_one_warp) {
  tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns,
                          &shared_storage.tmem_base_ptr);
}
__syncthreads();
tCtAcc.data() = shared_storage.tmem_base_ptr;

↪ blog/05_blackwell_mma_sm100.cu#L229-L241

Executing the MMA looks superficially familiar but with a crucial difference: tcgen05.mma is single-thread-issued. Only one thread per CTA issues each MMA instruction, and CuTe handles that under the elect_one_warp predicate:

// From 01_mma_sm100.cu:286–297
if (elect_one_warp) {
  for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
    gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCtAcc);
    tiled_mma.accumulate_ = UMMA::ScaleOut::One;
  }
  cutlass::arch::umma_arrive(&shared_storage.mma_barrier);
}
cute::wait_barrier(shared_storage.mma_barrier, mma_barrier_phase_bit);
mma_barrier_phase_bit ^= 1;

↪ blog/05_blackwell_mma_sm100.cu#L286-L297

The first MMA uses ScaleOut::Zero to clear the TMEM accumulator; subsequent MMAs flip to ScaleOut::One to add. umma_arrive is the UMMA-specific analog of warpgroup_commit_batch() on Hopper.

In the epilogue, TMEM must be explicitly read into registers using tcgen05.ld:

// From 01_mma_sm100.cu:303–321
TiledCopy tiled_t2r_copy = make_tmem_copy(SM100_TMEM_LOAD_32dp32b1x{}, tCtAcc);
ThrCopy   thr_t2r_copy   = tiled_t2r_copy.get_slice(threadIdx.x);

Tensor tDgC = thr_t2r_copy.partition_D(tCgC);
Tensor tDrC = make_fragment_like(tDgC);
copy(tDgC, tDrC);                                // gmem C -> regs

Tensor tDtAcc = thr_t2r_copy.partition_S(tCtAcc);
Tensor tDgD   = thr_t2r_copy.partition_D(tCgD);
Tensor tDrAcc = make_tensor<AccType>(shape(tDgD));
copy(tiled_t2r_copy, tDtAcc, tDrAcc);            // TMEM -> regs

axpby(alpha, tDrAcc, beta, tDrC);                // fused alpha/beta
copy(tDrC, tDgD);                                // regs -> gmem

↪ blog/05_blackwell_mma_sm100.cu#L303-L321

Tutorial 02 — Add TMA Load

Replace cooperative_copy with SM90_TMA_LOAD:

// From 02_mma_tma_sm100.cu:491–496
Copy_Atom tma_atom_A = make_tma_atom(
  SM90_TMA_LOAD{}, mA, sA_layout, select<0,2>(mma_tiler));

↪ blog/06_blackwell_mma_tma_sm100.cu#L491-L496

And in the mainloop one elected thread issues both TMAs and sets the expected transaction bytes:

// From 02_mma_tma_sm100.cu:309–313
if (elect_one_warp && elect_one_thr) {
  cute::set_barrier_transaction_bytes(shared_storage.tma_barrier, tma_transaction_bytes);
  copy(tma_atom_A.with(shared_storage.tma_barrier), tAgA(_,k_tile), tAsA);
  copy(tma_atom_B.with(shared_storage.tma_barrier), tBgB(_,k_tile), tBsB);
}

↪ blog/06_blackwell_mma_tma_sm100.cu#L309-L313

Tutorial 03 — Multicast TMA Across Cluster

In a cluster of multiple CTAs, each CTA needs A tiles rotated through different M-coords and B tiles rotated through different N-coords. But many CTAs share the same A along the N-axis of the cluster (and B along the M-axis). Hopper TMA multicast broadcasts a single gmem read to all CTAs that need it:

// From 03_mma_tma_multicast_sm100.cu:528–534
Copy_Atom tma_atom_A = make_tma_atom(
    SM90_TMA_LOAD_MULTICAST{},
    mA, sA_layout,
    select<0,2>(mma_tiler),
    size<2>(cluster_layout_vmnk));  // # CTAs in N-mode to multicast to

↪ blog/07_blackwell_mma_tma_multicast_sm100.cu#L528-L534

Inside the kernel each TMA gets a multicast mask — a 16-bit bitmap of destination CTAs:

// From 03_mma_tma_multicast_sm100.cu (mainloop)
uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cluster_layout_vmnk,
                                                         cta_in_cluster_coord_vmnk);
copy(tma_atom_A.with(shared_storage.tma_barrier, tma_mcast_mask_a),
     tAgA(_,k_tile), tAsA);
cute::cluster_sync();  // cluster-wide barrier

↪ blog/07_blackwell_mma_tma_multicast_sm100.cu

One DRAM read. Multiple SMs. Enormous bandwidth amplification — this is where a significant chunk of the Blackwell uplift comes from on K-heavy problems.

Tutorial 04 — 2SM MMA

A single 256×256 MMA is physically larger than one SM can hold. Blackwell solves this with 2SM MMA: two CTAs in a cluster share a TMEM region and each owns half of the accumulator. A leader CTA drives the MMA, a peer CTA contributes its half of the A data:

// From 04_mma_tma_2sm_sm100.cu:450–452
TiledMMA tiled_mma = make_tiled_mma(
    SM100_MMA_F16BF16_2x1SM_SS<TypeA, TypeB, TypeC,
                               256, 256,   // <-- 256x256, was 128x256
                               UMMA::Major::K, UMMA::Major::K>{});

↪ blog/08_blackwell_mma_tma_2sm_sm100.cu#L450-L452

TMA atoms now know about the 2-CTA participation and require the full mma_tiler:

// From 04_mma_tma_2sm_sm100.cu:531–537
Copy_Atom tma_atom_A = make_tma_atom_A_sm100(
    SM100_TMA_2SM_LOAD_MULTICAST{}, mA, sA_layout,
    mma_tiler,            // full MmaTiler_MNK
    tiled_mma,
    cluster_layout_vmnk);

↪ blog/08_blackwell_mma_tma_2sm_sm100.cu#L531-L537

Leader/peer asymmetry shows up in the mainloop — only the leader sets the transaction bytes, and the MMA arrives multicast to 2 CTAs:

// From 04_mma_tma_2sm_sm100.cu (mainloop)
auto elect_one_cta = get<0>(cta_in_cluster_coord_vmnk) == Int<0>{};
int tma_transaction_bytes =
    size<0>(cluster_layout_vmnk) * sizeof(make_tensor_like(tAsA))
  + size<0>(cluster_layout_vmnk) * sizeof(make_tensor_like(tBsB));
if (elect_one_cta) {
  cute::set_barrier_transaction_bytes(shared_storage.tma_barrier, tma_transaction_bytes);
}
cutlass::arch::umma_arrive_multicast_2x1SM(&shared_storage.mma_barrier, mma_mcast_mask_c);

↪ blog/08_blackwell_mma_tma_2sm_sm100.cu

Tutorial 05 — TMA Epilogue

By now C is still being read gmem → register and D still register → gmem. The final tutorial pushes even the epilogue onto the TMA engine. Shared memory is reorganized as a union so the mainloop’s A/B buffers and the epilogue’s C/D buffers share the same smem budget:

// From 05_mma_tma_epi_sm100.cu:132–139
alignas(128) union {
  alignas(128) struct {
    alignas(128) cute::ArrayEngine<TypeA, ...> A;
    alignas(128) cute::ArrayEngine<TypeB, ...> B;
  } mainloop;
  alignas(128) cute::ArrayEngine<TypeC, ...> C;
  alignas(128) cute::ArrayEngine<TypeD, ...> D;
} tensors;

↪ blog/09_blackwell_mma_tma_epi_sm100.cu#L132-L139

C and D use SM90_TMA_LOAD / SM90_TMA_STORE atoms. The epilogue loop is a pipeline unto itself:

// From 05_mma_tma_epi_sm100.cu:423–455 (condensed)
for (int epi_tile_idx = 0; epi_tile_idx < size<2>(tTR_tAcc); ++epi_tile_idx) {
  // TMA: gmem C -> smem C
  if (elect_one_warp && elect_one_thr) {
    cute::set_barrier_transaction_bytes(shared_storage.tma_barrier, tma_transaction_bytes);
    copy(tma_atom_C.with(shared_storage.tma_barrier, 0),
         tGS_gC(_,epi_tile_idx), tGS_sC);
  }
  cute::wait_barrier(shared_storage.tma_barrier, tma_barrier_phase_bit);

  // smem C -> reg + TMEM acc -> reg
  copy_aligned(tTR_sC, tTR_rC);
  copy(t2r_copy, tTR_tAcc(_,_,epi_tile_idx), tTR_rD);

  // alpha*acc + beta*C in registers
  axpby(beta, tTR_rC, alpha, tTR_rD);
  __syncthreads();

  // reg D -> smem D
  copy_aligned(tTR_rD, tTR_sD);
  tma_store_fence();
  __syncthreads();

  // TMA: smem D -> gmem D
  if (elect_one_warp && elect_one_thr) {
    copy(tma_atom_D, tSG_sD, tSG_gD(_,epi_tile_idx));
    tma_store_arrive();
    tma_store_wait<0>();
  }
}

↪ blog/09_blackwell_mma_tma_epi_sm100.cu#L423-L455

This is the complete SM100 GEMM primitive: TMA mainloop load → TMEM MMA → TMA epilogue store. With 2SM + multicast, one CTA pair can sustain near-peak throughput for all three.

Part 6 — Production CUTLASS: A Tour of Optimizations

Everything above is the CuTe-level primitive. Nobody writes a production kernel this way end-to-end — they compose these pieces via the CUTLASS 3.x collective builder API. This section is a guided tour of the cutlass/examples/ directory, explaining what each variant adds.

Collective Builder: the 5-line GEMM

From 49_hopper_gemm_with_collective_builder: most of a production GEMM turns into type plumbing.

// Excerpt from 49_collective_builder.cu
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
    cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
    Shape<_128,_128,_64>, Shape<_1,_1,_1>,
    cutlass::epilogue::collective::EpilogueTileAuto,
    ElementAccumulator, ElementCompute,
    ElementC, LayoutC, AlignmentC,
    ElementD, LayoutD, AlignmentD,
    EpilogueScheduleType,
    cute::conditional_t<UseCustomEVT, CustomEVT, DefaultOperation>
  >::CollectiveOp;

The builder inspects your types (e.g. Sm90, FP16, TileShape<128,128,64>) and picks the right TMA descriptors, swizzled smem layouts, MMA atom, pipeline depth, and warp specialization schedule. For 80% of use cases you don’t touch the internals at all. For the remaining 20% you hand-write a collective mainloop in CuTe (Part 3–5) and plug it into the same kernel skeleton.

Warp Specialization (SM90 + SM100)

48_hopper_warp_specialized_gemm demonstrates the pattern where different warps in a CTA have different jobs: producer warps issue TMAs, consumer warps issue WGMMAs. The handoff is via the PipelineState barriers from Part 4. Under KernelSchedule = KernelTmaWarpSpecializedCooperative, CUTLASS emits a kernel with this structure.

Why do this? Register pressure. If all threads do both TMA issuance and WGMMA consumption, each thread’s register file has to hold the union of both working sets, which hurts occupancy. Warp specialization lets producers have tiny register files (they mostly issue instructions) while consumers get the full budget for accumulators.

Epilogue Visitor Trees (EVT)

The other abstraction worth understanding is EVT — a type-level AST for describing the epilogue as a composition of nodes (Load, Compute, Broadcast, Store). You build the tree, and CUTLASS compiles it into a fused epilogue kernel. From example 49:

// Custom EVT: D = alpha * acc + beta * C
using CustomEVT =
  cutlass::epilogue::fusion::Sm90EVT<
    cutlass::epilogue::fusion::Sm90Compute<cutlass::homogeneous_multiply_add,
      ElementD, ElementCompute, RoundStyle>,
    cutlass::epilogue::fusion::Sm90ScalarBroadcast<ElementScalar>, // beta
    cutlass::epilogue::fusion::Sm90SrcFetch<ElementC>,
    cutlass::epilogue::fusion::Sm90EVT<
      cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies,
        ElementCompute, ElementCompute, RoundStyle>,
      cutlass::epilogue::fusion::Sm90ScalarBroadcast<ElementScalar>, // alpha
      cutlass::epilogue::fusion::Sm90AccFetch
    >
  >;

Read from the bottom up: alpha * acc, then beta * C + (alpha * acc). Each node lowers to a tiny device functor; the tree lowers to one fused pass that streams through the accumulator tile once. You can compose in per-row bias, activation functions, auxiliary output writes, and amax reductions — all without writing a new kernel.

Classic Fused Epilogue: Bias + ReLU

CUTLASS 2-era epilogue fusion (still used for simple cases):

// From 12_gemm_bias_relu
using EpilogueOp = cutlass::epilogue::thread::LinearCombinationRelu<
    ElementOutput,
    128 / cutlass::sizeof_bits<ElementOutput>::value,
    ElementAccumulator, ElementComputeEpilogue,
    cutlass::epilogue::thread::ScaleType::NoBetaScaling>;

Computes d = max(0, α·acc + bias) inline as the accumulator tile is written back. Everything is registers → registers → gmem in one pass, no intermediate smem staging.

FP8 and Blockwise Scaling

54_hopper_fp8_warp_specialized_gemm shows FP8 (E4M3) input with a richer epilogue fusion: per-row bias, ReLU, auxiliary output, and amax reduction for dynamic range tracking. The whole thing is a single EVT:

using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux<
    LayoutAux, cutlass::epilogue::thread::ReLU,
    ElementD, ElementCompute,
    ElementAux, ElementAmax, ElementBias, ElementC>;

67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling is the DeepSeek-style blockwise FP8 GEMM — instead of one scale per tensor, there’s one scale per (block_m, block_k) tile of A and one per (block_n, block_k) tile of B. Blockwise scaling keeps per-block dynamic range tight, which is what makes FP8 training stable.

// From 67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling.cu
using ScaleConfig = decltype(cutlass::detail::sm90_trivial_blockwise_scale_config(TileShape{}));
using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB());
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8Blockwise;

The scales ride along as extra TMA-loaded tensors and get applied during accumulation, not in the epilogue — that’s the whole point.

Mixed-Dtype GEMM (INT4 × FP16, etc.)

55_hopper_mixed_dtype_gemm is the kernel behind W4A16 LLM inference: A is FP16 activations, B is INT4 weights with FP16 scales and zero-points. CUTLASS upcasts B to FP16 in registers after the ldmatrix, then the MMA runs in FP16 as normal:

using CollectiveMainloopConvertOnly = typename cutlass::gemm::collective::CollectiveBuilder<
    ArchTag, OperatorClass,
    ElementB, LayoutB_Transpose, AlignmentB,  // <-- narrow type in B slot
    ElementA, LayoutA_Transpose, AlignmentA,
    ElementAccumulator,
    TileShape, ClusterShape,
    ...

Crucial detail: A and B are swapped in the builder. CUTLASS generally expects the wider type in the A slot, so to support A=FP16, B=INT4 with B’s dequantization happening on the fly, the kernel is really computing B^T * A^T = (A*B)^T, and transposing layouts accordingly. The transpose trick is why you see LayoutB_Transpose above.

Grouped GEMM — MoE’s workhorse

In a Mixture-of-Experts layer, every token is routed to one or a few experts, which means you have a batch of many small GEMMs of different shapes — infeasible to launch as separate kernels. 57_hopper_grouped_gemm packs all of them into a single kernel launch:

using ProblemShape = cutlass::gemm::GroupProblemShape<Shape<int,int,int>>;

struct CooperativeConfig {
  using KernelSchedule   = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum;
  using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative;
  using TileShape        = Shape<_256,_128,_128>;
  using ClusterShape     = Shape<_1,_2,_1>;
};

Each tile pulls its problem index from a scheduler, fetches its per-problem pointers and per-problem TMA descriptors, and runs the normal warp-specialized mainloop. MoE inference at scale is basically grouped GEMM + all-to-all.

StreamK Tile Scheduler

The classical “data-parallel” scheduler assigns one output tile (m, n) to one CTA, which then runs through all of K. When tile counts don’t divide evenly by the SM count, a few “tail” SMs do extra work and the rest idle. StreamK (example 47) breaks this assumption — it decomposes work along K and distributes it across all SMs, with an atomic epilogue reducing partial results.

using DeviceGemmStreamK = cutlass::gemm::device::GemmUniversal<
    ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
    ElementAccumulator,
    OperatorClass, ArchTag,
    ThreadblockShape, WarpShape, InstructionShape,
    EpilogueOp,
    cutlass::gemm::threadblock::ThreadblockSwizzleStreamK,   // <-- only difference
    NumStages, AlignmentA, AlignmentB>;

On problems where M*N doesn’t wave-align to the SM count (long skinny matmuls, odd batch sizes), StreamK can pick up 20-40% by eliminating the tail.

Fused Reductions: GEMM + Softmax, GEMM + LayerNorm

Examples 35 and 37 show how CUTLASS fuses reductions into or around a GEMM:

• 35_gemm_softmax — row-reduce the output of a GEMM to get max and sum, then normalize in a second pass. Basically the building block of attention.
• 37_gemm_layernorm_gemm_fusion — back-to-back GEMM with a LayerNorm baked between them. First GEMM’s epilogue emits row mean/variance; a tiny reduction kernel finalizes; second GEMM’s prologue applies the normalization as it loads.

Example 23 (ampere_gemm_operand_reduction_fusion) goes further: the GEMM output and a reduction of one of the operands along K come out of the same kernel in one pass.

These aren’t just convenience — attention, layernorm, and RMSNorm are memory-bound in the standard implementation, and fusing them into an adjacent GEMM often doubles wall-clock performance.

Blackwell Builder

Nothing about the high-level API changes between Hopper and Blackwell — you swap Sm90 for Sm100 and KernelTmaWarpSpecializedCooperative for its SM100 variant:

// From 70/71_blackwell_gemm*
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
    cutlass::arch::Sm100, OperatorClass,
    ElementA, LayoutA, AlignmentA,
    ElementB, LayoutB, AlignmentB,
    ElementAccumulator,
    MmaTileShape_MNK, ClusterShape_MNK,
    cutlass::gemm::collective::StageCountAutoCarveout<
      static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
    cutlass::gemm::collective::KernelScheduleAuto
  >::CollectiveOp;

All the tcgen05, TMEM, 2SM MMA mechanics from Part 5 are selected automatically by KernelScheduleAuto based on the types you pass in. This is the whole point of the collective builder: the hardest thing on SM100 (getting the pipeline and TMEM allocation right) becomes KernelScheduleAuto.

Putting It Together

Reading backward from what a framework user sees:

• PyTorch calls F.linear(x, w) → dispatches to cuBLAS or a CUTLASS extension.
• CUTLASS 3.x assembles a GemmUniversalAdapter from a CollectiveMainloop + CollectiveEpilogue + TileScheduler. The builders pick the right templates for your architecture and types.
• The CollectiveMainloop is a CuTe kernel that loops over K-tiles, uses TMA/WGMMA/tcgen05 under warp specialization, and produces an accumulator in registers or TMEM.
• The CollectiveEpilogue walks an EVT to apply scaling, bias, activations, and auxiliary reductions while streaming the accumulator back out — TMA epilogue on SM100.
• Under both sits CuTe — layouts, tiled copies, MMA atoms, swizzles — the algebra that makes any of this composable in the first place.

The skill ladder, in order of “can write kernels at this level”:

1. Read a CuTe layout and predict what print_latex would draw.
2. Write a TiledCopy / TiledMMA by hand, reason about bank conflicts via swizzle selection.
3. Write a three-stage software pipeline on SM80 (sgemm_sm80.cu).
4. Write a producer-consumer pipeline with TMA and WGMMA on SM90, and understand why WGMMA is async.
5. Write an SM100 kernel with TMEM allocation, 2SM MMA, and TMA epilogue.
6. Compose a fused epilogue as an EVT — GEMM + bias + activation + amax + aux output in one kernel.
7. Ship a CUTLASS extension module that selects the right schedule per shape via KernelScheduleAuto and handles the weird cases (odd shapes → StreamK; variable-shape batching → Grouped GEMM).

If you made it this far, the practical next step is to actually build and profile tiled_copy.cu, sgemm_sm80.cu, wgmma_tma_sm90.cu, and 05_mma_tma_epi_sm100.cu with Nsight Compute and walk through the memory-pipeline timeline. The gap between reading about TMA and seeing a tile of A show up in smem in one instruction, with the mbarrier lighting up exactly when expected, is where the abstractions stop being words and start being hardware.

🧩 Try the Puzzles

Reading is easy, writing is the test. I packaged the progression from this post — layouts → SM80 pipeline → Hopper WGMMA+TMA → Blackwell tcgen05 → 2SM MMA → TMA epilogue — into a standalone repo of 9 daily puzzles:

github.com/supercharleszhu/cutlass-puzzle

Each day is a self-contained CMake target with a puzzle.cu (critical blocks blanked out with // TODO:), a solution.cu (the full reference), and a README explaining what you’re building and why. The build system finds CUTLASS via a submodule, -DCUTLASS_DIR=, or $CUTLASS_DIR, and gates per-day targets by compute capability — so even without H100/B200 hardware you can still compile most days with -arch=sm_90a/sm_100a and diff your answer against the reference.

Rough pacing:

• Days 1–2 (any CUDA GPU / A100): layout algebra and the three-stream SM80 mainloop. This is the foundation of everything else.
• Days 3–4 (H100): the Hopper async dance — WGMMA + TMA + mbarrier pipelines. If you can finish Day 4 clean, you can read FlashAttention-3.
• Days 5–9 (B200): the full Blackwell progression — TMEM allocation, TMA load, cluster multicast, 2SM MMA, TMA epilogue. Each day adds exactly one primitive.

If you find a puzzle under-specified or over-leading, PRs welcome.

References

• CUTLASS GitHub — tutorials referenced in this post live at examples/cute/tutorial/ and examples/.
• CuTe Documentation — the canonical CuTe guide, worth reading front-to-back.
• Hopper Architecture Whitepaper — background on WGMMA, TMA, clusters.
• Blackwell Architecture Whitepaper — tcgen05, TMEM, 2SM MMA.
• FlashAttention 3 — the definitive worked example of warp-specialized CuTe on Hopper.