A visual mental model of how data-center Blackwell tcgen05 kernels move large Tensor Core accumulators out of the normal register file,
then turns TMEM into a pipeline buffer between tcgen05.mma and the epilogue.
Scope: tcgen05-capable data-center Blackwell, not every “Blackwell” GPU
TMEM ≠ TMA
128 lanes/datapaths × 512 columns × 32-bit cells
tcgen05.mma accumulates into TMEM
tcgen05.ld moves TMEM → registers
Double buffering overlaps MMA + epilogue
Hopper mental model
register-centric
GMEMglobal
→ TMA →
SMEMA/B tiles
→
WGMMATensor Core
→
RMEMaccumulator
Large output tiles imply many accumulator registers. Register pressure can constrain tile size,
occupancy, and how much fused epilogue work fits.
Primary pressure: “Can the participating threads hold the accumulator and still leave enough registers for the rest of the kernel?”
Blackwell mental model
TMEM-centric
GMEM
→ TMA →
SMEMA/B tiles
→
tcgen05.mma
→
TMEMaccumulator
→ tcgen05.ld + wait::ld →
RMEMepilogue chunk
For tcgen05 MMA, the destination/accumulator D lives in TMEM. Registers are brought into play when TMEM data is explicitly loaded for post-processing or storage.
New question: “How do I schedule mma → TMEM → ld → epilogue so the kernel stays in steady state?”
1. TMEM is not generic byte-addressed SRAM
Architecturally, the tcgen05 TMEM view is 128 lanes (often called datapaths in CUTLASS) × 512 columns × 32-bit cells. Allocation is by columns;
each allocated column spans all 128 datapaths.
Visual compression: the diagram renders 16 column groups rather than all 512 columns.
A 128-column allocation represents 128 × 128 × 4 B = 64 KiB of logical cell capacity. The two colored regions are only an allocation visual. For example, a 128×256 FP32 accumulator occupies 256 columns (128 KiB), so two full such accumulator stages consume all 512 columns.
2. TMEM changes the steady-state pipeline
GMEMnext operands
→
SMEMstaged A/B
→
MMA
→
TMEMC/D
→
RMEMsmall slice
→
SMEMstore staging
The full D accumulator tile no longer needs to stay live in registers. The epilogue can stream subtiles out of TMEM,
process them, and hand them to the store path.
3. Bigger epilogue chunk ≠ always faster
TMEM → RMEM chunk
larger
Epilogue iterations
fewer
Register footprint
higher
SMEM epilogue footprint
higher
Mainloop SMEM stages
may fall
Optimization becomes a balance between fewer TMEM reads and preserving enough SMEM / registers for a deep mainloop pipeline.
4. Double-buffered TMEM accumulators: a common optimization pattern
When a kernel has independent output tiles/stages, it can reserve separate TMEM accumulator regions so the Tensor Core
produces tile i+1 into one region while the epilogue consumes the already-completed tile i from the other.
The same TMEM region must not be read by the epilogue while MMA is still updating it.
t0
t1
t2
t3
t4
t5
t6
t7
TMA producer role
load operands 0
load operands 1
load operands 2
load operands 3
load operands 4
load operands 5
load operands 6
load operands 7
MMA issue role
tile 0 → buf 0
tile 1 → buf 1
tile 2 → buf 0
tile 3 → buf 1
tile 4 → buf 0
tile 5 → buf 1
tile 6 → buf 0
Epilogue role
consume buf 0
consume buf 1
consume buf 0
consume buf 1
consume buf 0
consume buf 1
Conceptual tile-level timeline only: each “MMA tile” may itself contain a K-loop / multiple tcgen05.mma instructions.
A completion barrier (commonly tcgen05.commit → mbarrier) must make the accumulator region visible to the consumer before
tcgen05.ld; the load itself is asynchronous and must complete before its registers are used.
Likewise, the producer cannot reuse a buffer until the epilogue has released it.
TMA producer role Keeps future operand tiles moving GMEM → SMEM. Exact warp count is kernel-specific.
MMA issue role An elected single thread issues tcgen05.mma; kernels often dedicate a warp/role around that instruction stream.
Epilogue consumer role A warpgroup/warps load their permitted TMEM lane windows, wait for loads, fuse math, and stage stores.
5. Why the whole accumulator should not be loaded at once
Example: output tile = 128 × 256, FP32 accumulator
128 × 256 × 4 B = 131,072 B = 128 KiB
If an implementation tried to materialize the whole 128 KiB tile in registers at once:
→ register pressure comes back
→ epilogue occupancy/flexibility suffers
Typical epilogue strategy:
TMEM full tile
↓ tcgen05.ld
↓ tcgen05.wait::ld
small RMEM subtile
↓
fuse bias / activation / scale / cast
↓
store
↓
next TMEM subtile
6. Layout becomes part of performance
MMA layout
→
TMEM layout
→
RMEM layout
→
Epilogue layout
Some production kernels choose operand/result swizzles and epilogue layouts so the TMEM readout composes more efficiently with the desired register arrangement; this is kernel-specific, not a universal TMEM requirement.
Optimization goal: minimize rearrangement between tcgen05.mma output ownership and the next consumer.
7. Block-scaled MMA: scale factors must be staged in TMEM
SMEMscale blocks
→ tcgen05.cp →
TMEMscale factors
→
block-scaled MMA
For tcgen05 block-scaled MMA (for example MXFP8/MXFP4/NVFP4), SFA/SFB scale factors are staged GMEM → SMEM → TMEM and consumed by MMA from TMEM. In tuned kernels, scale-factor movement can become a major bottleneck.
Scheduling opportunity: where the instruction ordering and independent buffers permit it, pipeline tcgen05.cp(scale[i+1]) with tcgen05.mma(tile[i]). Both use the tcgen05 asynchronous machinery, so correctness requires the documented ordering/barrier protocol.
8. Fused epilogue after TMEM
TMEMFP32 acc
→ tcgen05.ld + wait::ld →
RMEMsubtile
→
scale+bias activation/cast
→
SMEM
→ TMA →
GMEM
Do as much cheap post-processing as possible while the data is already in registers, instead of round-tripping through another kernel.
9. 1-CTA vs 2-CTA UMMA: operation scope, not “one shared TMEM pool”
CTA 0 TMEM
Has its own CTA TMEM view. In cta_group::2, allocation/deallocation are collective across the pair.
one tcgen05.mma issue can operate on current + peer CTA TMEM
Peer CTA TMEM
The peer must be launched and active. The 2-CTA MMA updates TMEM associated with both CTAs.
cta_group::2 is a cooperative two-CTA operation. Do not model it as two CTAs blindly sharing one flat 512-column address space.
PTX requires a consistent cta_group choice for tcgen05 operations in the kernel.
10. What to tune in a real Blackwell kernel
①
TMEM columns per accumulator tile Does the chosen tile leave room for double buffering?
②
TMEM acc 0/1 buffer count Can the epilogue consume a completed region while MMA produces into a different region?
③
TMEM → RMEM subtile size Too small = too many iterations; too large = register pressure.
④
Epilogue SMEM footprint Don’t steal so much SMEM that mainloop staging depth collapses.
⑤
Warp specialization Choose role counts for TMA, MMA issue, compute, and epilogue; the exact partition is workload-specific.
⑥
mbarrier / pipeline timing Async overlap only works if producer/consumer handoff is exact.
Scale-factor feed for block-scaled formats SFA/SFB need SMEM → TMEM staging; pipeline copies with MMA only under correct ordering and separate-buffer constraints.
⑨
1-CTA vs 2-CTA UMMA Choose based on tile shape and cooperative scheduling economics.
⑩
Steady-state utilization Optimize the flat middle of the timeline, not only fill/drain overhead.
One-picture summary
GMEMfuture data
TMA →
SMEMmulti-stage operands
→
tcgen05.mmaasync compute
→
TMEM acc stagesone or more regions
tcgen05.ld →
RMEMsmall epilogue slice
→
SMEMstore staging
TMA →
GMEMoutput
Core idea: On tcgen05-capable data-center Blackwell, TMEM removes the need to keep the full Tensor Core accumulator in the ordinary register file,
then turns that accumulator storage into a schedulable buffer that can overlap with both the next MMA and the current epilogue.
Audit notes — corrected in this revision
✓
Architecture scope TMEM here refers to the tcgen05 data-center Blackwell path; “Blackwell” broadly also includes products that do not expose TMEM.
✓
128 × 512 terminology PTX calls the 128 dimension “lanes”; current CUTLASS utilities often call them “datapaths.” The page now shows both terms.
✓
Double-buffer race removed The old timeline incorrectly overlapped epilogue reads with ongoing MMA writes to the same TMEM region. Buffers now alternate only after completion/release.
✓
Async completion shown tcgen05.ld is asynchronous; the page now calls out its completion wait and cross-role synchronization.
✓
2-CTA model corrected cta_group::2 spans current + peer CTA TMEM; it should not be drawn as a single undifferentiated shared pool.
✓
Block-scale path tightened Scale factors SFA/SFB are explicitly shown as SMEM → TMEM via tcgen05.cp before block-scaled MMA.
This page is a conceptual visualization. Exact supported MMA shapes, TMEM layouts, synchronization rules, and PTX syntax vary by instruction form and architecture target; consult the PTX ISA and CUTLASS docs for implementation details.