Warpgroup MMA Programming Guide#
Hopper (SM90a) introduced the warpgroup-level MMA PTX instruction family
wgmma.mma_async.sync.aligned. A warpgroup (128 threads / 4 warps)
cooperates on one asynchronous D = A * B + C matrix multiply-accumulate.
Key architectural characteristics:
Warpgroup scope: One MMA is issued collectively by a 128-thread warpgroup rather than by a single warp.
Asynchronous issue model: WGMMA instructions are ordered with
cute.nvgpu.warpgroup.fence(),commit_group(), andwait_group().Descriptor-based operand path: Operand B is sourced from staged shared memory. Operand A can be sourced either from shared memory descriptors or from registers via
OperandSource.Register accumulator: The accumulator lives in RMEM and serves as both the input C and output D of
cute.gemm().Architecture-specific operand layouts: F16/BF16 supports K-major and MN-major dense layouts when A comes from SMEM. FP8 and INT8 variants are K-major only.
The dense DSL op classes currently exposed are MmaF16BF16Op (F16/BF16),
MmaF8Op (FP8 E4M3/E5M2), and MmaI8Op (INT8/UINT8); see
Setting up the TiledMMA, MMA Ops for their full constructor parameters,
instruction K extents, and major-mode constraints.
This guide outlines the CuTe Python DSL programming model for WGMMA kernels: stage operands in SMEM, build fragment descriptors, launch asynchronous warpgroup MMAs, and stage the RMEM accumulator back to GMEM in the epilogue.
Global Memory (GMEM) to MMA data flow overview#
WGMMA instructions require us to stage B input operands in Shared Memory (SMEM), while A input operands can be sourced from either SMEM or registers (RMEM). SMEM operands are read asynchronously by the hardware via SMEM descriptors. The accumulator is always kept in registers (RMEM) of the warpgroup.
The diagram below traces the full data flow of a WGMMA GEMM kernel, for the most
common case where A and B matrices are stored in GMEM and both are staged through
SMEM (a_src=SMEM), and the output matrix –accumulated in RMEM– is written
back to GMEM through an SMEM staging buffer.
There are 3 parallel tracks where each has 2 sub-tracks. Three parallel tracks are for operands A, B, and C/D, respectively. The two sub-tracks are for copying data between different memory spaces and for MMA execution.
Operand A (and symmetrically Operand B):
First, we need to create SMEM tensors for A and B matrices:
sAandsB. These tensors are physically allocated tensors that are the destination of TMA copy and the source operands for the WGMMA instructions.Next the data copy flow creates the tensor views for copying data from GMEM to SMEM. It starts with
mAtensor that represents the matrix A in global memory. ThenmA→local_tile→gAoperation creates the local tile view of A that is the slice of A matrix needed to compute the given CTA’s output tile. Thentma_partition(tma, sA, gA)produces TMA viewstAsA,tAgA, and the loop copies tiles from GMEM into SMEM viacopy(tma, tAgA[k], tAsA[stage]).In parallel, the MMA flow turns the SMEM tensor into an iterable tensor of SMEM descriptors for the WGMMA instructions.
sA(the same shared-memory allocation written by TMA) →partition_A→tCsA(MMA-partitioned SMEM view) →make_fragment_A→tCrA(SMEM descriptor passed tocute.gemm()). Note that the SMEM descriptor is a view created from the SMEM tensor that is interpretable by the WGMMA instructions.
Accumulator C/D:
RMEM accumulator flow (gemm input/output):
partition_C(gC)→tCgC→make_rmem_tensor(tCgC.shape)→acc, which serves as both the accumulator input (C) and output (D) ofcute.gemm()(and the WGMMA instruction).Output flow (RMEM → SMEM → GMEM): After the main loop, the accumulator is type-converted and copied from registers to SMEM via
stmatrix(R2S copy), then stored to global memory via TMA store (S2G copy):mC→local_tile→gC→partition_C→tCgCon the destination side, andtRS_rAcc/tRS_sD/bSG_sD/bSG_gDviews drive the two copy stages.
Operand A Dataflow Path Operand B Dataflow Path Accumulator C/D Dataflow Path
─────────────────────── ─────────────────────── ─────────────────────────────
mA: (M, K) [GMEM] mB: (N, K) [GMEM] ┌──── RMEM ──────────┐
│ │ │ make_rmem_tensor() │
│ local_tile(mA, cta_tiler, coord) │ local_tile(mB, cta_tiler, coord) │ acc: accumulator │
▼ ▼ └───────┬────────────┘
gA: (BM, BK, k) [GMEM] gB: (BN, BK, k) [GMEM] │
│ │ acc:(MMA,MMA_M,MMA_N) [RMEM]
│ ┌──── SMEM ─────────┐ │ ┌──── SMEM ─────────┐ │
│ │ sA = alloc(layout)│ │ │ sB = alloc(layout)│ │ mC: (M, N) [GMEM]
│ └──┬────────┬───────┘ │ └──┬────────┬───────┘ │ │
│ │ │ │ │ │ │ │ local_tile
│ │ thr_mma.partition_A(sA) │ │ thr_mma.partition_B(sB) │ ▼
│ │ ▼ │ │ ▼ │ gC: (BM, BN) [GMEM]
│ │ tCsA:(MMA,MMA_M, │ │ tCsB:(MMA,MMA_N, │ │ partition_C
│ │ MMA_K,PIPE) [SMEM] │ │ MMA_K,PIPE) [SMEM] │ ▼
│ │ │ │ │ │ │ tCgC:(MMA,MMA_M,
│ │ make_fragment_A(tCsA) │ │ make_fragment_B(tCsB) │ MMA_N)
│ │ ▼ │ │ ▼ │ [GMEM] (epi dest)
│ │ tCrA:(MMA,MMA_M, │ │ tCrB:(MMA,MMA_N, │ │
│ │ MMA_K,PIPE) │ │ MMA_K,PIPE) │ │
│ │ [SMEM descriptors] │ │ [SMEM descriptors] │ │
│ │ └─────────────┐ │ │ └─────────────┐ │ │
╰─────┤ │ ╰─────┤ │ │ │
▼ │ ▼ │ │ │
tma_partition(tma, │ tma_partition(tma, │ │ │
sA, gA) │ sB, gB) │ │ │
→ tAsA, tAgA │ → tBsB, tBgB │ │ │
▼ │ ▼ │ │ │
┌───┴────────────────────┐ │ ┌──────┴─────────────────┐│ │ │
│ TMA copy loop (A path):│ │ │ TMA copy loop (B path):││ │ │
│ copy(tma, tAgA[k], │ │ │ copy(tma, tBgB[k], ││ │ │
│ tAsA[stage]) │ │ │ tBsB[stage]) ││ │ │
┌─▶│ (writes into sA; │ │ ┌──▶│ (writes into sB; ││ │ │
│ │ tCrA reads same sA) │ │ │ │ tCrB reads same sB) ││ │ │
│ │ repeat for next k/stage│ │ │ │ repeat for next k/stage││ │ │
│ └────────────────────────┘ │ │ └────────────────────────┘│ │ │
│ │ │ │ │ │ │ │
└────────┘ ▼ └─────────┘ ▼ ▼ │
└───────┬───────────────────────────────┴───────────────────┘ │
│ │
▼ │
┌──────────────────────────────────────────────┐ │
│ GEMM Loop: | │
│ warpgroup.fence() │ │
│ cute.gemm(tiled_mma, │ │
│ acc, D (output, RMEM), │ │
┌──▶ │ tCrA[stage], A (SMEM desc -> sA), │ │
│ │ tCrB[stage], B (SMEM desc -> sB), │ │
│ │ acc) C (accumulator, RMEM) │ │
│ │ warpgroup.commit_group() │ │
│ │ warpgroup.wait_group(n) │ │
│ └──────────────────────────────────────────────┘ │
│ │ │ │
└───────┘ | │
▼ │
Epilogue: │
tRS_rAcc = retile(acc) │
tRS_rD = type_convert(tRS_rAcc) │
│ │
▼ │
R2S: copy(tiled_copy_r2s, tRS_rD, tRS_sD) │
[RMEM → SMEM via stmatrix] │
│ │
▼ │
sC = alloc(epi_layout) [SMEM] │
bSG_sD, bSG_gD = tma_partition(tma_c, sC, gC) ◀───────────────────┘
│
▼
S2G: copy(tma_c, bSG_sD[stage], bSG_gD[coord])
[SMEM → GMEM via TMA store]
Naming convention:
cta_tiler = (BM, BN, BK) = CTA-wide tiler dimensions
mX= a global tensor, e.g., (M, K) for AgX= CTA-tiled GMEM slice, e.g., (BM, BK, k) for AsX= SMEM allocation, e.g., (BM, BK, PIPE) for AtAsA/tBsB= TMA-partitioned SMEM viewstAgA/tBgB= TMA-partitioned GMEM viewstCsX= MMA-partitioned SMEM view, e.g., (MMA, MMA_M, MMA_K, PIPE) for AtCrX= SMEM descriptor fragment, e.g., (MMA, MMA_M, MMA_K, PIPE) for Aacc= RMEM accumulator, (MMA, MMA_M, MMA_N)tCgC= MMA-partitioned GMEM, (MMA, MMA_M, MMA_N)tRS_rAcc/tRS_sD= epilogue retile views for R2S (RMEM → SMEM) copybSG_sD/bSG_gD= TMA-partitioned SMEM/GMEM views for epilogue storeMMA = warpgroup atom thread-value layout; MMA_M/MMA_N/MMA_K = repeat counts (e.g., BM/inst_M), k = outer K-tiles, PIPE = pipeline stages
Setting up the TiledMMA, MMA Ops#
As shown in the data flow overview, CuTe DSL provides many utilities to tile/partition the global memory tensors, and create fragment views of SMEM tensors for MMA instructions.
To utilize these functions, we need to setup the TiledMMA, MMA Ops first.
Creating a WGMMA Op#
A WGMMA op describes the hardware instruction to use, it has parameters like data types, instruction shape, operand A source (SMEM or RMEM), and operand major modes.
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import OperandMajorMode
import cutlass.cute.nvgpu.warpgroup as warpgroup
op = warpgroup.MmaF16BF16Op(
cutlass.Float16, # A/B element type
cutlass.Float32, # accumulator type
(64, 128, 16), # instruction shape (M, N, K)
warpgroup.OperandSource.SMEM, # A operand from shared memory
OperandMajorMode.K, # A is K-major
OperandMajorMode.K, # B is K-major
)
The key parameters are:
Instruction shape
(M, N, K): determines the size of one hardware MMA instruction. WGMMA requiresM = 64and8 <= N <= 256in steps of 8. K is fixed by the op class (16 for F16/BF16, 32 for FP8 and INT8).OperandSource:
SMEMreads A from a shared memory descriptor;RMEMreads A directly from registers.OperandMajorMode:
Kfor K-major (default),MNfor transposed layout. F16/BF16 supports both K-major and MN-major for A and B whena_src=SMEM; whena_src=RMEM, only B can be transposed. FP8 and INT8 are K-major only.
CuTe DSL provides implementation of the following WGMMA ops:
PTX name |
Python class |
Constructor parameters |
|---|---|---|
|
|
|
|
|
|
|
|
|
Creating a Tiled MMA#
A TiledMma tiles the WGMMA atom across the CTA tile. You can pass the op
directly or create an explicit atom first.
# Option 1: directly from op (common shorthand)
tiled_mma = cute.make_tiled_mma(op)
# Option 2: explicit atom creation
atom = cute.make_mma_atom(op)
tiled_mma = cute.make_tiled_mma(atom)
Spatial tiling with a repeat count#
A repeat tuple (M_rep, N_rep, K_rep) replicates the WGMMA atom across the
M, N, and K dimensions, producing a larger tiled MMA that covers a bigger CTA
tile with a single cute.gemm call. Each entry in the repeat tuple
corresponds to one warpgroup (128 threads / 4 warps), so (2, 1, 1)
uses two warpgroups — the standard configuration for large Hopper tiles:
atom = cute.make_mma_atom(op) # op shape: (64, 128, 16)
tiled_mma = cute.make_tiled_mma(
atom,
atom_layout_mnk=(2, 1, 1), # 2 warpgroups in M
)
WGMMA Atom make_tiled_mma(atom, (2, 1, 1))
+---------------+ +----------------+
| | | | ^
| 64 x 128 | | Atom (0,0,0) | |
| x 16 | --(2,1,1)--> | 64 x 128 | | 2 x M_atom
| | repeat | x 16 | | = 128
| | | [Warpgroup 0] | |
+---------------+ +----------------+ |
| | |
| Atom (1,0,0) | |
| 64 x 128 | |
| x 16 | |
| [Warpgroup 1] | v
+----------------+
<-- N_atom = 128 -->
K unchanged = 16
The Hopper dense GEMM examples
(examples/cute/hopper/kernel/dense_gemm/dense_gemm.py) use this pattern.
The helper sm90_utils.make_trivial_tiled_mma(...) selects the repeat count
automatically:
atom_layout_mnk = (2, 1, 1)when bothtile_M > 64andtile_N > 128(two warpgroups reduce register pressure).atom_layout_mnk = (1, 1, 1)otherwise (a single warpgroup suffices).
import cutlass.utils.hopper_helpers as sm90_utils
tiled_mma = sm90_utils.make_trivial_tiled_mma(
a_dtype,
b_dtype,
a_major_mode,
b_major_mode,
acc_dtype,
atom_layout_mnk=(2, 1, 1),
tiler_mn=(64, 128), # atom instruction shape (M, N)
)
Custom tile permutation with permutation_mnk#
make_tiled_mma also accepts an optional permutation_mnk argument that
controls how the tiled atom footprint is laid out across M, N, and K. At a
high level:
atom_layout_mnktells CuTe how many atoms (warpgroups) to replicate.permutation_mnktells CuTe how the final tiled footprint is ordered.
permutation_mnk is a tuple of layouts or integers that represent the
tile size and ordering of values along each dimension. When a mode’s
permutation size is larger than the atom layout’s natural coverage
(atom_layout x inst_shape), each warpgroup receives additional values
to fill the extended region — the warpgroup count stays the same, but each
warpgroup holds more data.
atom = cute.make_mma_atom(op) # op shape: (64, 128, 16)
tiled_mma = cute.make_tiled_mma(
atom,
atom_layout_mnk=(2, 1, 1),
permutation_mnk=(128, 256, 16), # extend N from 128 to 256
)
Without permutation — natural atom coverage (M = 128, N = 128):
C tile (M=128, N=128)
+----------------+
| | ^
| [Warpgroup 0] | |
| 64 x 128 | | 2 x inst_M
| | | = 128
+----------------+ |
| | |
| [Warpgroup 1] | |
| 64 x 128 | |
| | v
+----------------+
<--- N = 128 --->
(each warpgroup owns one (64, 128) atom)
With permutation_mnk = (128, 256, 16) — N extended to 256:
C tile (M=128, N=256)
+----------------+----------------+
| | | ^ N = 128 → 256:
| [Warpgroup 0] | [Warpgroup 0] | | atom pattern repeats
| 64 x 128 | 64 x 128 | | along N. Each warpgroup
| | | | now holds 2x the values
+----------------+----------------+ | along N (same threads,
| | | | more data).
| [Warpgroup 1] | [Warpgroup 1] | |
| 64 x 128 | 64 x 128 | |
| | | v
+----------------+----------------+
<------------ N = 256 ------------>
| atom coverage | value repeat |
Why WGMMA typically does not need permutation_mnk: The WGMMA
instruction already has a large N dimension (64, 128, or 256), so the natural
atom coverage is wide enough that no permutation is needed to align with SMEM
swizzle widths. The Hopper
dense GEMM examples (dense_gemm.py, dense_gemm_persistent.py) use
atom_layout_mnk alone without permutation_mnk.
When permutation_mnk is not provided (default), the tile ordering is
sequential and no permutation is applied.
Partitioning Tensors#
Before computing, partition the CTA-tiled tensors according to the tiled MMA layout. WGMMA partitioning is warpgroup-oriented: each warpgroup (128 threads / 4 warps) receives its own slice of the CTA tile, sized to match the SMEM descriptors and register accumulators that the WGMMA instruction expects.
2-warpgroup example
GEMM (M, N, K) = (512, 768, 256), tile_shape_mnk = (128, 256, 64),
F16 WGMMA atom = (64, 256, 16), atom_layout_mnk = (2, 1, 1),
num_stages = 4, 2 warpgroups = 256 threads.
Global matrices:
mA: (M, K) = (512, 256) mB: (N, K) = (768, 256) mC: (M, N) = (512, 768)
K=256 K=256 N=768
|<--------->| |<--------->| |<----------------->|
+-----------+ +-----------+ +---+---+---+-------+
| | ^ | | ^ | | | | | ^
| mA | | M=512 | mB | | N=768 | | | | | | M=512
| | v | | v | | | | | v
+-----------+ +-----------+ +---+---+---+-------+
Tiling with tile_shape_mnk = (BM, BN, BK) = (128, 256, 64) gives
M/BM = 4 tiles, N/BN = 3 tiles, K/BK = 4 tiles:
mA tiled into (M/BM x K/BK) mB tiled into (N/BN x K/BK) mC tiled into (M/BM x N/BN)
= (4 x 4) blocks = (3 x 4) blocks = (4 x 3) blocks
BK=64 x4 BK=64 x4 BN=256 x3
|<--->| |<--->| |<------>|
+-----+-----+-----+-----+ +-----+-----+-----+-----+ +--------+--------+--------+
| | | | | ^ | | | | | ^ | (0,0) | (0,1) | (0,2) | ^
| | | | | |128 | | | | | |256 | | | | |128
+-----+-----+-----+-----+ v +-----+-----+-----+-----+ v +--------+--------+--------+ v
| | | | | ^ | | | | | ^ | (1,0) | (1,1) | (1,2) | ^
| | | | | |128 | | | | | |256 | | | | |128
+-----+-----+-----+-----+ v +-----+-----+-----+-----+ v +--------+--------+--------+ v
| | | | | | | | | | | (2,0) | (2,1) | (2,2) |
+-----+-----+-----+-----+ +-----+-----+-----+-----+ +--------+--------+--------+
| | | | | | (3,0) | (3,1) | (3,2) |
+-----+-----+-----+-----+ +--------+--------+--------+
Each CTA picks one (M-tile, N-tile) coordinate.
For example, CTA at tile_coord = (1, 0, :).
After local_tile — one CTA’s tile (k = K/BK = 256/64 = 4):
gA: (BM, BK, k) = (128, 64, 4) gB: (BN, BK, k) = (256, 64, 4) gC: (BM, BN) = (128, 256)
BK=64 BK=64 BN=256
|<----->| |<----->| |<--------->|
+-------+-- +-------+-- +-----------+
| |.. | |.. | | ^
BM= | gA | k=4 BN= | gB | k=4 BM= | gC | | 128
128 | | 256 | | 128 | | v
+-------+ +-------+ +-----------+
SMEM tensors sA and sB include a pipeline staging dimension:
sA: (BM, BK, PIPE) = (128, 64, 4) sB: (BN, BK, PIPE) = (256, 64, 4)
get_slice(warp_group_thread_layout(warp_group_idx)) — each
warpgroup receives its slice of the tiled MMA footprint.
With atom_layout_mnk = (2, 1, 1) and inst shape (64, 256, 16),
the tiled MMA covers (2x64, 1x256, 16) = (128, 256, 16) which
exactly matches the CTA tile in M and N. Each warpgroup owns one
64-row slice of M:
sA (one pipeline stage, BM=128, BK=64):
Warpgroup 0's slice Warpgroup 1's slice
inst_K inst_K inst_K inst_K
=16 =16 =16 =16
|<--->|<--->|<--->|<--->| |<--->|<--->|<--->|<--->|
+-----+-----+-----+-----+ ^ +-----+-----+-----+-----+ ^
| 0 | 1 | 2 | 3 | |64 | 0 | 1 | 2 | 3 | |64
+-----+-----+-----+-----+ v +-----+-----+-----+-----+ v
|<-- MMA_K = BK/inst_K = 4 -->| |<-- MMA_K = 4 ---------->|
MMA_M = 64/64 = 1 MMA_M = 64/64 = 1
gC (BM=128, BN=256):
+---------------------------+ ^
| Warpgroup 0: 64 x 256 | | 64
| | |
+---------------------------+ v
| Warpgroup 1: 64 x 256 | ^
| | | 64
+---------------------------+ v
<--------- N = 256 -------->
MMA_M = 64/64 = 1, MMA_N = 256/256 = 1
After partition (per warpgroup):
tCsA: (MMA, MMA_M, MMA_K, PIPE) = (MMA, 1, 4, 4)— MMA_M = BM / (atom_M x inst_M) = 128 / (2x64) = 1, MMA_K = BK / inst_K = 64 / 16 = 4tCsB: (MMA, MMA_N, MMA_K, PIPE) = (MMA, 1, 4, 4)— MMA_N = BN / (atom_N x inst_N) = 256 / (1x256) = 1, MMA_K = 4tCgC: (MMA, MMA_M, MMA_N) = (MMA, 1, 1)— MMA_M = 1, MMA_N = 1
The first mode MMA contains the atom’s thread x value layout — it
encodes which registers within a warpgroup hold which matrix elements.
The remaining modes are repeat counts that tile the atom across the
full CTA tile.
Note
Because the WGMMA instruction shape is large (64 x {64..256}), the tiled MMA footprint typically covers the entire CTA tile in M and N with just one or two warpgroups. This means MMA_M and MMA_N are often 1. The MMA_K dimension is where the repeat count is non-trivial (BK / inst_K iterations per pipeline stage).
1-warpgroup example (contrast)
For a smaller tile (128, 128, 64) with atom_layout_mnk = (1, 1, 1),
inst shape (64, 128, 16), and num_stages = 4,
the tiled MMA covers only (64, 128, 16).
Now a single warpgroup must iterate over two atom-blocks along M:
tCsA: (MMA, MMA_M, MMA_K, PIPE) = (MMA, 2, 4, 4)— MMA_M = 128 / (1x64) = 2tCsB: (MMA, MMA_N, MMA_K, PIPE) = (MMA, 1, 4, 4)— MMA_N = 128 / (1x128) = 1tCgC: (MMA, MMA_M, MMA_N) = (MMA, 2, 1)
# Based on examples/cute/hopper/kernel/dense_gemm/dense_gemm.py
@cute.kernel
def kernel(tiled_mma: cute.TiledMma, ...):
tidx, _, _ = cute.arch.thread_idx()
# CTA-tiled global tensors
gA_mkl = cute.local_tile(
mA_mkl, tile_shape_mnk, tile_coord_mnkl, proj=(1, None, 1)
)
gB_nkl = cute.local_tile(
mB_nkl, tile_shape_mnk, tile_coord_mnkl, proj=(None, 1, 1)
)
gC_mnl = cute.local_tile(
mC_mnl, tile_shape_mnk, tile_coord_mnkl, proj=(1, 1, None)
)
# Warpgroup-oriented slicing (128 threads per warpgroup)
warp_group_idx = cute.arch.make_warp_uniform(
tidx // num_threads_per_warp_group # 128
)
warp_group_thread_layout = cute.make_layout(
mma_warp_groups, # e.g. 2
stride=num_threads_per_warp_group, # 128
)
thr_mma = tiled_mma.get_slice(
warp_group_thread_layout(warp_group_idx)
)
# Partition C from global
tCgC = thr_mma.partition_C(gC_mnl) # (MMA, MMA_M, MMA_N)
# Partition A/B from staged SMEM
tCsA = thr_mma.partition_A(sA) # (MMA, MMA_M, MMA_K, PIPE)
tCsB = thr_mma.partition_B(sB) # (MMA, MMA_N, MMA_K, PIPE)
Pre and Post-Conditions for Partitioning#
The inputs of
partition_A,partition_B, andpartition_Cshould be at least rank-2 tensors.The output layout is constrained by the selected MMA atom:
For A, the output has layout
(MMA, MMA_M, MMA_K, ...).For B, the output has layout
(MMA, MMA_N, MMA_K, ...).For C, the output has layout
(MMA, MMA_M, MMA_N, ...).
Partitioning reasons about layout, not memory space or element type. When
a_src=OperandSource.RMEM, the same tiled MMA shape still determines the logical A footprint, but A is materialized as a register fragment rather than a shared-memory descriptor.
Making Fragments#
Fragments are the tensors that the WGMMA instruction operates on. For dense WGMMA:
Fragment A: an SMEM descriptor when
a_src=OperandSource.SMEM, or an RMEM register fragment whena_src=OperandSource.RMEM.Fragment B: an SMEM descriptor pointing into staged shared memory buffers.
Fragment C (accumulator): an RMEM tensor that serves as both the input C and output D of
cute.gemm().
WGMMA fragments for A and B are SMEM descriptors — the hardware reads directly from shared memory. There is no explicit SMEM → RMEM copy step for operands A and B. The accumulator, however, still lives in per-thread registers (RMEM).
Creating fragment descriptors and accumulator fragments#
Fragment creation has two parts:
1. A and B fragment descriptors
make_fragment_A and make_fragment_B take the MMA-partitioned SMEM
views (tCsA / tCsB) and produce descriptor tensors that the WGMMA
instruction consumes. Each descriptor points to one tile within a pipeline
stage in shared memory.
# MMA-partitioned SMEM views (see "Partitioning Tensors")
tCsA = thr_mma.partition_A(sA) # (MMA, MMA_M, MMA_K, PIPE)
tCsB = thr_mma.partition_B(sB) # (MMA, MMA_N, MMA_K, PIPE)
# SMEM descriptor fragments consumed by cute.gemm()
tCrA = tiled_mma.make_fragment_A(tCsA) # (MMA, MMA_M, MMA_K, PIPE)
tCrB = tiled_mma.make_fragment_B(tCsB) # (MMA, MMA_N, MMA_K, PIPE)
Continuing the 2-warpgroup example from Partitioning Tensors
(F16 atom = (64, 256, 16), tile_shape_mnk = (128, 256, 64),
atom_layout_mnk = (2, 1, 1), num_stages = 4):
tCsA: (MMA, MMA_M=1, MMA_K=4, PIPE=4)
tCsB: (MMA, MMA_N=1, MMA_K=4, PIPE=4)
make_fragment_A(tCsA) -> tCrA: (MMA, 1, 4, 4)
make_fragment_B(tCsB) -> tCrB: (MMA, 1, 4, 4)
Each element of tCrA/tCrB is an SMEM descriptor — one per
(MMA_K, PIPE) pair. The hardware reads SMEM directly via the
descriptor; no explicit SMEM -> RMEM load is needed.
tCrA per warpgroup (4 pipeline stages, 4 K-blocks each):
|<-- MMA_K = BK/inst_K = 4 -->|
stage 0: +------+------+------+------+
| k=0 | k=1 | k=2 | k=3 | inst_M=64 (MMA_M=1)
+------+------+------+------+
stage 1: +------+------+------+------+
| k=0 | k=1 | k=2 | k=3 | inst_M=64
+------+------+------+------+
stage 2: +------+------+------+------+
| k=0 | k=1 | k=2 | k=3 | inst_M=64
+------+------+------+------+
stage 3: +------+------+------+------+
| k=0 | k=1 | k=2 | k=3 | inst_M=64
+------+------+------+------+
Similarly for tCrB with shape (MMA, MMA_N=1, MMA_K=4, PIPE=4).
Note
WGMMA fragments for A and B are SMEM descriptors — the hardware
reads SMEM directly, so there is no ldmatrix retiling step required
before cute.gemm().
When A comes from registers (``OperandSource.RMEM``)
In fused kernels, the output of one MMA can become the A operand of the
next. The second TiledMma is created with
a_src=OperandSource.RMEM, and make_fragment_A is not used.
Instead:
The accumulator’s C layout
(MMA, MMA_M, MMA_N)is converted to the A layout(MMA, MMA_M, MMA_K)expected by the secondTiledMma.The accumulator values are type-converted and stored into an RMEM tensor with the A layout.
The resulting RMEM tensor is passed directly to
cute.gemm()as the A operand — no SMEM descriptor is involved.
See the Hopper FMHA example (examples/cute/hopper/kernel/attention/fmha.py) for the complete pattern.
2. C fragment (accumulator)
The accumulator lives in per-thread registers (RMEM). Its shape is derived
from the partitioned C layout. The accumulator starts at zero before the K
loop and is updated in-place by each cute.gemm() call.
# Partition C from global (see "Partitioning Tensors")
tCgC = thr_mma.partition_C(gC_mnl) # (MMA, MMA_M, MMA_N)
# Allocate RMEM accumulator with the same shape
acc_shape = tCgC.shape
acc = cute.make_rmem_tensor(acc_shape, cutlass.Float32)
acc.fill(0.0)
For the same running example:
tCgC: (MMA, MMA_M=1, MMA_N=1)
make_rmem_tensor(tCgC.shape, Float32) -> acc: (MMA, 1, 1)
The accumulator stays in RMEM for the entire main loop.
cute.gemm() reads A/B from SMEM descriptors and accumulates into acc.
+-----------------------------------+
| acc: (MMA, 1, 1) in RMEM |
| 64 x 256 elements per warpgroup |
| Float32 |
+-----------------------------------+
Creating SMEM layouts for A and B#
The SMEM layouts define how A and B tiles are staged in shared memory,
including swizzling for bank-conflict-free descriptor access. The helper
functions in cutlass.utils.hopper_helpers handle the details.
Host side (@cute.jit):
import cutlass.utils.hopper_helpers as sm90_utils
# Create SMEM layouts (includes swizzle + staging)
a_smem_layout = sm90_utils.make_smem_layout_a(
a_layout, # LayoutEnum — row-major or col-major
tile_shape_mnk, # CTA tile (M, N, K)
a_dtype, # element type (e.g. Float16)
num_stages, # pipeline depth
)
b_smem_layout = sm90_utils.make_smem_layout_b(
b_layout,
tile_shape_mnk,
b_dtype,
num_stages,
)
epi_smem_layout = sm90_utils.make_smem_layout_epi(
c_dtype,
c_layout,
epi_tile,
epi_stage,
)
make_smem_layout_a and make_smem_layout_b are convenience helpers that
build a complete, staged SMEM layout in four steps:
Extract the operand tile shape. For A the
(M, K)portion oftile_shape_mnkis kept viacute.slice_; for B the(N, K)portion.Determine the major mode. The major mode (K-major or MN-major) is read from the layout enum (
a_layout.is_k_major_a()). The major-mode dimension size is used for swizzle selection.Select and materialise the swizzle atom. A heuristic (
get_smem_layout_atom) picks the widest swizzle whose contiguous size (in bits) evenly divides the major-mode dimension:Swizzle
Contiguous bits
SW128
1024 (128 B)
SW64
512 (64 B)
SW32
256 (32 B)
Interleave
128 (16 B)
make_smem_layout_atomthen combines the chosen swizzle with a compact outer layout into aComposedLayout(swizzle, outer).Tile to the operand shape and append the staging dimension.
cute.tile_to_shapebroadcasts the atom to the full(M_or_N, K)shape withnum_stagesappended. Theorderargument controls which dimension is contiguous:(0, 1, 2)for K-major (K innermost),(1, 0, 2)for MN-major (MN innermost).
For the running F16 example (tile_shape_mnk = (128, 256, 64),
num_stages = 4, K-major A, K-major B):
A operand (K-major, tile = (M=128, K=64)):
major_mode_size = 64
64 * 16 bits = 1024 bits → SW128
atom = make_smem_layout_atom(K_SW128, Float16)
tile_to_shape(atom, (128, 64, 4), order=(0,1,2))
-> a_smem_layout: ComposedLayout with shape (128, 64, 4)
B operand (K-major, tile = (N=256, K=64)):
major_mode_size = 64
64 * 16 bits = 1024 bits → SW128
atom = make_smem_layout_atom(K_SW128, Float16)
tile_to_shape(atom, (256, 64, 4), order=(0,1,2))
-> b_smem_layout: ComposedLayout with shape (256, 64, 4)
Kernel side (@cute.kernel):
The layout and swizzle are passed to shared-memory allocation. The result
is a ComposedLayout whose .outer is the logical layout and .inner
is the swizzle:
# Based on examples/cute/hopper/kernel/dense_gemm/dense_gemm.py
sA = storage.sA.get_tensor(
a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
)
sB = storage.sB.get_tensor(
b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
)
After allocation:
sAhas shape(BM, BK, PIPE) = (128, 64, 4).sBhas shape(BN, BK, PIPE) = (256, 64, 4).
These are the staged SMEM tensors consumed by partition_A /
partition_B and make_fragment_A / make_fragment_B
(see Making Fragments).
Note
If you need finer control, you can build layout atoms directly with
cute.nvgpu.warpgroup.make_smem_layout_atom(...) and compose the final
SMEM layout manually via cute.tile_to_shape.
Executing the GEMM (Main Loop)#
The main loop iterates over K-tiles. The WGMMA-specific part of each iteration is the fence / gemm / commit / wait sequence:
acc.fill(0.0)
tiled_mma.set(cute.nvgpu.warpgroup.Field.ACCUMULATE, True)
for k_tile in cutlass.range(k_pipe_mmas, k_tile_cnt, 1, unroll=1):
# ... wait for TMA load (pipeline details in dense_gemm.py) ...
cute.nvgpu.warpgroup.fence()
tile_crd = (None, None, None, consumer_read.index)
cute.gemm(tiled_mma, acc, tCrA[tile_crd], tCrB[tile_crd], acc)
cute.nvgpu.warpgroup.commit_group()
cute.nvgpu.warpgroup.wait_group(k_pipe_mmas)
# ... release buffer & advance pipeline (see dense_gemm.py) ...
cute.nvgpu.warpgroup.wait_group(0)
Key points:
fence()orders prior SMEM writes before WGMMA issue.commit_group()publishes queued WGMMA instructions as a group.wait_group(n)waits until at mostngroups remain in flight.wait_group(0)after the loop drains all work before the epilogue.Field.ACCUMULATE—Trueaccumulates (D += A*B),Falseoverwrites (D = A*B). The dense GEMM setsTrueand zero-fillsaccso the first iteration computes0 + A*B.
Complete Workflow#
Putting it all together, a typical Hopper WGMMA GEMM has this structure.
The MMA-relevant steps are highlighted; see dense_gemm.py for the full
kernel including TMA, pipeline, and epilogue details.
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import OperandMajorMode
import cutlass.cute.nvgpu.warpgroup as warpgroup
import cutlass.utils.hopper_helpers as sm90_utils
# --- Host side (@cute.jit) ---
# 1. MMA op + tiled MMA
op = warpgroup.MmaF16BF16Op(
cutlass.Float16, cutlass.Float32, (64, 128, 16),
warpgroup.OperandSource.SMEM, OperandMajorMode.K, OperandMajorMode.K,
)
tiled_mma = cute.make_tiled_mma(op)
# 2. SMEM layouts
a_smem_layout = sm90_utils.make_smem_layout_a(a_layout, tile_shape_mnk, a_dtype, num_stages)
b_smem_layout = sm90_utils.make_smem_layout_b(b_layout, tile_shape_mnk, b_dtype, num_stages)
# 3. TMA copy atoms + kernel launch (see dense_gemm.py)
# --- Kernel side (@cute.kernel) ---
# 4. Allocate SMEM
smem = cutlass.utils.SmemAllocator()
storage = smem.allocate(SharedStorage)
sA = storage.sA.get_tensor(
a_smem_layout.outer, swizzle=a_smem_layout.inner) # (BM, BK, PIPE)
sB = storage.sB.get_tensor(
b_smem_layout.outer, swizzle=b_smem_layout.inner) # (BN, BK, PIPE)
# 5. CTA-tiled global tensors
gA_mkl = cute.local_tile(mA_mkl, tile_shape_mnk, tile_coord, proj=(1, None, 1))
gB_nkl = cute.local_tile(mB_nkl, tile_shape_mnk, tile_coord, proj=(None, 1, 1))
gC_mnl = cute.local_tile(mC_mnl, tile_shape_mnk, tile_coord, proj=(1, 1, None))
# 6. Warpgroup slice, partition & make fragments
warp_group_idx = cute.arch.make_warp_uniform(tidx // num_threads_per_warp_group)
warp_group_thread_layout = cute.make_layout(mma_warp_groups, stride=num_threads_per_warp_group)
thr_mma = tiled_mma.get_slice(warp_group_thread_layout(warp_group_idx))
tCsA = thr_mma.partition_A(sA) # (MMA, MMA_M, MMA_K, PIPE)
tCsB = thr_mma.partition_B(sB) # (MMA, MMA_N, MMA_K, PIPE)
tCrA = tiled_mma.make_fragment_A(tCsA) # SMEM descriptor
tCrB = tiled_mma.make_fragment_B(tCsB) # SMEM descriptor
tCgC = thr_mma.partition_C(gC_mnl) # (MMA, MMA_M, MMA_N)
acc = cute.make_rmem_tensor(tCgC.shape, acc_dtype)
# 7. TMA pipeline setup + prefetch (see dense_gemm.py)
# 8. Main loop — fence / gemm / commit / wait
acc.fill(0.0)
tiled_mma.set(cute.nvgpu.warpgroup.Field.ACCUMULATE, True)
for k_tile in cutlass.range(k_pipe_mmas, k_tile_cnt, 1, unroll=1):
# ... wait for TMA load ...
cute.nvgpu.warpgroup.fence()
tile_crd = (None, None, None, consumer_read.index)
cute.gemm(tiled_mma, acc, tCrA[tile_crd], tCrB[tile_crd], acc)
cute.nvgpu.warpgroup.commit_group()
cute.nvgpu.warpgroup.wait_group(k_pipe_mmas)
# ... release buffer, advance pipeline ...
cute.nvgpu.warpgroup.wait_group(0)
# 9. Epilogue: RMEM → SMEM (stmatrix) → GMEM (TMA store)
# ... (see dense_gemm.py)
See also:
Dense GEMM example:
examples/cute/hopper/kernel/dense_gemm/dense_gemm.pyPersistent GEMM example:
examples/cute/hopper/kernel/dense_gemm/dense_gemm_persistent.pyFMHA example (RMEM A path):
examples/cute/hopper/kernel/attention/fmha.pyHelper utilities:
cutlass.utils.hopper_helpers