Matmul II: register blocking, warp tiling, tensor cores

1. The ratio, and the only faster memory

The tiled kernel of chapter 7 does one FMA per two shared-memory loads, and an SM can issue four warp-wide FMAs but only one warp-wide shared load per cycle. To use the FMA units, each value that comes out of shared memory has to be used several times, and the only place to keep a value between uses is a register: the fastest storage on the chip, private to one thread, read by the FMA units directly and at full speed.

The trick that makes this possible is to look at matrix multiply differently. The dot-product view says C[i][j] = ∑_k A[i][k]·B[k][j]: one output at a time. The outer-product view swaps the loops: for each k, column k of A times row k of B is a whole matrix of products, and C is the sum of those K matrices. A thread that owns a TM×TN block of C only needs, at each k, TM values of A’s column and TN values of B’s row, and every pair of them is one of its FMAs.

So the design of a fast GEMM is a nest of the same idea at three levels, each trading on-chip storage for reuse: a block tile in shared memory (reuse across the block, chapter 7), a warp tile (section 5), and a thread tile in registers (sections 2–4). The rest of the chapter adds them one at a time, as lab 7’s kernels v4 to v7.

2. v4: one column of outputs per thread

The smallest step: each thread computes TM = 8 outputs stacked in one column of the tile. Per step k it loads one value of B into a register and reuses it for all 8 of its outputs, reading the 8 A values it needs from shared memory. The block computes a 64×64 tile with 512 threads, and each phase loads a 64×8 slice of A and an 8×64 slice of B, exactly one element of each per thread:

for (int k = 0; k < BK; ++k) {                 // BK = 8
    float b = Bs[k][tcol];                        // one load...
    for (int i = 0; i < TM; ++i)                  // ...reused 8 times
        acc[i] += As[trow * TM + i][k] * b;       // the same As word for the whole warp: broadcast
}

Nine loads for eight FMAs: 1.125 per FMA instead of 2, which moves the shared-memory ceiling from 1/8 of peak to about 1/4.5. The As loads are broadcasts, since all 32 lanes of a warp share trow, and the 8 accumulators live in registers. Boehm measured this step at 36% of cuBLAS, almost three times the tiled kernel, for a change of a dozen lines.

3. v5: a block of outputs per thread

Now the full outer product of figure 1: each thread owns an 8×8 block, so the block tile grows to 128×128 with 256 threads, and the inner loop becomes

for (int k = 0; k < BK; ++k) {
    for (int i = 0; i < TM; ++i) regM[i] = As[trow * TM + i][k];   // 8 loads
    for (int j = 0; j < TN; ++j) regN[j] = Bs[k][tcol * TN + j];   // 8 loads
    for (int i = 0; i < TM; ++i)
        for (int j = 0; j < TN; ++j)
            acc[i][j] += regM[i] * regN[j];                        // 64 FMAs
}

0.25 loads per FMA: shared memory can now supply the FMA units with room to spare, and the ceiling of chapter 7 is gone. The block tile got bigger too: 128×128 with BK = 8 means each phase loads 128 + 128 values per k for 2×128×128 FLOPs, a global intensity of 32 FLOP/B if nothing hit in L2 (chapter 4’s “register tiling” point). Each thread now loads four elements of each tile per phase, in a strided loop, which keeps the loads coalesced.

The price: registers

64 accumulators, 16 operand registers, plus indices and addresses: expect around 100 registers per thread (-Xptxas -v prints the exact number when you build). At that size an SM’s 64K registers hold only two 256-thread blocks: 16 warps, 33% occupancy. Chapter 2 would have called that a problem. It is not, and why it is not is one of the most useful ideas in GPU performance.

Occupancy exists to hide latency: while one warp waits, another issues. But latency can also be hidden inside a warp, if the warp has many independent instructions to issue while it waits. The 64 FMAs of one step are all independent of each other, and the loads of the next step can be issued before they are done. That is instruction-level parallelism doing the job of thread-level parallelism. Vasily Volkov made this point in a 2010 talk titled “Better performance at lower occupancy”, and every fast GEMM since lives at low occupancy with a large register tile. Chapter 2’s callout, “occupancy is a means, not a goal”, is this.

4. v6: wider loads and a transposed tile

