Fourth post in the Foundations series. WGMMA has come up twice already without a real explanation: it was the "warp-level MMA" row in CuTe Layouts's hierarchy diagram, and it did the actual matrix multiplies inside FlashAttention-3's warp-specialized kernels in the previous post. This one closes that gap.

From warp to warpgroup

Ampere-era matrix multiplies use mma.sync — one instruction, issued and executed by a single warp (32 threads), synchronously: the warp issues it and waits right there for the result.

Hopper's wgmma.mma_async changes both halves of that. It's issued by a warpgroup — four contiguous warps, 128 threads, with the first warp's rank a multiple of four — executing the instruction collectively. And it's asynchronous: the warpgroup issues it and can go on to do other work immediately; the actual multiply-accumulate happens in the background, and reading the result requires an explicit wait.

Diagram showing four warps of 32 threads combining into one 128-thread warpgroup that issues a single wgmma.mma_async instruction, producing a 64xN output tile held in registers

Where the operands have to live

WGMMA has a specific, asymmetric rule about where its inputs come from, and it's not arbitrary:

  • Operand B must always be in shared memory.
  • Operand A can be in shared memory or registers.
  • The accumulator (C) is always in registers.

This connects directly to what makes the async model work at all. Hopper's TMA (Tensor Memory Accelerator — covered when it came up in What is CUTLASS) loads tiles from global memory into shared memory asynchronously, with no thread involvement. Requiring WGMMA's operands to come from shared memory means a producer warp can be off doing TMA loads for the next tile while a consumer warpgroup is still computing on shared memory data that's already there — the operand placement rule is what lets loading and computing overlap without stepping on each other, which is exactly the "warp specialization" pattern FlashAttention-3 uses.

What one instruction actually computes

A single wgmma.mma_async isn't a full GEMM — it computes one tile, with fixed constraints on the shape: M is always 64, N ranges from 8 to 256 (in multiples of 8), and K is 16 for 16-bit inputs like FP16. A real kernel issues many of these, looping over the K dimension of the full problem and accumulating.

Register fragments: where the output actually lands

The accumulator tile doesn't sit in one place — it's split across all 128 threads' registers, and how it's split isn't "the first 32 elements go to thread 0." According to Colfax Research's WGMMA tutorial, for a 64×64 output tile each thread ends up holding 32 values in a specific repeating pattern — thread 0, for instance, holds the values at coordinates (0,0), (0,1), (8,0), (8,1), and that same four-value group repeats every 8 columns across the tile.

That's not a random layout — it's structured enough that the tutorial describes it by factoring the 32 values into shape (2, 2, 8). If that looks familiar, it should: a Shape like that is exactly a CuTe Layout. The register fragment is a Layout — the same (Shape, Stride) idea from two posts ago, just describing "which register on which thread" instead of "which offset in memory."

The synchronization pattern

In practice, code doesn't hand-write raw wgmma.mma_async PTX — CuTe exposes it through a small set of functions that mirror the issue-then-wait structure directly:

cute::warpgroup_arrive();
cute::gemm(tiled_mma, tCrA(...), tCrB(...), tCrC);
cute::warpgroup_commit_batch();
cute::warpgroup_wait<0>();

warpgroup_arrive() issues a fence, gemm(...) issues the actual wgmma.mma_async instructions, commit_batch() groups the pending async operations together, and warpgroup_wait<0>() blocks until all of them (zero groups still pending) have completed. Only after that call is it safe to read the accumulator.

Why any of this is worth knowing

None of this is exposed by cuBLAS, and it's easy to treat "WGMMA" as just a name that shows up in changelogs. The mechanism is what actually explains why Hopper-era kernels look structurally different from Ampere-era ones: warp specialization (some warps moving data, others computing) only makes sense because WGMMA is asynchronous and warpgroup-wide rather than synchronous and per-warp, and the operand placement rules are exactly what let that overlap happen safely. This is the same idea the Foundations series keeps landing on — each generation's programming model changes because a specific hardware capability changed underneath it, not the other way around.

References

Warpgroup diagram above is original artwork made for this post.