Matmul, one optimization at a time

Matmul is the op worth understanding in detail, for three reasons. It is 90–95% of the FLOPs in a transformer. It sits far above the roofline ridge, so it is the one place where arithmetic efficiency is the whole game. And its optimization ladder is the canonical demonstration of every idea in Parts 1–3, in order, each with a measurable payoff.

We are computing C = A B with A of shape M×K, B of K×N, C of M×N. The bare product is 2MNK FLOPs against a compulsory (MK + KN + MN)·s bytes: read A and B, write C. BLAS's more general C ← αAB + βC adds another MN·s read when β ≠ 0. For 4092³ in FP32, the bare product is 137.0 GFLOP against 134 MB read plus 67 MB written = 201 MB; Boehm's BLAS-form benchmark also reads C, for 268 MB total. The A6000's 38.7-TFLOP/s peak gives a 3.54 ms arithmetic floor, while Boehm's measured 23.25-TFLOP/s cuBLAS result takes 5.89 ms; his simplified 30-TFLOP/s estimate is 4.57 ms. Even the larger 268 MB transfer is only about 0.34 ms at peak bandwidth. A good GEMM is decisively compute-bound. Every rung below is a step toward actually achieving that.

__global__ void sgemm_naive(int M, int N, int K, const float* A, const float* B, float* C)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;   // row of C
    int y = blockIdx.y * blockDim.y + threadIdx.y;   // col of C
    if (x < M && y < N) {
        float acc = 0.f;
        for (int k = 0; k < K; ++k)
            acc += A[x*K + k] * B[k*N + y];
        C[x*N + y] = acc;
    }
}

C = A @ B      # the thing we are trying to match

Problem one: it re-reads everything. Each of the MN threads reads a whole row of A and a whole column of B, so the kernel moves 2MNK elements — an arithmetic intensity of 1 FLOP per element, which Part 3 says is hopeless.

Problem two: half the accesses are uncoalesced. The variable named x is derived from threadIdx.x, so it varies fastest within a warp — and it indexes the row of both A and C. The 32 lanes of a warp therefore read 32 addresses K×4 bytes apart: 32 separate sectors instead of 4.

Fixing the mapping so that consecutive lanes read consecutive columns takes the kernel from 309 to 1,986 GFLOP/s. GMEM throughput goes from 15 GB/s to 110 GB/s. Nothing else changed.

The next rung stages tiles in shared memory. Each block owns a 32×32 output tile and walks K in chunks of 32: load a 32×32 slab of A and of B into SMEM, synchronize, do 32 multiply-adds per thread out of SMEM, synchronize, advance.

__shared__ float As[32*32], Bs[32*32];
float acc = 0.f;
for (int kb = 0; kb < K; kb += 32) {
    int aCol = kb + tx, bRow = kb + ty;
    As[ty*32 + tx] = (row < M && aCol < K) ? A[row*K + aCol] : 0.f;
    Bs[ty*32 + tx] = (bRow < K && col < N) ? B[bRow*N + col] : 0.f;
    __syncthreads();
    for (int k = 0; k < 32; ++k)
        acc += As[ty*32 + k] * Bs[k*32 + tx];
    __syncthreads();                            // before overwriting the tile
}
if (row < M && col < N) C[row*N + col] = acc;

Why so little? Because the arithmetic intensity is still low. With a 32×32 tile each result costs K/16 global accesses, down from 2K — a 32× reduction in DRAM traffic — but each thread now does two shared-memory loads per FMA. The bottleneck simply moved from HBM to SMEM. Nsight shows 66% occupancy, so occupancy is not the problem either.

This is the step that matters. Instead of one accumulator per thread, each thread holds an 8×8 block of accumulators in registers, and the inner loop becomes an outer product: load 8 values of A and 8 of B from shared memory, and do 64 FMAs.

// BM=BN=128, BK=8, TM=TN=8  →  256 threads, each owning an 8×8 tile of C.
float acc[TM][TN] = {0.f};     // 64 registers of accumulator per thread
float regA[TM], regB[TN];

for (int kb = 0; kb < K; kb += BK) {
    loadTilesToSmem(As, Bs, A, B, kb);
    __syncthreads();

    for (int k = 0; k < BK; ++k) {
        // 16 SMEM loads ...
        for (int i = 0; i < TM; ++i) regA[i] = As[k*BM + threadRow*TM + i];
        for (int j = 0; j < TN; ++j) regB[j] = Bs[k*BN + threadCol*TN + j];
        // ... feeding 64 FMAs, all independent -> deep ILP
        for (int i = 0; i < TM; ++i)
            for (int j = 0; j < TN; ++j)
                acc[i][j] += regA[i] * regB[j];
    }
    __syncthreads();
}

# Triton expresses the same thing as a block-level dot product.
# You choose the tile; the compiler chooses the register blocking.
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, mask=offs_k[None, :] < K - k*BLOCK_K, other=0.0)
    b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k*BLOCK_K, other=0.0)
    acc = tl.dot(a, b, acc, input_precision="ieee")  # match FP32 SGEMM semantics
    a_ptrs += BLOCK_K * stride_ak
    b_ptrs += BLOCK_K * stride_bk

