Matrix Multiplication: The Core Operation

Strip away the vocabulary of transformers, attention, feed-forward, projections, and what remains underneath is almost entirely one operation performed billions of times: matrix multiplication. Every linear layer, every query-key-value projection, every attention score, every output logit is a matmul or a small composition of them. If you understand the shape rule for matrix multiplication and its FLOP cost, you understand where a transformer spends essentially all of its compute.

The operation and its shape rule

For a matrix A of shape (m, k) and B of shape (k, n), the product C = A @ B has shape (m, n), and each entry is a dot product of a row of A with a column of B:

C[i, j] = sum_k  A[i, k] * B[k, j]

The rule that matters in practice: the inner dimensions must match (A's columns equal B's rows), and that shared dimension disappears from the output while the two outer dimensions survive. In a transformer's linear layers, the "shared dimension" is d_model or the hidden width, and it is the axis being contracted, summed over and discarded, on every forward pass. When code raises a dimension-mismatch error, it is this rule being violated.

Why it dominates the compute budget

A matmul between an (m, k) and (k, n) matrix costs 2 * m * k * n floating-point operations (one multiply and one add per term in the sum, for each of the m * n output entries). For a transformer, the dominant matmuls are the linear projections in attention and the feed-forward block, each roughly O(seq_len * d_model^2) per layer. This is the basis of the widely used approximate rule that a transformer forward pass costs about 2 * N FLOPs per token, where N is the parameter count, because almost every parameter participates in exactly one multiply-add per token passed through it (Kaplan et al., 2020, arXiv:2001.08361). Attention itself contributes a separate O(seq_len^2 * d_model) term that this per-parameter approximation does not capture, which is why it becomes the dominant cost only at long context lengths (see context-windows-long-context).

Batched matmul

Transformers do not run a single matmul at a time; they run a batch of them, one per sequence, and the framework fuses these into a single batched matmul call: (batch, m, k) @ (batch, k, n) -> (batch, m, n), looping over the batch dimension internally while sharing kernel launch overhead. GPUs are matrix-multiply accelerators first and everything else second: tensor cores exist specifically to execute matmuls at far higher throughput than general arithmetic, which is why frameworks go to considerable lengths to express as much of a model's computation as possible as large, well-shaped matmuls rather than many small elementwise operations.

scores = Q @ K.transpose(-2, -1)   # (batch, heads, seq, d_head) @ (batch, heads, d_head, seq)
                                     # -> (batch, heads, seq, seq)

When it falls down

• Small matmuls waste hardware. A GPU's tensor cores need matrices large enough to saturate their pipelines; a matmul with a tiny inner dimension or batch size runs at a fraction of peak throughput, which is one reason batch size and hidden width are tuned with hardware utilisation in mind, not just statistical considerations.
• FLOPs are not wall-clock time. A matmul-heavy model can still be memory-bandwidth bound if activations do not fit in fast on-chip memory, so counting FLOPs alone overstates how "expensive" an operation actually is; see the arithmetic-intensity framing in flash-attention.
• Precision changes both cost and correctness. Lower-precision matmuls (bf16, fp8) run faster on modern accelerators but accumulate rounding error differently, so numerically sensitive contractions (like softmax normalisers) are often kept in higher precision even inside an otherwise low-precision matmul.
• Associativity is not guaranteed in floating point. (A @ B) @ C and A @ (B @ C) are mathematically equal but can produce slightly different floating-point results, which matters for reproducibility across hardware and library versions.

Further reading

• Golub and Van Loan, Matrix Computations - the standard reference for matrix multiplication and numerical linear algebra.
• Kaplan et al., 2020, Scaling Laws for Neural Language Models, arXiv:2001.08361 - the 2N FLOPs-per-token approximation derived from matmul cost.