v5 fixed the ratio but still issues one instruction per float moved: 16 LDS per step for 64 FFMA, plus the global loads of the tiles, one float each. Instruction issue is itself a shared resource (one instruction per scheduler per cycle), so v6 moves 16 bytes per instruction wherever it can:

• Global loads as float4: each thread loads 4 consecutive floats of A or B with one 16-byte instruction. For the B tile, a warp’s 32 such loads cover one 512-byte row segment; for the A tile, 32-byte pieces of 16 rows. Either way every sector fetched is fully used.
• The A tile stored transposed, as As[BK][BM]. Now for a fixed k, a thread’s 8 values of A are contiguous, and regM is two LDS.128 instead of eight LDS. regN was already contiguous in Bs: two more LDS.128. Four load instructions per 64 FMAs.

// one step k: 4 LDS.128 feed 64 FFMA
*reinterpret_cast<float4*>(&regM[0]) = *reinterpret_cast<const float4*>(&As[k][trow * 8]);
*reinterpret_cast<float4*>(&regM[4]) = *reinterpret_cast<const float4*>(&As[k][trow * 8 + 4]);
*reinterpret_cast<float4*>(&regN[0]) = *reinterpret_cast<const float4*>(&Bs[k][tcol * 8]);
*reinterpret_cast<float4*>(&regN[4]) = *reinterpret_cast<const float4*>(&Bs[k][tcol * 8 + 4]);

Vector loads come with a condition: the address must be a multiple of 16 bytes. A row of A starts at row * K * 4 bytes, which is 16-byte aligned for every row only if K is a multiple of 4; the same goes for N and B. So lab 7’s v6 and v7 come in two variants, chosen once per launch: with VEC they use 16-byte loads wherever a group of 4 is fully inside the matrix and fall back to four guarded scalar loads at the edges; without it (an odd K or N) they always use the scalar path. The test shapes include both kinds.

5. v7: warp tiling

In v5 and v6, threads are numbered across the block tile row by row: thread t owns rows 8×(t/16) and columns 8×(t mod 16). A warp, threads 0–31, therefore owns a strip 16 rows tall and 128 columns wide. At each step k the warp as a whole reads 16 values of A and 128 values of B from shared memory to produce its 2,048 outputs.

The amount of shared memory a warp reads per step is the perimeter of its region (height + width), while the work it does is the area (height × width). For a fixed area, a square minimizes the perimeter. So v7 gives each warp an explicit square: the 128×128 block tile is split into four 64×64 warp tiles, one per warp (128 threads), and each thread’s 128 outputs are arranged inside its warp’s square.

The indexing is the only hard part. With a 64×64 warp tile and 32 lanes of 128 outputs each, v7 splits the warp tile into WNITER = 4 vertical strips of 16 columns; in each strip, a lane owns an 8×4 block (lane / 4 picks the 8 rows, lane mod 4 picks the 4 columns). Per step k a thread loads 8 values of A (two LDS.128) and 4 × 4 of B (four LDS.128) and does 8 × 16 = 128 FMAs:

// lab 7, sgemm_warptile: BM = BN = 128, BK = 16, WM = WN = 64, WNITER = 4, TM = 8, TN = 4
const int warpRow = warp / 2, warpCol = warp % 2;          // which 64 x 64 square
const int tRowInWarp = lane / 4, tColInWarp = lane % 4;    // which 8 x 4 block in each strip
for (int k = 0; k < BK; ++k) {
    load 8 floats:   As[k][warpRow*64 + tRowInWarp*8 + 0..7]                 -> regM[0..7]
    for ws in 0..3:  Bs[k][warpCol*64 + ws*16 + tColInWarp*4 + 0..3]         -> regN[ws*4 + 0..3]
    acc[8][16] += outer(regM, regN)                                          // 128 FMAs
}

Besides the traffic, the warp tile matters for a second reason, and it is the one that makes this section the bridge to the rest of the chapter: a warp is the unit a tensor core works on. A tensor-core instruction takes a small tile of A and B held collectively by the 32 lanes of a warp. Every tensor-core GEMM is therefore organized exactly like v7: block tile in shared memory, warp tile per warp, and inside the warp, fragments in registers.

Explore the hierarchy. The presets are lab 7’s kernels; change the shapes and watch registers, occupancy, loads per FMA and shared traffic per output move:

6. Tensor cores