Notice what this does to the numbers from Part 1. Each thread now uses 64 registers for accumulators alone, plus operand registers and addresses — well over 100. At 128 registers per thread the SM fits 65536 / (32 × 128) = 16 warps: a third of the A6000's 48-warp maximum, a quarter of Hopper's 64. That low occupancy is the deliberate purchase. This is Volkov's tradeoff in its natural habitat: you spend occupancy to buy instruction-level parallelism, and the 64 independent FMAs in the inner loop cover the latency that the missing warps would have covered.

Vectorized access (Part 2, §2.3) is the next 10%: transpose the A tile as it is stored into shared memory so that both the global loads and the shared loads can be 128-bit. Boehm measures 18,237 GFLOP/s, 78.4%, and notes the SMEM transpose alone is worth about 3%.

Warptiling inserts a level between the block tile and the thread tile. It exists because a warp is a real hardware unit: the 32 threads of a warp issue together and their shared-memory accesses are coalesced together, so giving a warp a contiguous rectangle of the output tile makes those accesses regular. The hierarchy becomes:

That is 21,779 GFLOP/s, 93.7% of cuBLAS. It is also exactly the hierarchy CUTLASS uses, and exactly the level at which tensor cores plug in — in Part 6 the innermost level stops being an FMA and becomes an mma.sync or wgmma instruction, but the three enclosing levels are unchanged.

Look again at the structure of the kernel: load tile → barrier → compute → barrier → load next tile. The tensor cores or FMA pipes are idle for the whole load, and the memory system is idle for the whole compute. Perfect overlap changes a serial cost L+C into roughly max(L,C): a gain of up to 2×, reached only when load and compute take equal time.

The fix is double buffering, and then multi-stage pipelining — the structure from Part 2, §2.7. The exact operand path changes with the instruction generation:

• Ampere: cp.async fills circular shared-memory stages; ldmatrix double-buffers register fragments while mma.sync consumes the other set.
• Hopper: TMA fills circular shared-memory stages and the common WGMMA SS form reads A and B there through descriptors, with no ldmatrix stage. The WGMMA RS form can instead keep A in registers while B stays in shared memory.

Two problems remain, and neither is fixed by touching the mainloop.

5.6.1 L2 rasterization

CUDA does not guarantee the order in which blocks execute, and a wave need not be exactly one block per SM. But schedulers commonly draw nearby block IDs into the same period of execution, so the mapping is still a useful locality heuristic. With row-major tile IDs, a wave can span a long row of the output: its blocks reuse one row-panel of A but touch many different column-panels of B. Re-mapping nearby IDs onto a compact 2D region gives both operands a better chance of reuse in L2; measure the result rather than depending on a launch-order guarantee.

# Grouped ordering — five lines, worth >10% on some hardware.
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m

// CUTLASS calls this a threadblock swizzle; it maps consecutive
// block indices onto "packed two-dimensional regions" of the output
// to maximise last-level-cache reuse.
//   cutlass/gemm/threadblock/threadblock_swizzle.h

5.6.2 Not enough tiles: split-K and Stream-K

The other problem is the shape of LLM decode. With M = 1 and a 128×256 tile, an 8192×8192 weight matrix produces 32 output tiles for 132 SMs. Three quarters of the GPU is idle, and no inner-loop optimization changes that.

The available parallelism is along K, and it is usable because the partial sums can be reassociated. Because floating-point addition is not associative, that changes rounding; it becomes run-to-run nondeterministic only if the implementation also uses an order that can vary, such as contending atomics. A fixed reduction tree can be deterministic. Split-K partitions the reduction dimension across blocks, each computing a partial product, followed by a reduction kernel. Stream-K generalizes it: instead of assigning tiles to blocks, it assigns an even share of the total inner-loop iterations to a fixed set of blocks, then fixes up the partial sums.

The Stream-K paper (Osama et al., 2023) reports up to 14× over data-parallel CUTLASS with the same blocking, and 6.7× over cuBLAS, measured across 32,824 GEMM shapes — from a single tile configuration per precision. The mechanism is the quantization problem from Part 1. Colfax's worked example: on an H100 PCIe with 114 SMs, going from 114 output tiles to 115 roughly halves device utilization — and Stream-K removes that cliff.

When the mainloop ends, the whole output tile is sitting in registers. Work done there is usually much cheaper than a separate kernel, because it avoids writing 128×128 floats to HBM and reading them back. It is not literally free: extra instructions, registers, shared memory or barriers can reduce throughput and occupancy. Standard fusions:

• bias add, scaling, and the residual add
• activation — GELU, SiLU, and the SwiGLU gate multiply
• quantization back down to FP8/FP4, including computing the per-block scale (Part 8)
• the transpose or layout change needed by the next op

CUTLASS expresses these as epilogue visitor trees so that arbitrary elementwise DAGs can be attached without writing a new kernel. In Triton you just write the arithmetic after the k loop, before tl.store. On Hopper the epilogue usually stages the result back through shared memory with stmatrix so that a TMA store can write it out in one swizzled transfer — Part 6.