An MMA — matrix multiply-accumulate — is the instruction that computes \(D = A \times B + D\). Tensor Cores are the hardware units that run it. The important Blackwell change is not simply a faster MMA. It is a new home for \(D\). Hopper writes every running result into the register file, the private storage of CUDA threads. Blackwell writes it into Tensor Memory, a new 256 KB on-chip SRAM on each streaming multiprocessor (SM), the GPU's compute tile.
Datacenter Blackwell kernels such as FlashAttention-4 use tcgen05.mma, TMEM, and often a 2-CTA MMA. So do matrix multiplies from CUTLASS, NVIDIA's linear-algebra library, on SM100 (Blackwell's architecture number). Their Hopper predecessors use wgmma.mma_async. For the attention algorithm itself, start with How Flash Attention Speeds Up; this post is only the hardware handoff underneath it.
1.The one-line evolution
A, B, and D in registers
A, B from SMEM · D in registers
A, B from SMEM/TMEM · D in TMEM
Ampere, Hopper, and Blackwell are successive NVIDIA GPU generations. SM80, SM90, and SM100 are their architecture numbers. Ampere's mma.sync is a warp-level instruction: 32 threads carry both the inputs \(A,B\) and the result \(D\) in registers.
Hopper stopped keeping most of \(A\) and \(B\) in registers. A Tensor Memory Accelerator (TMA) copies tiles from HBM into shared memory, and WGMMA reads them there. But the accumulator — the running matrix \(D\) — still lands in registers split across 128 threads.
Blackwell moves \(D\) out of registers too. The Tensor Core reads \(A\) and \(B\) from shared memory or TMEM and updates \(D\) in TMEM. CUDA threads no longer have to own the matrix while the multiply is running.
TMA is the copy engine: HBM ↔ shared memory. It exists on Hopper and Blackwell. TMEM is Blackwell's new 256 KB SRAM for MMA results. A B200 has both. The names are easy to mix up; they are different hardware.
2.Read an MMA shape
Before following the memory, read one Hopper MMA. The long name encodes the tile and the math. BF16 is a 16-bit floating-point format. The tail .f32.bf16.bf16 means: multiply BF16 × BF16, add into FP32. .sync.aligned means the 128 threads start together.
wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16
\[ D_{64\times128}\mathrel{+}= A_{64\times16}B_{16\times128} \]The instruction produces — or updates — the entire 64×128 output tile. The 16 is not an output dimension. Each issue is one execution of the instruction. It contributes 16 terms to every dot product in \(D\), then the next issue adds the next 16.
If the inner dimension is 256 — later, that will be the attention head width — the kernel issues the instruction \(256/16=16\) times. Every issue consumes the next 16-wide slices of \(A\) and \(B\) and adds them into the same 64×128 \(D\).
Changing K changes how many times the instruction must run. It does not change the output shape M×N.
3.Hopper's data path
What are we computing? One ordinary matrix multiply, not attention yet. The point is only where \(A\), \(B\), and \(D\) live while a WGMMA runs.
For 32-bit (FP32) accumulation, a 64×128 \(D\) contains 8,192 values:
\[ \frac{64\times128\text{ FP32 values}}{128\text{ threads}} =64\text{ registers per thread}. \]Logically \(D\) is one matrix. Physically each thread owns 64 scattered elements. The complete tile exists only when those 128 slices are considered together.
FP32 \(D\) for an m64nNk16 WGMMA. This count is only the accumulator — not addresses, softmax state, or another pipeline stage.
4.Why registers run out
Now apply the same MMA to attention. FlashAttention-3 (FA3) is the Hopper attention kernel. One example: 64 query rows, 128 key columns, head width 256. FA3 picks these tile sizes per kernel; they are not fixed.
The 128 threads that run MMA and softmax are the consumer warpgroup. Separate producer warps only copy data with TMA.
S is the score matrix \(QK^\top\). P is that same tile after softmax. O is the output \(PV\).
\[ \begin{aligned} S&=QK^\top,\\ P&=\operatorname{softmax}(S),\\ O&\mathrel{+}=PV. \end{aligned} \]Q [64×256] @ Kᵀ [256×128] = S [64×128]
P [64×128] @ V [128×256] = O [64×256]
Each bar is that tile alone against the 255-register cap — they are not all live at once. BF16 packs two values in one 32-bit register.
Hopper does not have to copy P to shared memory
WGMMA has two forms, named by where \(A\) and \(B\) come from. SS (shared-shared) reads both from shared memory. RS (register-shared) reads \(A\) from registers and \(B\) from shared memory. FlashAttention-3 uses the RS form for \(PV\): softmax writes \(P\) into registers, and that same WGMMA reads those registers as input \(A\) while writing \(O\). \(V\) stays in shared memory as input \(B\).
The optimized Hopper path is S in registers → softmax in registers → P stays in registers → RS-WGMMA reads P. Descriptions that always send \(P\) to shared memory are reference implementations, not the optimized FA3 path.
The problem is owning two matrices at the same time. While that RS-WGMMA reads \(P\), its \(O\) result must also live in registers. In this example that is \(32+128=160\) payload registers per thread. FA3's two-stage pipeline also keeps the next score tile, \(S_{\text{next}}\), ready while the current \(PV\) runs — another 64 registers:
\[ P\ (32)+O\ (128)+S_{\text{next}}\ (64)=224 \quad\text{logical payload registers per thread.} \]That leaves at most 31 registers under the 255-register cap in this simplified count. It is a lower bound on how many values must be live together, not the compiled kernel's register count. Real allocation depends on tile sizes, how many consumer warpgroups run, compiler scheduling, and whether every slice is live in the same instruction. Row maxima, sums, pointers, descriptors, and temporaries still need space. The FA3 paper names the extra \(S_{\text{next}}\) buffer as the register cost that limits deeper pipelining.
Hopper works around this rather than simply spilling to slower memory: shrink the query or key tile, overlap only some stages, pick the number of consumer warpgroups, and call setmaxnreg so TMA producer warps give registers to the MMA/softmax consumers. Good kernels avoid those spills, but the register budget still constrains the schedule. The FA3 paper's machine-code study includes a 64×192 \(QK^\top\) WGMMA and a 64×128 \(PV\) WGMMA — not one universal tile.
5.Blackwell's TMEM path
Datacenter Blackwell adds 256 KB of physical SRAM to each SM. It is new silicon, separate from the existing 256 KB register file and the L1/shared-memory pool. NVIDIA calls it Tensor Memory.
Think of TMEM as a 128-row by 512-column spreadsheet of 4-byte cells. Software carves out a rectangle of columns, uses it, and later gives it back. Ordinary shared-memory loads cannot see it.
lanes = rows · software-managed · 32-column allocation units · tcgen05 only
A CTA is a CUDA thread block — the same “thread block” from the shared-memory definition. The scheduler places one CTA on an SM. A CTA pair is two blocks that cooperate; section 7 is the 2-CTA MMA.
Blackwell's replacement for WGMMA is tcgen05.mma. CUTLASS, NVIDIA's matrix-multiply library, also calls it UMMA. In the MMA equation \(D = AB + C\), \(C\) is the running result going in and \(D\) is the result after this instruction. They often occupy the same storage, so people write C/D. Where the tiles live:
| Hopper WGMMA | Blackwell tcgen05.mma | |
|---|---|---|
| A | registers or SMEM | TMEM or SMEM |
| B | SMEM | SMEM |
| C / D, the running result | registers | TMEM |
| Who issues the instruction | 128-thread warpgroup | one thread for a CTA or CTA pair |
| Largest BF16 tile in one instruction | 64×256×16 | 128×256×16 (1-CTA) or 256×256×16 (2-CTA) |
The “largest tile” row is the biggest \(M \times N \times K\) one instruction can cover. People also call that the MMA atom.
The Tensor Core writes \(D\) directly into TMEM. No CUDA thread owns a slice of \(D\) while the MMA runs. Later, a warp uses tcgen05.ld to pull a needed slice into registers for ordinary arithmetic or the final store to HBM.
Ordinary instructions such as ld.shared, ldmatrix, and cp.async cannot address TMEM. On SM100, each CTA sees 128 lanes × 512 columns of 32-bit cells. One warp performs tcgen05.alloc; allocation uses 32-column units and a power-of-two column count. Every allocation must later be released with tcgen05.dealloc.
6.One attention tile, side by side
What is being computed? Both columns run the same \(QK^\top\), softmax, and \(PV\) update. The comparison shows only where the intermediate matrices live. The Hopper column continues the 64×128, head-256 example; FA3 supports other tiles. The Blackwell column follows FlashAttention-4's 128×128 score tile.
Hopper · FA3-style
TMA → SMEM
64 regs / thread
32 regs / thread
128 regs / thread
Blackwell · FA4-style
TMA → SMEM
64 KB in TMEM
128 FP32 + up to 64 packed BF16 regs
tcgen05.mma
stays in TMEM
Blackwell does not remove registers. Softmax is ordinary per-element arithmetic, so FA4 still gives each softmax thread one 128-element score row. The physical tcgen05.ld and tcgen05.st operations are warp-collective: all 32 threads of a warp execute the load or store together, and each warp sees one quarter of its CTA's TMEM. A softmax thread may hold 128 FP32 scores plus up to 64 registers of packed BF16 output, plus temporaries.
The difference is lifetime. FA4 stores the first three quarters of \(P\) to TMEM and launches those MMAs, then handles the last quarter separately. The \(PV\) MMA reads \(P\) from TMEM while \(O\) accumulates in TMEM. The 128×128 \(S\) tile and the \(O\) tile therefore do not sit in every softmax thread's private registers for the whole pipeline.
tcgen05.ld and tcgen05.st still copy softmax rows between TMEM and registers. Softmax needs \(e^x\). On B200 that exponential hardware can become slower than the larger MMA path, which is why FA4 also reschedules softmax and partly emulates \(e^x\).
7.Two CTAs, one cluster
TMEM keeps \(D\) off the register file on one SM. FlashAttention-4 and Blackwell matrix multiplies also group two CTAs so one tcgen05.mma can span two SMs. A CTA is still a whole thread block, not a thread.
The software name for that pair is a thread-block cluster. The hardware neighborhood they must share is a graphics processing cluster (GPC) — several SMs with a fast local interconnect. The GPC name is leftover from graphics chips. On a B200 it is just “these SMs can talk to each other.”
Hopper already had clusters. Two Hopper CTAs in the same GPC can read each other's shared memory — distributed shared memory (DSMEM). Their MMA is still one-SM WGMMA. The new Blackwell instruction is a 2-CTA tensor-core MMA: the hardware consumes both CTAs' shared-memory \(B\) tiles and both SMs' TMEM in one operation.
What is being computed? One 256×128 output tile \(D = AB\). \(A\) has 256 rows. \(B\) has 128 columns. The example teaches why splitting \(B\) is the win, and why splitting \(A\) is only a work partition. It is the 2-CTA mechanism, not a full attention kernel.
Hopper · two independent CTAs
Blackwell · one 2-CTA MMA
\(A\) is half on both sides for the same reason: each CTA owns different output rows, so it only needs those rows of \(A\). That is ordinary tiling.
\(B\) is different. Every output row needs all columns of \(B\). On Hopper, two CTAs computing different row-halves each stage the whole \(B\) tile. Blackwell's 2-CTA MMA lets the pair share \(B\): each CTA stages half, and the tensor cores consume the combined tile. That roughly halves per-CTA shared-memory traffic for \(B\).
| Instruction | Who cooperates | Where |
|---|---|---|
Ampere mma.sync | 1 warp · 32 threads | 1 SM |
Hopper wgmma | 1 warpgroup · 128 threads | 1 SM |
| Blackwell 1-CTA MMA | 1 CTA · cta_group::1 | 1 SM |
| Blackwell 2-CTA MMA | 1 CTA pair · cta_group::2 | 2 SMs · 1 GPC |
CTA 1 does not issue a shared-memory load into CTA 0's half of \(B\) during the MMA. One thread in the leader CTA issues the 2-CTA instruction. The peer CTA must stay resident: its shared memory and TMEM are operands of that instruction. DSMEM — an explicit load from the other CTA's shared memory — is a separate Hopper-era path. FlashAttention-4 still uses DSMEM in the backward pass to exchange half of the softmax-gradient tile.
Measured on one B200
What was timed? The same BF16 matrix multiply \(C = AB\), not attention. CUTLASS's Python dense-GEMM example, with TMA store and FP32 accumulation. One representative tile per mode — not an autotuned sweep. 8 warmup iterations, then the mean of 30 timed launches. Cache left warm.
A 1-CTA MMA cannot use \(M = 256\). The 2-CTA path can. The comparison is that 256×128 tile against the 1-CTA tiles a single SM can issue.
4096×4096×4096 BF16 on one B200. Shorter bar is faster. 2-CTA is 1.20× the 128×128 1-CTA kernel.
- 2-CTA is 1.20× faster than 1-CTA 128×128 on 4096³: 98.5 µs vs 117.7 µs (1396 vs 1168 TFLOP/s).
- Widening N is not a substitute. 1-CTA 128×256, the same output area as 256×128, was slower on this size: 129.0 µs.
- The ratio holds at 8192³, a bit smaller: 888.5 µs vs 1024.9 µs, or 1.15×. Widening the 1-CTA tile to 128×256 closed part of that gap (937.5 µs) and still lost to 2-CTA.
Both 1-CTA and 2-CTA kernels passed the CUTLASS reference check on 256×256×512. This is a kernel microbenchmark, not FlashAttention-4. FA4 uses the same 2-CTA MMA inside a larger pipeline; isolating TMEM, softmax, and 2-CTA from each other needs a controlled kernel comparison.
8.How many threads are on one SM?
Both H100 (SM90) and B200 (SM100) can hold at most 64 resident warps = 2,048 resident threads per SM. Resident means “currently living on that SM,” not “issued this MMA.” Occupancy is how full that capacity is. A kernel can report low occupancy and still be fast if it is using registers or shared memory instead of filling every thread slot.
| Level | Hopper H100 | Blackwell B200 |
|---|---|---|
| Warp | 32 threads | 32 threads |
| Who issues one MMA | 4 warps = 128 threads | 1 thread issues for a CTA / CTA pair |
| Maximum resident | 64 warps = 2,048 threads / SM | 64 warps = 2,048 threads / SM |
| Register file | 65,536 × 32-bit registers / SM | 65,536 × 32-bit registers / SM |
| Per-thread limit | 255 registers | 255 registers |
| New Tensor Memory | none | 256 KB / SM |
A kernel often keeps far fewer than 2,048 threads resident because registers and shared memory run out first. For example, 256 threads at 240 registers each consume 61,440 of the SM's 65,536 registers. There is no room for another block of the same size, even though 1,792 theoretical thread slots remain.
\[ \text{resident threads from registers} \leq \left\lfloor\frac{65{,}536}{\text{registers per thread}}\right\rfloor, \]Warps, thread-block size, shared memory, and how registers are rounded also cap occupancy. That is why “2,048 threads per SM” and “one 128-thread warpgroup issues WGMMA” are both true — and why an optimized attention kernel can report low occupancy without spilling to slower memory.
9.Limits and the final mental model
- TMEM is datacenter-Blackwell-specific. The
tcgen05path here targets SM100/SM100a, such as B200 and GB200. Consumer Blackwell SM120 does not expose this TMEM path. - TMEM is not general scratch space. Ordinary CUDA arithmetic cannot run on it. Softmax, activations, conversions, and other elementwise work still load values into registers first.
- Low occupancy can remain. A Blackwell kernel may still fill shared memory or use many registers for non-MMA stages. TMEM removes accumulator pressure; it does not make other resources infinite.
- 2-CTA MMA is not “Hopper clusters, but faster.” Hopper clusters already enable DSMEM. Blackwell adds a tensor-core instruction whose operands live on two SMs. The CTA group size, 1 or 2, must stay fixed for TMEM and tensor-core ops in that kernel.
- TMEM or 2-CTA alone is not a benchmark result. FA4 combines TMEM with larger MMA tiles, staged \(P\) stores, two-CTA mode, and a redesigned softmax. The GEMM numbers in section 7 isolate 2-CTA for one CUTLASS tile pair. They are not FA4's end-to-end speedup.
Ampere: threads carry the inputs and the result. Hopper: shared memory feeds an asynchronous 128-thread MMA, but threads still carry the result. Clusters exist, but the MMA stays on one SM. Blackwell: one thread dispatches the MMA, the Tensor Core keeps \(D\) in TMEM, and a CTA pair on one GPC can share that MMA so each block stages only half of \(B\).
References
- NVIDIA PTX ISA —
wgmma.mma_async. - NVIDIA PTX ISA —
tcgen05.mma, includingcta_group::1andcta_group::2. - NVIDIA CUDA Programming Guide — thread-block clusters.
- NVIDIA Hopper Tuning Guide — occupancy, registers, and shared-memory limits.
- NVIDIA Blackwell Tuning Guide — SM100 occupancy and memory limits.
- Shah et al., FlashAttention-3 — WGMMA/softmax pipelining and register pressure.
- Shah et al., FlashAttention-4 — TMEM pipeline, 128×128 tiles, staged P transfer, and 2-CTA MMA.
- NVIDIA CUTLASS — Blackwell SM100 GEMMs. The 2-CTA numbers use the CuTe-DSL
dense_gemmexample from CUTLASS 4.6.2. - NVIDIA Technical Blog — Blackwell Ultra SM and TMEM.
How to cite this post
Dong, S. (2026). From WGMMA to TMEM: Hopper-to-Blackwell MMA Evolution.
https://simondong1.github.io/hopper-to-blackwell-mma.html