Matrix Multiplication (GEMM)

Definition

Nearly everything an AI model computes is matrix multiplication. In a single layer of Llama 3 70B, more than 99.8% of the arithmetic is matrix multiplication — the softmax, the normalisation and the activation functions, the parts with the interesting names, together come to less than a fifth of one percent. "AI compute" and "matrix-multiply throughput" have become the same phrase for a good reason. An NVIDIA H100 will perform 989 trillion floating-point operations per second on work you express as a matrix multiply, and 67 trillion on work you do not.

The operation itself you could do on paper. To multiply an m×k matrix A by a k×n matrix B, take each of the m rows of A and each of the n columns of B and compute their dot product: multiply the k pairs of numbers, add the results, and you have filled one of the m·n cells of the output. Each cell costs k multiplications and k−1 additions, so the whole multiplication costs, to a very good approximation:

2 · m · n · k floating-point operations.

That is the only formula on this page you need, and everything below is an application of it. Hardware people call the operation GEMM, for general matrix multiply — the name of the routine in BLAS, the linear-algebra interface standardised in 1990, decades before deep learning existed. The fastest code path on your GPU is named after a Fortran subroutine.

How It Works

From a prompt to a pile of matrices

A prompt is tokenized into a sequence of tokens, and each token becomes an embedding: a vector of d numbers, where d = 8,192 for Llama 3 70B. A 1,000-token prompt is therefore a 1,000 × 8,192 matrix. The model's learned weights are matrices too. Running the model on the prompt is, almost literally, multiplying the first by the second, eighty times over.

Stacking the tokens into one matrix rather than processing them one at a time is not tidy bookkeeping — it is the whole performance story. Every cell of the output is a dot product that depends on nothing any other cell is doing, so a machine with thousands of arithmetic units can compute them all at once. Matrix multiplication is embarrassingly parallel, which is why an industry could build chips that do little else.

One transformer layer, counted

Take a real model rather than a hypothetical one. Meta published Llama 3's architecture: 80 layers, model dimension 8,192, 64 attention heads of 128 dimensions each, 8 key/value heads (grouped-query attention, so K and V project down to 8 × 128 = 1,024), and a SwiGLU feed-forward network of inner width 28,672. Apply 2·m·n·k with m = 1 token, and one layer's weight multiplications come out like this:

Seven matrices, 1.71 GFLOPs per token. Attention then does two more matrix multiplications that are not against weights at all — the query against every stored key, and the resulting scores against every stored value. At an 8,192-token context that is 4 × 8,192 × 8,192 ≈ 0.27 GFLOPs, bringing the layer to roughly 1.98 GFLOPs per token.

Now price everything that is not a matrix multiply. Two RMSNorms over 8,192 numbers, rotary embeddings on the query and key vectors, a SiLU and an elementwise multiply over 28,672 numbers, a softmax over 8,192 attention scores in each of 64 heads, two residual additions: about 3 million FLOPs, against 1.98 billion. 0.15%. The layer is a matrix multiplier with some seasoning.

Where the FLOPs actually land

Two things fall out of that table, and both of them shape real systems.

The feed-forward network is the model. Its three matrices are 1.41 of the 1.71 GFLOPs spent on weights — 82% — against 18% for all four attention projections combined. Attention gets the papers; the FFN gets the silicon time. This is the entire reason mixture-of-experts works: routing each token to 2 of 8 feed-forward experts cuts into precisely the term that dominates the bill, which is why an MoE model can carry hundreds of billions of parameters and still cost like a much smaller one.

The whole model is just this, eighty times. 80 × 1.71 GFLOPs = 136.9 GFLOPs, plus a final projection from 8,192 dimensions to the 128,256-token vocabulary (2 × 8,192 × 128,256 = 2.1 GFLOPs), gives 139 GFLOPs per token. Compare that with the rule of thumb that a forward pass costs 2 FLOPs per parameter: Llama 3 70B has 70.55 billion parameters, of which 1.05 billion are the input embedding table — a lookup, not a multiply. 2 × (70.55 − 1.05) billion = 139 GFLOPs. The numbers agree exactly, because the "2 FLOPs per parameter" heuristic that FLOPs explains is not an approximation of something more complicated. It is the matmul count, wearing a disguise.

The second cost: moving the operands

A matrix multiply has two prices, not one: doing the arithmetic, and fetching the numbers to do it on. 2·m·n·k grows with the product of all three dimensions while the operands you must read grow only with their sums, so a big square matmul is a bargain and a skinny one is a disaster — and generating one token at a time makes every matmul as skinny as it gets. That asymmetry is the subject of the memory wall, and it is where the second half of every inference bill comes from.

Real-World Applications

