CPU Matmul Kernels and BLAS

Why GEMM is the transformer workhorse

Open a transformer’s forward pass and count the FLOPs: the overwhelming majority live in dense matrix multiplies. Each attention block projects the input X: [T, d] through W_Q, W_K, W_V (three matmuls), recombines heads through W_O (one more), and the feed-forward network does two big ones — up-projection to [T, 4d] and back down to [T, d]. Even the softmax attention scores are QK^T and (·)V, two more matmuls. Layer-norm, softmax, and activations are real but small; they are element-wise and cost O(elements), while a matmul costs O(elements × contraction dimension).

This is why the entire performance story of CPU inference collapses onto one kernel. If GEMM runs at ten percent of the machine’s peak, your model runs at roughly ten percent of peak, no matter how clever the rest of the code is. It is also why library authors pour years into a single routine (sgemm and friends): a transformer is, to a first approximation, a pipeline of GEMM calls with thin glue between them, and optimizing the glue before the GEMM is polishing the doorknobs on a house with no walls.

What GEMM actually computes

The BLAS definition is C ← α·A·B + β·C, where A: [M, K], B: [K, N], and C: [M, N]. The α and β scalars let one routine also do accumulation (β=1) and scaling; for a plain matmul take α=1, β=0. The three dimensions have names worth internalizing: M and N are the output rows and columns, and K is the contraction (or reduction) dimension that is summed away.

The FLOP count is exact and worth memorizing: every output element C[i,j] is a dot product of length K, which is K multiplies and K adds, so 2K FLOPs. There are M·N outputs, giving

FLOPs(GEMM) = 2 · M · N · K

The data touched, by contrast, is only M·K + K·N + M·N elements. Compute grows as a product of three dimensions while data grows as a sum of pairwise products — that asymmetry is the reason GEMM can be made compute-bound, and the reason a bad implementation squanders it.

The naive triple loop and why it crawls

The textbook implementation is three nested loops:

for i in range(M):          # output row
    for j in range(N):      # output col
        acc = 0
        for k in range(K):  # contraction
            acc += A[i,k] * B[k,j]
        C[i,j] = acc

It is correct, and on a modern core it might reach 1–3 GFLOP/s against a peak north of 100 GFLOP/s. The arithmetic is not the problem — the access pattern is. Assume row-major storage. As k advances in the inner loop, A[i,k] walks contiguously (good), but B[k,j] jumps by a full row of N elements every step (bad). For N = 4096 FP32, that is a 16 KB stride — a different cache line every single access, so B is effectively read from DRAM the whole time.

Worse, nothing is reused. The inner dot product loads 2K values to produce one output, then throws the loaded data away. Each element of B gets fetched M separate times across the outer loops. The CPU’s vector units and fused-multiply-add pipelines sit almost idle, starved by a memory system the code actively fights.

The memory hierarchy is the whole game

To see why blocking works you have to hold the cost of a memory access in your head. The numbers vary by chip, but the ratios are stable and brutal:

A DRAM miss can cost fifty times an L1 hit. Meanwhile a core with AVX and FMA can retire dozens of FLOPs per cycle. If every multiply waits ~200 cycles for its operand, the arithmetic units are idle 99% of the time — which is exactly the naive loop’s fate. The entire craft of a fast GEMM is therefore data reuse: arrange the computation so that once a value has been dragged up the hierarchy into a fast level, you do as much arithmetic with it as possible before it is evicted. Every technique that follows — reordering, blocking, packing, register tiling — is a different lever on that one idea.

Arithmetic intensity and the roofline

The concept that makes this quantitative is arithmetic intensity (AI): the ratio of FLOPs performed to bytes moved from memory.

AI = FLOPs / bytes moved   (units: FLOP/byte)

The roofline model plots achievable performance against AI. Below a cutoff you are memory-bound: performance rises with AI along a slope equal to memory bandwidth, because bytes are the bottleneck. Above it you are compute-bound: performance flattens at the machine’s peak FLOP/s. The corner — the ridge point — sits at AI = peak_FLOPs / peak_bandwidth. For a core with, say, 100 GFLOP/s and 25 GB/s of usable bandwidth, the ridge is at AI = 4 FLOP/byte; you must do at least four FLOPs per byte fetched to have any hope of being compute-bound.

Here is the punchline for GEMM. Its intrinsic AI, if each matrix is read exactly once, is enormous — roughly K/3 for square problems — so GEMM is inherently a compute-bound operation. The naive loop destroys that by re-reading B from DRAM M times, crashing its effective AI back below the ridge. Blocking is the art of recovering the AI the math already gave you.

Loop reordering: the cheapest win

Before any blocking, simply changing the loop order helps, because it changes which array is walked contiguously. The naive i, j, k order strides badly through B. Swap to i, k, j:

for i in range(M):
    for k in range(K):
        a = A[i,k]              # scalar, loaded once
        for j in range(N):      # contiguous in B and C
            C[i,j] += a * B[k,j]

Now the inner loop over j streams along a row of B and a row of C, both contiguous in row-major memory. The value A[i,k] is hoisted into a register and broadcast across the whole inner loop. This access pattern is also exactly what a SIMD unit wants — contiguous loads it can vectorize — so the compiler can auto-vectorize the inner loop. Reordering alone, with no other change, commonly buys a several-fold speedup.

But it is not enough. The i, k, j loop still streams all of B and C through cache for each i, and for large matrices they do not fit. To go further you must stop thinking in single rows and start thinking in blocks.

Cache blocking: tiling the problem

