The Cost of Matrix Multiplication — Where Almost All of AI's Compute Goes
Why GEMM is everything: the anatomy of O(n³), memory bandwidth and arithmetic intensity, what actually makes a GPU fast, and the intuition behind tiling — ending with you able to estimate a model's training and inference FLOPs yourself.
Key points
- Modern computers have arithmetic units far faster than memory bandwidth, so runtime depends on how many bytes you move, not just how many operations you perform.
- Matrix multiplication (GEMM) of m×k by k×n costs 2mkn FLOPs, or 2n^3 for square matrices, and dominates AI compute.
- Arithmetic intensity I = FLOPs / bytes moved decides memory bound versus compute bound, capping performance at min(peak compute, I × memory bandwidth).
- GEMM has I ≈ n/6 and improves with size, while a matrix-vector product is fixed at I ≈ 0.5, which is why batching helps.
- Tiling raises effective intensity to I ≈ T/6, and the estimates are about 2N FLOPs per token for inference, 6N for training, so C ≈ 6ND.
Sources: Attention Is All You Need · cuDNN: Efficient Primitives for Deep Learning · BLAS (Basic Linear Algebra Subprograms) · Parallel Thread Execution ISA (NVIDIA CUDA Documentation) and 2 more (below)
The metaphor: the machine and the trucks
A factory installs a ferociously fast machine tool: it can cut a thousand parts per second. The trucks delivering raw stock, however, arrive with ten parts per second. This factory's output is set by the number of trucks, not by the machine. The machine stands idle 99% of the time.
Modern computers are that factory. Arithmetic units are ferociously fast; the bandwidth that carries data out of memory has not kept pace. Counting how many operations you perform therefore does not predict runtime. You have to count how many bytes you move at the same time.
And in AI, the overwhelming majority of that time goes into a single operation: matrix multiplication.
Why matrix multiplication is everything
Take a Transformer apart and what falls out is almost entirely matrix products: the projections to Q, K and V, the attention scores, the weighted sum over values, the output projection, and the two linear layers of the FFN. Convolutions, too, are typically lowered to a matrix product via im2col or an equivalent rearrangement.
That is why the library world gave this one operation a dedicated name — GEMM, general matrix multiply — and has been optimizing it for decades. Strip the marketing off "making AI faster" and a large fraction of it is "making GEMM faster".
The anatomy of
Multiply an matrix by a matrix to get an matrix . One output cell is defined as
Read in words: "row of against column of , multiplied elementwise and summed" — a single dot product. There are cells, each costing multiplications and roughly additions. Counting a multiply and an add as one operation each,
which says nothing more exotic than: cells in the output (), times multiply-adds needed to fill one cell (), times two operations per multiply-add.
That is for square matrices. Double and the work grows eightfold. That is the content of , and it is also what people mean by a model's "compute".
Numbers give it scale. A square product at costs about FLOPs — over a hundred billion operations for one multiplication. A Transformer stacks dozens of layers of those and repeats them per token.
It is worth feeling, with your hands, that the atom of all this is a dot product.
That covers the operation count. But runtime is not determined by alone — the trucks from the opening are about to matter.
Arithmetic intensity: how much work per byte delivered
Here is the factory, written as a formula. Define the arithmetic intensity:
Read in words: the numerator is how much work you did, the denominator is how many bytes you paid to move in order to do it. Think of as fuel economy for data. It is the ratio of work done to data delivered — how many times you managed to reuse each byte you paid to fetch. An operation with low stalls on bandwidth no matter how fast the arithmetic units are (memory bound). With high you get to use the units at full tilt (compute bound). Achievable performance is capped by
which says your real speed is pinned by whichever is lower: what the arithmetic units are capable of, or what the bandwidth can keep them fed with. The first term is the machine's ceiling, the second is "bytes per second times work per byte". Plotting those two lines gives you the roofline model.
Run the numbers for GEMM. For square matrices you move three of them, elements, which in fp32 is bytes. The work is , so
Put in words, the arithmetic intensity of a matrix product is nothing but its side length divided by six. Intensity grows linearly with : the bigger the matrix, the more times each fetched byte gets reused. That is the distinguishing property of matrix multiplication. At , — comfortably compute bound.
Contrast a matrix-vector product (GEMV): FLOPs against roughly bytes, so , fixed. Each element of the matrix is used exactly once, leaving nothing to reuse. This is why single-stream LLM token generation is slow and batching helps so dramatically: batching converts GEMV into GEMM and raises .
What actually makes a GPU fast
Two things: sheer count of arithmetic units, and instructions cut to the shape of a matrix product.
A CPU is built to finish one instruction as quickly as possible, spending enormous silicon area on branch prediction and speculation. A GPU assumes the opposite workload — the same operation applied across a flood of data — so it strips control logic and packs the area with arithmetic units instead. Recent accelerators go further, with dedicated units that perform a small matrix multiply-accumulate as a single instruction. It is a deliberate bet on operations with high .
Which is where intensity bites back. Multiply your arithmetic units by a thousand and any low- operation still stalls on bandwidth. A GPU is fast because matrix multiplication has high arithmetic intensity, not because computation in general got faster. That asymmetry drives real implementation decisions.
Tiling: cut the matrix into cache-sized blocks
assumed the ideal: , and all resident in fast memory. In reality a matrix does not fit in cache. A naive triple loop walks a column of on every pass of the inner loop, going out to main memory each time, and the effective intensity collapses to around 0.25.
The fix is tiling (blocking). Cut the matrices into blocks so that one block each of , and live in fast memory simultaneously. Once they are resident, that block performs operations on elements of data. Across the whole matrix this gives
which says the number governing your fuel economy is no longer the side of the whole matrix, , but the side of the block you managed to keep resident, . Bigger is better, bounded by the capacity of cache — or shared memory on a GPU. What BLAS implementations have polished for decades is, at bottom, exactly this partitioning, applied at several levels down to registers. The algorithm stays while data movement alone changes runtime by factors of several to an order of magnitude — the most blatant instance of the "outside the complexity class" effects from the complexity article.
There are also ways to attack the itself. Strassen's algorithm uses divide and conquer to reach and does get used for very large matrices, at some cost in numerical stability. The theoretical exponent has been pushed below 2.4, but with constants so vast that those are "galactic algorithms" — provably better, never run.
Estimating FLOPs yourself
This is the destination. First, a tool.
def gemm_cost(m, k, n, bytes_per_elem=4):
flops = 2 * m * k * n # one multiply-add = 2 FLOP
moved = (m * k + k * n + m * n) * bytes_per_elem
return flops, flops / moved # (FLOPs, arithmetic intensity)
print(gemm_cost(4096, 4096, 4096)) # → intensity 682.7 FLOP/byte
print(gemm_cost(4096, 4096, 1)) # → intensity 0.5 FLOP/byte
Same matrix , and the intensity differs by more than a factor of a thousand depending on whether the other operand is a matrix or a vector. The value of batching falls straight out of the arithmetic.
For a whole model, estimate from the parameter count . Each parameter contributes a multiply and an add, so a forward pass costs about FLOPs per token. Training adds a backward pass that computes gradients with respect to both inputs and weights, costing roughly twice the forward, for a total of about FLOPs per token (backpropagation). Training on tokens therefore costs
which says, in words: about six operations for every pairing of one parameter with one token. The bill for a training run is essentially model size multiplied by how much text you push through it. This approximation ignores attention's term and normalization, and holds while , i.e. while the sequence length stays below about . At long context the attention term stops being negligible — remember that caveat.
Put numbers in. With and tokens, FLOPs. On a setup sustaining FLOP/s — one PFLOP/s — that is seconds, about 1.3 years. One line of arithmetic explains why training runs are spread across hundreds or thousands of accelerators.
Four things that matter in practice
1. Suspect the batch size When inference is slow, look at first. Batch-1 generation is GEMV and memory bound, so adding compute does nothing. Batching is the only lever that works.
2. Track effective FLOPs over peak FLOPs Back out effective performance from measured time and divide by peak. A low ratio tells you the units are idling, and that the problem is data movement rather than arithmetic.
3. Lowering precision is buying bandwidth Moving from fp32 to fp16 or int8 does not only speed up arithmetic — it halves or quarters the bytes moved, which raises . That is why quantization pays off most in memory-bound regimes.
4. Do not write your own GEMM Tiling, prefetching and register allocation carry decades of accumulated tuning inside existing BLAS implementations. What you should be shaping is the form in which you call GEMM.
Summary
- AI compute reduces almost entirely to GEMM: FLOPs for by , or when square
- Operation count alone does not set runtime. Arithmetic intensity = FLOPs / bytes moved decides memory bound versus compute bound
- GEMM has and improves with size; GEMV is stuck at . That gap is why batching works
- Tiling raises effective intensity toward , bounded by cache. Runtime moves by an order of magnitude with the complexity unchanged
- Rules of thumb: FLOPs per token for inference, for training — enough to size a run yourself
Next: how this computation gets split across many units — the limits of parallelism and the communication cost of distributed training.
Comments
Sign in to comment