Tensor cores: a GPU reorganised around one kernel. In 2017 NVIDIA shipped the V100 with 640 "tensor cores", units that do nothing but a 4×4×4 matrix multiply-accumulate — 64 fused multiply-adds per clock, each. The result: 125 TFLOPS of matrix throughput against 15.7 TFLOPS of ordinary FP32 on the same chip. Six years later the H100 widened the gap to 989.5 TFLOPS (dense FP16 on tensor cores) against 67 TFLOPS (FP32 on the general-purpose cores) — a factor of about 15. Read that ratio the unflattering way: code on an H100 that cannot be expressed as a matrix multiplication runs at roughly 7% of the chip's arithmetic potential. The hardware did not get faster in general. It got faster at one operation, and a transformer is 99.8% that operation.

The TPU: a systolic array with a computer attached. Google's first Tensor Processing Unit, deployed in its datacenters from 2015, is built around a 256 × 256 grid of 8-bit multiply-accumulate units — 65,536 MACs, running at 700 MHz, for a peak of 92 TOPS. Weights flow in from one edge, activations from another, and partial sums ripple through the grid without a trip to memory between steps. That single unit occupies 24% of the die. Google's stated reason for building it was blunt: multi-layer perceptrons and LSTMs — both stacks of matrix multiplies — were 90% of its datacenter inference workload, and convolutional networks, which the academic architecture community was busy accelerating, were 5%. A TPU is what happens when you take the FLOP census seriously and design the chip around the answer.

This is the payoff. Two of the most consequential pieces of silicon of the last decade — one from a graphics company, one from a search company — are, structurally, matrix multipliers with supporting cast. The chips were shaped by a single kernel. When you ask why AI hardware looks the way it does, this is the whole answer: the GPU you can buy, the software that drives it, the number formats, the memory hierarchy, all of it descends from an operation you could do on graph paper.

Key Concepts

• GEMM vs. GEMV: multiplying a matrix by another matrix (GEMM) reuses each loaded weight across every token in the batch; multiplying it by a single vector (GEMV) uses each weight exactly once. Same maths, wildly different economics — and the reason inference servers batch requests together rather than serving them one at a time.
• Accumulate wider than you multiply: tensor cores multiply in FP16 or FP8 but sum the products in FP32, because adding thousands of small products in low precision loses the small ones entirely. This is why quantizing weights to 8 or 4 bits does not destroy accuracy the way naive intuition says it should.
• Shapes are not free: the hardware multiplies in fixed tiles (a tensor core's 4×4 block, a TPU's 256×256 array). A dimension that does not divide evenly into the tile wastes the remainder of it.

Challenges

A badly shaped matrix silently costs you a quarter of your compute. When Andrej Karpathy padded nanoGPT's vocabulary from 50,257 to 50,304 — adding 47 tokens that mean nothing and are never used — training got about 25% faster, purely because 50,304 is a multiple of 64 and the matmul now landed on an efficient kernel path. Nothing about the model changed. This is the most common way a team loses performance without ever seeing an error: the shapes are legal, the results are correct, and the GPU is quietly running at half throttle.

Decoding turns every GEMM into a GEMV. Processing a 10,000-token prompt multiplies a 10,000-row matrix by the weights and saturates the tensor cores. Generating token 10,001 multiplies a single row by those same weights — the identical volume of weights read from memory to do 1/10,000th of the arithmetic. The chip is then idle almost all of the time, waiting on memory rather than maths, and no amount of extra TFLOPS helps. See the memory wall.

The FLOP map is only true for the context you counted it at. In the layer above, the attention product was 14% of the work at an 8,192-token context. But it grows linearly with context while the weight matmuls do not, so setting 4 × 8,192 × L equal to 1.71 GFLOPs gives a crossover at L ≈ 52,200 tokens: past roughly 52k of context, attention against the KV cache costs more than all seven weight matrices combined. A performance model built at 8k context will mislead you badly at 128k.

Optimising the wrong matmul. The feed-forward matrices are 82% of the weight FLOPs, so effort spent speeding up the attention projections is bounded above by an 18% win and usually delivers a fraction of that. Count first, then optimise; the census above took ten minutes.

Future Trends

The clearest direction is narrower numbers. Every halving of the operand format roughly doubles the matrix throughput of the same silicon — an H100 does 989.5 dense TFLOPS in FP16 and 1,979 in FP8 — and hardware vendors have been walking that ladder down through FP8 toward 4-bit formats, with the accumulation still done in something wider. The interesting question is no longer whether the chip can do the multiply, but how few bits the operands can survive.

The second is that alternatives to attention do not escape any of this. State-space models, linear attention and the various sub-quadratic schemes change which matrices get multiplied and how the cost scales with context; none of them replaces matrix multiplication with something else, because nothing else runs this fast on hardware built for exactly one thing. The kernel shaped the chips, and now the chips shape the architectures.