Blocking (or tiling) partitions the three loops so that a small sub-problem — a block of A times a block of B accumulating into a block of C — fits inside a fast cache level and is reused fully before eviction:

for i0 in range(0, M, Mc):
  for k0 in range(0, K, Kc):
    for j0 in range(0, N, Nc):
      # multiply the tiles that fit in cache:
      #   A[i0:i0+Mc, k0:k0+Kc]  (Mc x Kc)
      #   B[k0:k0+Kc, j0:j0+Nc]  (Kc x Nc)
      # accumulate into C[i0:i0+Mc, j0:j0+Nc]

The reuse is the point. Once a Kc × Nc tile of B is resident in cache, every one of the Mc rows of the A tile multiplies against it, so each loaded B value is used Mc times rather than once. Symmetrically the A tile is reused across the Nc columns. The bytes come from DRAM once per tile but feed O(Mc) or O(Nc) FLOPs each — the effective arithmetic intensity climbs back toward the intrinsic value, and the operation moves rightward across the roofline toward the compute-bound ceiling.

Blocking for every cache level at once

Real high-performance GEMM does not block once; it nests blocks so that each level of the memory hierarchy holds the working set of the loop that lives there. The design used by GotoBLAS, OpenBLAS, and BLIS chooses three cache block sizes — Mc, Nc, Kc — plus two tiny register-tile dimensions Mr, Nr, and assigns each a home:

The dimensions are tuned so each tile just fits its level with room for the others. A rough L1 budget, for instance, must simultaneously hold the Kc × Nr B-sliver, a column of the A-block, and leave the Mr × Nr C-tile pinned in registers. This is why the block constants in a BLAS library look like magic numbers: they are solved against the specific cache sizes, associativity, and vector width of a target micro-architecture.

Packing into contiguous panels

Even a perfectly sized tile has a hidden enemy: stride. A Mc × Kc sub-block of a big matrix is not contiguous in memory — consecutive rows are K elements apart — so streaming it still scatters across cache lines and thrashes the TLB. The fix is packing: before the hot loops run, copy each tile into a fresh, contiguous scratch buffer laid out in exactly the order the micro-kernel will read it.

Packing costs an extra copy, but it pays for itself many times over. Inside the packed buffer, the micro-kernel’s loads are unit-stride and cache-line-aligned; there are no page-crossing penalties and hardware prefetchers lock on perfectly. The A-block is typically packed so its Mr-row micro-panels are stored back-to-back, and the B-panel so its Nr-column micro-panels are contiguous. Because the packed A-block is reused across the entire N dimension and the packed B-panel across the entire M dimension, the one-time packing cost is amortized over a huge amount of arithmetic. This copy-to-pack step is one of the biggest reasons a tuned GEMM crushes a naive tiled loop that skips it.

The micro-kernel: register accumulation

At the very center sits the micro-kernel, the small hand-tuned routine that does the actual multiplying and where nearly all the FLOPs happen. It computes one Mr × Nr tile of C — small enough to live entirely in vector registers — by looping over the shared Kc dimension with a sequence of rank-1 updates:

# C_tile (Mr x Nr) is held in registers the whole time
for k in range(Kc):
    a_col = packA[:, k]     # Mr values (broadcast)
    b_row = packB[k, :]     # Nr values (a vector)
    C_tile += outer(a_col, b_row)   # Mr*Nr FMAs

The decisive trick is that C_tile never leaves the register file. Across all Kc iterations the partial sums accumulate in-register, so the Mr × Nr outputs are written back to memory exactly once, at the end. Only the thin A and B slivers stream in from L1. Each k-step moves Mr + Nr elements and performs Mr × Nr fused multiply-adds — the register tile is precisely the lever that pushes arithmetic intensity above the roofline ridge.

SIMD and FMA inside the kernel

The micro-kernel earns its peak by mapping onto the core’s vector hardware, the subject of this series’ sibling article on SIMD and BLAS, so the sketch here is deliberately brief. Two instruction-set features do the heavy lifting. SIMD registers hold a vector of lanes — AVX-512 packs 16 FP32 values — so one instruction multiplies sixteen products at once; the Nr dimension of the tile is chosen as a multiple of that width. FMA (fused multiply-add) computes d = a·b + c in a single instruction at full throughput, which is exactly the rank-1 accumulation shape.

Two more details unlock the peak. The micro-kernel keeps several independent accumulator registers so the deeply pipelined FMA units never stall waiting on a dependency — each cycle issues work to an accumulator whose previous result has already retired. And a_col is delivered by a broadcast load, splatting one scalar across all lanes. The result: a well-built micro-kernel sustains something close to vector_width × FMA_units × 2 FLOPs per cycle — the machine’s advertised peak.

Threading over tiles

Everything so far is single-core. A modern CPU has many cores, and GEMM parallelizes cleanly because different output tiles are independent — computing C[i0.., j0..] never needs C[i1.., j1..]. The standard approach parallelizes one or two of the outer blocking loops, handing each thread its own set of Mc-blocks or Nc-panels of the output.

The subtlety is memory, not correctness. Threads share the L3 cache and the DRAM controllers, so naive parallelism can turn a compute-bound kernel back into a bandwidth-bound one when every core hammers memory at once. Good implementations parallelize the loop over Nc so that cores share the same packed A-block (kept in L2 or L3) while streaming different B-panels, maximizing reuse of the shared bytes. On multi-socket or NUMA systems there is a further rule: keep each thread’s data on its local memory node, because a cross-socket DRAM access can cost twice a local one. Threading multiplies throughput by core count only when the memory system can still feed every core — the roofline applies per-machine, not per-core.