Tensor Cores, layouts, and swizzling

A tensor core is a fixed-function matrix-multiply unit. You hand it a small tile of A, a small tile of B and an accumulator, and it produces D = A·B + C in a few cycles. That is the whole idea, and it is worth roughly 15× the FP32 throughput of the same SM.

The complexity is entirely in the interface. The unit does not read "a matrix" — it reads specific bits out of specific registers of specific lanes, or specific bytes at specific shared-memory offsets. Almost every difficult thing in a modern GEMM or attention kernel is a consequence of getting data into that exact shape without wasting bandwidth on the way.

The instruction has been redesigned in every architecture, and the direction of travel is consistent: the operands move further away from the threads, and the issuing scope gets wider. That is not aesthetics. Each generation roughly doubled tensor-core throughput while the register file stayed at 256 KB per SM, so the register file stopped being able to hold or feed the operands.

Start with the one you can hold in your head: mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32. One warp computes a 16×8 output from a 16×16 A and a 16×8 B. Every one of those 128 output elements, and all 384 input elements, lives in a named register of a named lane.

Two derived indices do all the work:

Write g and t for those. Then, straight from PTX §9.7.15.5.8:

Reading a table of index formulas is not the same as seeing it. Pick a lane:

Two structural facts fall out of the picture, and both matter downstream:

• Four consecutive lanes own one row of A, two elements each. Lanes 4g…4g+3 cover columns 0–7 of row g. Four lanes × 2 bf16 elements = 16 bytes. Hold that number.
• The accumulator is already in the right shape to be an A operand for the next matmul — almost. C/D uses the same (g, 2t) pattern as A's first register. This is why FlashAttention can feed P = softmax(S) straight into the second GEMM without a shared-memory round trip.

// One warp: D[16x8] = A[16x16] * B[16x8] + C[16x8], bf16 in, fp32 accumulate.
uint32_t a[4], b[2];      // each holds 2 packed bf16
float    c[4] = {0,0,0,0};

asm volatile(
  "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
  "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
  : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
  : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));

Now the problem. Your tile is sitting in shared memory in some sane row-major order. The fragment layout above is not row-major, not column-major, and not any stride pattern a normal load naturally produces — lane 5 needs elements (1,2), (1,3), (9,2), (9,3), (1,10), (1,11), (9,10), (9,11). Ordinary shared-memory loads require multiple strided scalar or vector loads per lane plus a pile of index arithmetic; the exact instruction count depends on their width.

ldmatrix is a warp-collective gather that produces exactly that layout in one instruction:

The calling convention is unusual and worth stating precisely. Each lane supplies an address, and that address is the start of a matrix row, not the lane's own data. For .x4, lanes 0–7 give the eight row addresses of the first 8×8 matrix, lanes 8–15 the second, and so on. The hardware then redistributes: lane t receives row t/4, columns 2(t\,\%\,4) and 2(t\,\%\,4)+1 — identical to the mma.sync A-fragment mapping. That is the entire point of the instruction.

Hopper adds stmatrix (sm_90), the inverse, used in epilogues to push accumulators back into shared memory in a TMA-friendly layout. The matching transpose, movmatrix — an 8×8 register-to-register transposition across the warp, handy in attention backward passes where you need S^T without a shared-memory round trip — is older than people assume: it has existed since Turing (sm_75+, PTX §9.7.15.5.17).

Here is the collision, stated as plainly as possible.

An ldmatrix.m8n8.b16 operand is 8 rows × 16 bytes. Store your tile row-major with a K dimension of 64 bf16 values, and each row is 128 bytes, so the eight addresses the warp supplies are 128 bytes apart. Bank = (byte / 4) mod 32, and 128 bytes is exactly 32 banks — so all eight addresses land on the same four banks. An 8-way conflict, on every single operand load, for the rest of your kernel's life.

Padding can rotate the bank mapping, but it does not inherently destroy alignment: adding 16 bytes makes the stride 144 bytes, still a multiple of 16. Smaller arbitrary padding can violate the alignment rules, and any padding costs footprint and complicates the descriptor/layout contract. Production kernels therefore usually permute the layout with an XOR that preserves every 16-byte chunk and matches the layouts TMA and WGMMA describe. From Part 2:

Every mode that TMA and wgmma accept uses M = 4 and S = 3: the low 16 bytes are never permuted and the pattern repeats every 8 rows. So the 128-byte mode reduces to

and the eight addresses of an ldmatrix operand land in eight different bank groups. The figure below is the same explorer as Part 2, repeated here because this is where it earns its keep:

Three consequences worth memorising:

• The 8×16 B no-swizzle atom is exactly one ldmatrix.m8n8 operand. The swizzled atoms are the same 8 rows widened to 32, 64 or 128 bytes.
• Prefer 128 B swizzle when the tile permits it. Its atom covers one 128-byte contiguous region on the global side and a full conflict-free warp access on the shared side.
• But 128 B swizzle caps the innermost TMA box at 128 bytes — 64 elements in bf16. A BLOCK_K = 128 bf16 K-major tile needs two TMA boxes, or a narrower swizzle. Violating this rule is one common cause of CUDA_ERROR_INVALID_VALUE from cuTensorMapEncodeTiled.

wgmma.mma_async changes three things at once. It is issued by 128 threads (four contiguous warps whose first warp-rank is a multiple of 4). It is asynchronous — it returns immediately and you wait on a group. And its B operand, and optionally its A operand, are read directly from shared memory by the tensor core, with no ldmatrix and no registers involved.

What you pass instead of a register is a 64-bit matrix descriptor: a packed structure holding the shared-memory base address, two strides, and the swizzle mode. The hardware walks the tile itself.

// SS form: both operands described by shared-memory descriptors.
wgmma.fence.sync.aligned;                       // order prior writes vs. the async read
fence.proxy.async;                              // generic proxy -> async proxy

wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16
      {d0,...,d127}, a-desc, b-desc, 1, 1, 1, 0, 0;
//                       scale-d ^  ^ scale-a  ^ scale-b  ^ trans-a ^ trans-b

wgmma.commit_group.sync.aligned;                // close the batch
wgmma.wait_group.sync.aligned 1;                // keep 1 group in flight

// The layout atoms CuTe uses to keep SMEM indexing and the descriptor consistent.
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<bfloat16_t>;
auto smem_layout = tile_to_shape(SmemLayoutAtom{}, Shape<_128,_64>{});

// Swizzle<3,4,3> is the 128B mode: 8 rows x 128 B, chunk ^= row.
//   include/cute/atom/mma_traits_sm90_gmma.hpp

The accumulator cost is worth noting: for m64nNk16 with FP32 accumulate, each of the 128 threads holds N/2 registers. At N=256 that is 128 registers per thread for the accumulator alone — half the architectural maximum. That is the pressure that eventually pushed the accumulator out of the register file entirely.

If the tensor core reads its operands straight out of shared memory, then the kernel's remaining job is to keep shared memory full. Hopper's answer is the Tensor Memory Accelerator: a single thread issues one instruction describing a multidimensional tile copy, and dedicated hardware generates every address, detects out-of-bounds elements, applies the swizzle, and signals completion. Out-of-bounds loads are filled with zero or the configured OOB-NaN value; stores are discarded rather than clamped to an edge element.

The descriptor is built on the host:

CUtensorMap map;
cuTensorMapEncodeTiled(&map,
    CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, /*rank*/ 2, A_gmem,
    globalDim, globalStrides,            // strides: multiples of 16 B, < 2^40
    boxDim, elementStrides,              // boxDim[0]*sizeof(elem) <= swizzle size
    CU_TENSOR_MAP_INTERLEAVE_NONE,
    CU_TENSOR_MAP_SWIZZLE_128B,          // must match the SMEM layout atom
    CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
    CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
// Pass it into the kernel as: const __grid_constant__ CUtensorMap map

// One elected thread arms the barrier and issues the copy.
if (elect_one()) {
  mbarrier.arrive.expect_tx.shared::cta.b64 _, [bar], TILE_BYTES;
  cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes
      [smem], [tensorMap, {k_off, m_off}], [bar];
}
// Everyone waits on the phase bit. try_wait may return early, hence the loop.
wait: mbarrier.try_wait.parity.shared::cta.b64 p, [bar], phase;
      @!p bra wait;

auto tma = make_tma_copy(SM90_TMA_LOAD{}, gmem_tensor, smem_layout);
// ... in the kernel
copy(tma.with(barrier, mcast_mask), tAgA(_,_,k), tAsA(_,_,stage));

6.6.1 mbarriers count bytes, not just threads

A __syncthreads() counts arrivals. An mbarrier also counts transactions — bytes that asynchronous engines have promised to deliver. The handshake is:

1. A thread does mbarrier.arrive.expect_tx [bar], TILE_BYTES, which both arrives and adds TILE_BYTES to the barrier's transaction count.
2. The TMA engine decrements that count as bytes land in shared memory.
3. The barrier flips phase only when both the pending arrivals and the transaction count reach zero (PTX §9.7.14.16.8).

Consumers wait on the phase parity — a single bit that alternates each time the barrier completes. For a circular buffer of S stages, the barrier at stage s has completed k/S phases by iteration k, so the parity to wait for is (k/S) & 1. That one line is most of the bookkeeping in a Hopper mainloop.

TMA has two further tricks that matter for GEMM. Multicast: with a cluster (Part 2, §2.8), one cp.async.bulk.tensor….multicast::cluster delivers the same tile into the shared memory of every CTA named in a 16-bit mask — one HBM read feeding several SMs. And im2col mode, which unrolls convolution windows during the copy, so a convolution can use a plain GEMM mainloop.

Once loads are issued by one thread and the math is issued by a warpgroup, it stops making sense for every warp to do both. Hopper kernels split the block into producer and consumer warpgroups: the producer does nothing but issue TMA copies and manage barriers; the consumers do nothing but wgmma and the epilogue.

The producer needs almost no registers. The consumers need as many as they can get. So the ISA lets a warpgroup hand registers back:

The register pool is per-CTA. .inc blocks until enough registers are free, the value must be between 24 and 256 and a multiple of 8, and every warp of the warpgroup must execute the same instruction (PTX §9.7.20.5). Done right, a producer warpgroup costs you 4 warps' worth of slots and almost no register file.

The ping-pong part is specific to attention and is covered in Part 7; for a plain GEMM the two consumer warpgroups simply work on different halves of the output tile. FlashAttention-3 measures each contribution separately: 582 TFLOP/s with warp specialization but no GEMM–softmax pipelining, 570 with pipelining but no warp specialization, and 661 with both.

tcgen05.mma takes the trend to its conclusion. It is issued by a single thread. B always comes from shared memory; A comes from shared memory or from Tensor Memory; and the accumulator lives only in Tensor Memory. The MMA itself is asynchronous, like everything else at this generation: completion is tracked on an mbarrier by tcgen05.commit (PTX §9.7.17.12.1) — the same handshake TMA uses.

TMEM is a dedicated 256 KB scratchpad per SM, organised as 128 lanes × 512 columns of 32-bit cells. It is allocated dynamically at runtime by one warp, in units of 32 columns, power-of-two sized, and it must be explicitly freed before the kernel exits. Access is restricted by warp: warp 0 of a warpgroup can only touch lanes 0–31, warp 1 lanes 32–63, and so on — so reading back a full accumulator requires a whole warpgroup.

Assembling everything: a Hopper GEMM mainloop, one iteration, warp-specialized, with a multi-stage circular buffer.

// ---- setup ----
// SMEM: S >= 2 stages x (A tile + B tile), laid out with Layout_K_SW128_Atom
// mbarriers: full[S] (producer -> consumer), empty[S] (consumer -> producer)

if (warpgroup_id == PRODUCER) {
    setmaxnreg_dec(40);
    for (k = 0; k < num_k_tiles; ++k) {
        s = k % S;
        wait(empty[s], parity_of(k/S - 1));       // stage free?
        if (elect_one()) {
            expect_tx(full[s], TILE_BYTES);
            tma_load(A_smem[s], tensor_map_A, {k*BK, m0}, full[s]);
            tma_load(B_smem[s], tensor_map_B, {k*BK, n0}, full[s]);
        }
    }
} else {                                        // CONSUMER
    setmaxnreg_inc(232);
    wgmma_fence();                                  // once, after ordinary acc/A-reg writes
    for (k = 0; k < num_k_tiles; ++k) {
        s = k % S;
        wait(full[s], parity_of(k/S));               // bytes landed?
        wgmma_mma_async(acc, desc(A_smem[s]), desc(B_smem[s]));
        wgmma_commit_group();
        if (k > 0) {
            wgmma_wait_group<1>();                   // group k-1 is complete
            old = (k - 1) % S;
            arrive(empty[old]);                      // only now may TMA overwrite it
        }
    }
    wgmma_wait_group<0>();                       // drain the last group before reading acc
    if (num_k_tiles > 0)
        arrive(empty[(num_k_tiles - 1) % S]);      // release the final operand stage
    epilogue(acc);        // stmatrix -> SMEM -> TMA store
}

# Compact algorithmic form; it does not contractually select TMA or warp specialization.
@triton.autotune(configs=[
    triton.Config({'BLOCK_M':128,'BLOCK_N':256,'BLOCK_K':64,'GROUP_M':8},
                  num_stages=3, num_warps=8),
], key=['M','N','K'])
@triton.jit
def matmul(...):
    ...
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        a = tl.load(a_ptrs); b = tl.load(b_ptrs)
        acc = tl.dot(a, b, acc)
        a_ptrs += BLOCK_K * stride_ak; b_ptrs += BLOCK_K * stride_bk

6.9.1 So when do you need the CUDA version?

For a plain GEMM: almost never. Triton's generated Hopper kernels are competitive, and cuBLAS/CUTLASS is right there.

You reach for the low level when the schedule is the optimization — when you need a specific producer/consumer split, a ping-pong between two consumer warpgroups, an in-kernel FP8 transpose, or TMEM allocated in a particular shape. That is the FlashAttention-3 territory in Part 7, and it is the reason the fastest attention kernels are still CUTLASS.