Everything so far used the FP32 lanes, the “CUDA cores”: one multiply-add per lane per cycle. Since Volta (2017), each SM also has tensor cores: units that multiply small matrices. A tensor-core instruction is issued by a whole warp and computes D = A·B + C on tiles whose elements are spread across the 32 lanes’ registers. On your card the basic FP16 instruction is mma.sync.m16n8k16: A is 16×16, B is 16×8, C and D are 16×8, so one warp instruction performs 2,048 multiply-adds, where an FFMA performs 32.

The table behind that figure, in the units of chapter 1:

Two things in that table deserve a comment. First, on a GeForce card the tensor-core advantage for FP16 with FP32 accumulation is “only” 2× the CUDA cores; on an H100 it is about 15×. The halving is a product decision (NVIDIA runs FP32 accumulation at half rate on consumer cards), not physics, and it means that on your card the bytes saved by FP16 matter as much as the FLOPs. Second, the ridge moves: at 61.7 TFLOPS and 672 GB/s, a kernel needs 92 FLOP/B to be compute-bound (chapter 1), so tensor-core GEMMs need bigger tiles and more reuse than FP32 ones to reach their roof.

Where the operands live: fragments

A tensor-core instruction has no pointer operands. Its A, B, C and D tiles are in registers, distributed over the lanes of the warp: for m16n8k16, each lane holds 8 halves of A (in 4 registers), 4 halves of B (2 registers) and 4 floats of C and D. Which lane holds which element is fixed by the PTX specification, and it is not an obvious layout. Click a lane:

Two consequences. The 32 lanes together hold each tile exactly once: the warp cooperates, and no single thread “has” a row. And getting data from shared memory into this layout is its own problem: a plain load per element would be slow, so there is a dedicated instruction, ldmatrix (LDSM in SASS), that loads 8×8 blocks of 16-bit values from shared memory straight into the fragment layout, each lane supplying the address of one row.

Three ways to program them

WMMA is the simplest and it compiles to the same tensor-core instructions, but because it hides the layout, you cannot do elementwise work on a fragment that depends on the element’s position (masking, per-row scaling) without a round trip through shared memory. That is why serious kernels drop down to mma.sync.

7. Feeding the tensor cores: t1 to t3

A tensor core consumes operands fast: at 256 multiply-adds per SM per clock, an m16n8k16 (2,048 multiply-adds, 768 bytes of A and B) finishes in 8 clocks per SM. If every fragment came from shared memory once per MMA, that is 96 bytes per clock, most of shared memory’s 128. So the whole hierarchy of this chapter reappears, one level up: fragments must be reused from registers across several MMAs (the warp tile), tiles must be reused from shared memory across warps (the block tile), and the loads must overlap the math. Lab 7’s hgemm takes the three steps:

• t1: one warp per 16×16 tile, fragments straight from global memory. The tensor-core version of chapter 7’s naive kernel: no shared memory, no reuse beyond what L1 and L2 find. Every mma_sync waits for its own global loads.
• t2: 128×128 block tile in shared memory, 8 warps of 64×32. Each phase copies a 128×32 tile of A and a 32×128 tile of B into shared memory with 16-byte loads; each warp loads 4 A fragments and 2 B fragments per 16-wide step and does all 8 MMAs, so each A fragment is used twice and each B fragment four times from registers. Rows of the shared tiles are padded by 8 halves, so that the 16 rows of a fragment do not start in the same banks.
• t3: t2 with asynchronous copies into a double buffer. In t2, every phase is “load, wait, compute, wait”: while the tile loads, the tensor cores idle. t3 keeps two buffers and, while the warps compute on one, the copies for the next phase are already in flight into the other.

What makes t3 possible is cp.async: a copy from global to shared memory that the thread issues and then does not wait for. It does not pass through registers, so the thread can go straight on to issuing tensor-core instructions. In CUDA C++ it is __pipeline_memcpy_async(dst, src, 16); copies are grouped with __pipeline_commit(), and __pipeline_wait_prior(n) waits until at most n groups are still in flight. The main loop of t3:

load_tiles_async(buffer 0, tile 0);  commit();
for (int t = 0; t < tiles; ++t) {
    if (t + 1 < tiles) load_tiles_async(buffer (t + 1) % 2, tile t + 1);   // start the next copy
    commit();
    wait_prior(1);          // everything except the newest group has landed: tile t is in
    __syncthreads();        // ...for every thread's copies, not just mine
    compute(buffer t % 2);  // 8 warps x 8 MMAs x 2 steps
    __syncthreads();        // nobody reads buffer t % 2 any more: iteration t + 1 may refill it
}

