Systems

Matrix multiplication on a GPU

A GEMM is exactly 2MNK operations whatever runs it. Watch a GPU tile one, then compare three ways of building the hardware for it, CUDA cores, tensor cores and a TPU systolic array, and see why the fast ones are so hard to feed.

Size
Tile
Unit
Loading 3D view

Three set-upsAlmost every layer of a model is a matrix multiplication whose arithmetic is fixed, so what really separates CUDA cores, tensor cores, and a TPU systolic array is how many bytes each one has to move for every multiply-add.

Narrated with James Phoenix's AI voice.

2.15 GFLOP
Work
136.3 MB32x less than naive
Moved from HBM
15.8FLOP/B, ridge 295
Arithmetic intensity
0 / 1024
Tiles done
32x
Reuse
40.7 µsmemory bound
Tensor cores time

At 32 x 32 tiles the bf16 tensor cores path needs 295 FLOP per byte to stay busy and this schedule only reaches 15.8. The unit idles waiting for HBM. Bigger tiles, or reuse through L2, are the only fixes.

Almost everything a language model does is this operation. Attention scores, the value mix, every feed-forward layer, the final projection to the vocabulary: matrix multiplication, or GEMM (general matrix multiply) in the language of the libraries that run it. When people say a model needs a GPU, this is the thing they mean it needs a GPU for. It is worth knowing exactly what the GPU does with it.

The work itself cannot be reduced. Each element of the output is a dot product of a row of A with a column of B, which is K multiplies and K adds, and there are M x N of them. Two times M times N times K, full stop. A 4096-wide square GEMM is 137 billion operations on any hardware ever built. What differs between a slow implementation and a fast one is not the arithmetic. It is how many bytes have to move for each operation, and how many operations can happen at once.

Start with bytes, because that is where naive code dies. If every thread computes one output element and fetches its own row of A and column of B from the GPU's main memory (HBM, the stacked chips beside the die), then each element costs 2K reads for 2K operations: one operation per element read, a fraction of a FLOP per byte. HBM on an H100 moves 3.35 terabytes a second, which sounds enormous and is nowhere near enough. The cores would spend nearly all their time waiting. Set the tile control to 16 and watch what happens instead: a block of threads claims a 16 x 16 tile of C, loads a 16-row slab of A and a 16-column slab of B into shared memory (the fast, small, on-chip scratchpad every core in the block can read), and every loaded element is reused 16 times before it is thrown away. Traffic falls by the tile side. At 128 it is 128 times less than naive. That single idea, tile and reuse, is most of what a CUDA kernel or a Triton kernel is.

The arithmetic intensity stat is the honest score: operations per byte moved. Every compute unit has a ridge, the intensity at which its arithmetic and its memory bus take the same time, and it is simply peak FLOP/s divided by bandwidth. For the fp32 CUDA cores on an H100 the ridge is about 20 FLOP per byte, and a 64-wide tile clears it comfortably. Switch to tensor cores and the ridge jumps to nearly 300. A tensor core takes a whole 4 x 4 x 4 block of multiply-adds in one instruction, sixty-four operations where a CUDA core does one, and the die has hundreds of them. They are so fast that no single level of tiling from HBM can keep them fed. Watch the stat turn red as you switch: the kernel that was compute-bound on CUDA cores is memory-bound on tensor cores with the same tiles.

That is the real engineering problem of a modern GEMM, and it is why the memory hierarchy on the right has four levels rather than two. Real kernels climb it. Registers hold the accumulators for a thread's own outputs. Shared memory holds the current tiles. The 50 MB L2 cache catches the row-block of A that the next tile over is about to ask for, so most of the "HBM traffic" this page counts is actually served from L2 at several times the bandwidth. Kernels also use rectangular tiles like 256 x 128, split K across blocks, and on Hopper multicast one tile to a cluster of blocks. Every one of those tricks is a way to raise intensity past the ridge.

The systolic array is the other answer to the same question, and it is how a TPU does it. Pick that unit and step through the ticks. There is no shared memory and no scheduling at all: a fixed grid of multiply-add cells, A entering from the left and B from the top, each value passed one cell along per tick. Cell (i, j) accumulates C[i][j] and never touches memory mid-computation. The whole product arrives after K + 2(n - 1) ticks. It is rigid, which is its weakness (a shape that does not fit the grid wastes cells) and its strength (no cache misses, no address arithmetic, nearly every transistor doing arithmetic). The 4 x 4 here is a toy. The array in a TPU is 128 x 128 or larger.

Two things the picture simplifies. The traffic counter measures bytes into the streaming multiprocessors with no L2 in the way, which is the pessimistic bound; the L2 reuse described above is why real tensor-core kernels do reach their peak on large shapes. And the walk visits output tiles one at a time, where a real GPU has 132 multiprocessors each working a different tile at once, with several blocks per multiprocessor so one can compute while another waits on a load. The bytes and the FLOPs are exact. The order and the overlap are the simplification.

The maths

The work is fixed
Cij  =  k=1KAikBkj,FLOPs  =  2MNKC_{ij} \;=\; \sum_{k=1}^{K} A_{ik}\,B_{kj}, \qquad \text{FLOPs} \;=\; 2MNK

Every output element is a dot product of length K: K multiplies and K adds. The Work counter is this number and nothing on the page changes it. Tiling, tensor cores and systolic arrays only change what each operation costs in bytes and clocks.

Bytes moved with square tiles
bytes    2MNKbt  +  MNb\text{bytes} \;\approx\; \frac{2MNK\,b}{t} \;+\; MN\,b

Each t x t output tile reads a t x K row-block of A and a K x t column-block of B once, then writes itself. With t = 1 this is the naive kernel that re-reads a whole row and column per output. The tile control is t; the "moved from HBM" counter is this expression with ragged edges charged in full.

Arithmetic intensity and the ridge
I  =  FLOPsbytes    tb,I  =  peak FLOP/sbandwidthI \;=\; \frac{\text{FLOPs}}{\text{bytes}} \;\approx\; \frac{t}{b}, \qquad I^{*} \;=\; \frac{\text{peak FLOP/s}}{\text{bandwidth}}

Intensity is how much work each byte buys, and it grows with the tile side. The ridge is the intensity at which the compute unit and the memory bus take the same time. Below it the unit waits on memory; above it memory waits on the unit. The stat shows both, and the colour flips when they cross.

Systolic array latency
T  =  K+2(n1) ticksT \;=\; K + 2(n-1)\ \text{ticks}

Cell (i, j) meets a[i][k] and b[k][j] on tick i + j + k. The far corner sees its last pair on tick K - 1 + 2(n - 1), so the whole product is ready one tick later. The Tick counter in the systolic view runs to exactly this.

Related terms

More visualisations

Building with language models?

These explainers come out of the work. If you want the same thinking applied to your own system, that is what I do.

See how I can help