TV Layout & MMA Atom — Demystified

Thread 0 owns 4 values in its registers after the MMA executes: Register Matrix Position Why these specific cells? ──────── ───────────────── ───────────────────────── v0 → C[ 0, 0] NVIDIA's hardware wiring v1 → C[ 1, 0] decided this pattern. v2 → C[ 8, 0] You can't change it. v3 → C[ 9, 0] It's in the silicon. Visually in the 16×8 matrix: col 0 row 0: [v0] ← pair 1 row 1: [v1] ← row 2: . ... . (other threads' elements) row 7: . row 8: [v2] ← pair 2 (8 rows below pair 1) row 9: [v3] ← ... . row 15: .

The exact TV_layout_C for mma.sync.m16n8k16 (accumulator D/C): Shape: ( (4, 8), (2, 2) ) Stride: ( (16, 1), (8, 64) ) ─┬──── ─┬──── T V (threads) (values) Both T and V are hierarchical (nested tuples): T = (4, 8) means: thread_id is decomposed as (T_outer, T_inner) T_inner = tid % 8 (which column, 0-7) T_outer = tid / 8 (which row-group, 0-3) V = (2, 2) means: value_index is decomposed as (V_inner, V_outer) V_inner = vid % 2 (which row within a pair) V_outer = vid / 2 (which half: top or bottom)

The 16×8 matrix is stored in row-major as a flat array: offset = row * 8 + col Now trace each stride: T_inner stride = 1 tid % 8 contributes offset += (tid%8) * 1 → this IS the column! (col = tid % 8) T_outer stride = 16 tid / 8 contributes offset += (tid/8) * 16 16 in row-major(8 cols) = 2 rows → row += (tid/8) * 2 V_inner stride = 8 vid % 2 contributes offset += (vid%2) * 8 8 in row-major(8 cols) = 1 row → row += (vid%2) * 1 V_outer stride = 64 vid / 2 contributes offset += (vid/2) * 64 64 in row-major(8 cols) = 8 rows → row += (vid/2) * 8

Formula (derived from shape + stride): offset(tid, vid) = (tid%8)*1 + (tid/8)*16 + (vid%2)*8 + (vid/2)*64 Converting to (row, col) in the 16×8 matrix: col = tid % 8 row = (tid / 8) * 2 + (vid % 2) + (vid / 2) * 8 Verification: (tid=0, vid=0): col=0, row= 0*2 + 0 + 0*8 = 0 → C[ 0, 0] ✓ (tid=0, vid=1): col=0, row= 0*2 + 1 + 0*8 = 1 → C[ 1, 0] ✓ (tid=0, vid=2): col=0, row= 0*2 + 0 + 1*8 = 8 → C[ 8, 0] ✓ (tid=0, vid=3): col=0, row= 0*2 + 1 + 1*8 = 9 → C[ 9, 0] ✓ (tid=5, vid=1): col=5, row= 0*2 + 1 + 0*8 = 1 → C[ 1, 5] ✓ (tid=8, vid=0): col=0, row= 1*2 + 0 + 0*8 = 2 → C[ 2, 0] ✓ (tid=24, vid=2): col=0, row= 3*2 + 0 + 1*8 = 14 → C[14, 0] ✓ (tid=31, vid=3): col=7, row= 3*2 + 1 + 1*8 = 15 → C[15, 7] ✓

If the mapping were simply "stride s_t per thread, stride s_v per value": offset = tid * s_t + vid * s_v Then threads and values would need to be evenly spaced — but they're NOT: Thread 0 → col 0 (offset 0) Thread 1 → col 1 (offset 1) ... Thread 7 → col 7 (offset 7) Thread 8 → row 2 (offset 16) ← NOT offset 8! A single flat stride can't express this jump. The hierarchical shape (4, 8) lets CuTe encode the two different stride patterns: - within a group of 8: stride 1 (columns) - between groups: stride 16 (skip 2 rows)

// After mma.sync, results are in
// registers d0, d1, d2, d3.
// You must KNOW the mapping to store:
int lane = threadIdx.x % 32;
int col = lane % 8;
int row_base = (lane / 8) * 2;

// Store pair 1 (top half)
C[(row_base + 0) * N + col] = d0;
C[(row_base + 1) * N + col] = d1;
// Store pair 2 (bottom half, +8)
C[(row_base + 8) * N + col] = d2;
C[(row_base + 9) * N + col] = d3;

// Change to mma.m16n8k8?
// ALL index math must be rewritten!
// Change to wmma.m16n16k16?
// Completely different layout!

# The atom KNOWS the mapping
atom = MmaF16BF16Op(f16, f32, (16,8,16))

# partition_C uses the atom's TV layout
thr_mma = tiled_mma.get_slice(tidx)
tCgC = thr_mma.partition_C(gC)

# Copy results to global memory
cute.copy(copy_atom, tCrC, tCgC)
# ^^^ correct for ANY atom!

# Change to a different MMA?
# Just swap the atom.
# partition_C and copy auto-adapt.
# Zero index math to update.