8. What cuBLAS does that we do not

v7 and t3 contain the ideas that matter. The remaining distance to cuBLAS is made of many smaller ones, each worth a few percent, and knowing their names is useful even before you need them:

• Autotuning. The best BM, BN, BK, warp tile and thread tile depend on the GPU and the matrix shape. cuBLAS ships many kernel variants and a heuristic that picks one per call; Boehm gained about 6 points of cuBLAS on his card by searching tile sizes alone.
• Deeper pipelines: three or four shared-memory stages instead of two, so that a load that is slower than one phase of math is still hidden.
• Swizzled shared-memory layouts instead of padding: XOR-permuting the columns of each row so that both the stores into the tile and the ldmatrix loads out of it are conflict-free, without wasting memory.
• ldmatrix and hand-placed mma.sync, with fragments double-buffered in registers, so that the loads for the next k step overlap the MMAs of this one.
• Block rasterization: launching blocks in an order (for example, in groups along a column) that maximizes L2 hits between blocks that run at the same time.
• Different algorithms for different shapes: split-K (several blocks share one output tile and add their partial sums) when M·N is too small to fill 48 SMs, and dedicated GEMV kernels when M = 1, which is decode (chapter 15).

The last point matters most for this course. The square 4096³ GEMM of the benchmark is prefill’s shape. Decode’s shape is M = 1 to a few dozen, where none of this chapter’s compute tricks apply and the kernel is a memory-bound GEMV: the game there is bytes, and Part IV’s quantized GEMV kernels are where it is played.

9. Predictions, and what the profiler should show

At 4096³, from the mechanisms and from Boehm’s measured ratios on a different card; ranges, not promises:

Two comparisons to make when the numbers are in. Your best FP32 kernel against cuBLAS FP32, which is the chapter’s exam. And cuBLAS HGEMM against cuBLAS SGEMM on the same shape: that ratio is what switching the spine engine’s linear layers to FP16 is worth, before any quantization.

10. Lab 7: register blocking and tensor cores

cd ~/unibo-lessons/mlsys/labs/lab07-matmul-registers
make               # read the -Xptxas -v lines: registers and spills per kernel
make test          # sgemm2 and hgemm: all PASS
make bench         # the two tables at 1024, 2048, 4096 cubed
make sass          # FFMA, LDS, LDS.128, HMMA, LDSM per kernel
make ncu-table     # FP32 kernels: shared loads, wavefronts, conflicts, stalls
make ncu-tensor    # tensor-core kernels: Speed of Light, pipes, memory

Check your understanding

Sources and further reading

• S. Boehm, How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog (2022): kernels v4–v7 follow his kernels 4, 5, 6 and 10; the bank-conflict and autotuning remarks are his findings.
• V. Volkov, Better Performance at Lower Occupancy (GTC 2010), and V. Volkov, J. Demmel, Benchmarking GPUs to Tune Dense Linear Algebra (SC 2008): register blocking and instruction-level parallelism instead of occupancy.
• NVIDIA, CUTLASS: Efficient GEMM in CUDA: the block / warp / thread tile hierarchy, software pipelining, and why each level exists.
• NVIDIA, CUDA C++ Programming Guide: warp matrix functions (WMMA) and asynchronous data copies (__pipeline_memcpy_async); PTX ISA: mma.sync.m16n8k16 fragment layouts (the lane mapping in the fragment explorer), ldmatrix, cp.async, and the sm_90a / sm_100a-only instructions.
• NVIDIA, RTX Blackwell GPU Architecture whitepaper (2025): dense tensor throughput per format for the GeForce RTX 50 series, including the half-rate FP32 accumulation on GeForce parts.
• M. Fasi, N. Higham, M. Mikaitis, S. Pranesh, Numerical behavior of NVIDIA tensor cores (PeerJ CS, 2021): how tensor cores round inside an MMA, the reason for lab 7’s looser FP16 tolerance.
• Hwu, Kirk, El Hajj, Programming Massively Parallel Processors, 4th ed. (2023), ch. 6: thread coarsening, the general form of register blocking.