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.

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 $O(n^3)$

Multiply an $m \times k$ matrix $A$ by a $k \times n$ matrix $B$ to get an $m \times n$ matrix $C$. One output cell is defined as

Read in words: "row $i$ of $A$ against column $j$ of $B$, multiplied elementwise and summed" — a single dot product. There are $mn$ cells, each costing $k$ multiplications and roughly $k$ additions. Counting a multiply and an add as one operation each,

$$\text{FLOPs} = 2\,m\,k\,n$$

which says nothing more exotic than: cells in the output ($mn$), times multiply-adds needed to fill one cell ($k$), times two operations per multiply-add.

That is $2n^3$ for square matrices. Double $n$ and the work grows eightfold. That is the content of $O(n^3)$, and it is also what people mean by a model's "compute".

Numbers give it scale. A square product at $n = 4096$ costs about $1.4 \times 10^{11}$ 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 $2n^3$ 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 $I$ 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 $I$ stalls on bandwidth no matter how fast the arithmetic units are (memory bound). With high $I$ you get to use the units at full tilt (compute bound). Achievable performance is capped by

$$\text{performance} \le \min\big(\text{peak compute},\ \ I \times \text{memory bandwidth}\big)$$

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.