┌────────────────────────────────────────────────────────────────────────────────┐ │ │ │ Layer 1: Hardware (NVIDIA ISA) ────────────────────── 🔒 FIXED BY HARDWARE │ │ │ │ "mma.sync.m16n8k16 gives thread 0 elements at (0,0),(1,0),(8,0),(9,0)" │ │ Fixed. Can't change. Different per instruction. │ │ │ │ What's fixed: │ │ • Atom shape (16, 8, 16) in MNK │ │ • Threads per atom (32 = 1 warp) │ │ • Values per thread (4) │ │ • TV_layout strides (16, 1, 8, 64) — the wiring │ │ • Input data format (fp16) and accumulator format (fp32) │ │ │ ├────────────────────────────────────────────────────────────────────────────────┤ │ │ │ Layer 2: Atom (CuTe encoding) ────────────────────── 🔒 FIXED (wrapper) │ │ │ │ TV_layout_C = ((4,8), (2,2)) : ((16,1), (8,64)) │ │ Just a CuTe-readable encoding of the hardware truth. │ │ No user choice here — one-to-one mapping from Layer 1. │ │ │ │ What's fixed: │ │ • Everything from Layer 1, written as CuTe Layout objects │ │ • TV_layout_A, TV_layout_B, TV_layout_C — all defined by ISA │ │ │ ├────────────────────────────────────────────────────────────────────────────────┤ │ │ │ Layer 3: TiledMMA (user-level) ──────────────── 🎨 USER DESIGN CHOICES │ │ │ │ This is where YOUR workload meets the hardware. │ │ You choose how to tile the atom across your CTA tile. │ │ │ │ ┌──────────────────────────────────────────────────────────────────────────┐ │ │ │ Inputs from your problem (not free choices): │ │ │ │ │ │ │ │ • CTA tile size (BM=128, BN=128, BK=32) ← tuning knob, │ │ │ │ but constrained by SMEM size, register budget, occupancy │ │ │ │ • Total threads per block (128) ← must match atoms_layout × 32 │ │ │ │ • Problem size (M, N, K) ← given by the user's matrix dims │ │ │ │ • Data layout of A, B, C in memory ← given by caller │ │ │ └──────────────────────────────────────────────────────────────────────────┘ │ │ │ │ ┌──────────────────────────────────────────────────────────────────────────┐ │ │ │ Your free design parameters: │ │ │ │ │ │ │ │ 1. atoms_layout = (2, 2, 1) │ │ │ │ "How many atoms, arranged how?" │ │ │ │ → 2 in M, 2 in N, 1 in K = 4 atoms = 4 warps = 128 threads │ │ │ │ │ │ │ │ Coverage per step (atoms_layout × atom_shape): │ │ │ │ M: 2 atoms × 16 rows/atom = 32 rows │ │ │ │ N: 2 atoms × 8 cols/atom = 16 cols │ │ │ │ K: 1 atom × 16 depth/atom = 16 reduction │ │ │ │ ─────────────────────────────────────── │ │ │ │ One step covers: 32 × 16 = 512 output elements │ │ │ │ (but need 128×128 = 16384 total, so we iterate) │ │ │ │ │ │ │ │ Tradeoff: more atoms = more parallelism but more threads │ │ │ │ fewer atoms = fewer threads but more work per thread │ │ │ │ │ │ │ │ 2. permutation = (32, 32, 16) │ │ │ │ "How much total coverage per thread across all iterations?" │ │ │ │ → permutation_M = 32, permutation_N = 32, permutation_K = 16 │ │ │ │ │ │ │ │ These must be ≥ atoms_layout × atom_shape: │ │ │ │ M: permutation_M(32) = atoms(2) × atom_M(16) = 32 ✓ exact │ │ │ │ N: permutation_N(32) = atoms(2) × atom_N(8) × 2 = 32 │ │ │ │ ↑ the ×2 means each atom "visits" 2 N-positions │ │ │ │ │ │ │ │ Iterations to cover the full CTA tile: │ │ │ │ M: BM / permutation_M = 128 / 32 = 4 iterations in M │ │ │ │ N: BN / permutation_N = 128 / 32 = 4 iterations in N │ │ │ │ Total: 4 × 4 = 16 iterations to cover full 128×128 │ │ │ │ │ │ │ │ Tradeoff: larger permutation = more elements per thread │ │ │ │ = more registers = potentially lower occupancy │ │ │ │ │ │ │ │ 3. BK (tile depth) = 32 │ │ │ │ "How much K-reduction per SMEM load?" │ │ │ │ → BK / atom_K = 32 / 16 = 2 MMA calls per K-tile │ │ │ │ → Total K iterations: K / BK (outer loop over K) │ │ │ │ │ │ │ │ Tradeoff: larger BK = fewer SMEM loads but more SMEM used │ │ │ └──────────────────────────────────────────────────────────────────────────┘ │ │ │ │ ┌──────────────────────────────────────────────────────────────────────────┐ │ │ │ Derived automatically (you don't compute these): │ │ │ │ │ │ │ │ • Thread→element mapping for all threads (from atom TV + tiling) │ │ │ │ • partition_C/A/B output shapes and strides │ │ │ │ • Grid dimensions = ceil(M/BM) × ceil(N/BN) │ │ │ │ • Number of K-iterations = K / BK │ │ │ │ • Registers per thread = V × M_rest × N_rest (from partition_C) │ │ │ └──────────────────────────────────────────────────────────────────────────┘ │ │ │ └────────────────────────────────────────────────────────────────────────────────┘