This started as a set of annotated notes on Google DeepMind's "How To Scale Your Model", and grew into an evergreen reference of my own, now with interactive visualizations built in.
A dot product of two length-P vectors costs 2P floating point operations: P multiplies and P adds. Below, x and y are the inputs and r is the result: tap it to see exactly which named cells of x and y it came from.
Stretch that into a matrix-vector product and you're doing one dot product per row, so an N×P matrix A against a length-P vector x costs 2NP. Tap any cell of the result y and you'll see it lights up one row of A and all of x, because that row is the only thing that fed it.
Stretch again into a full matrix-matrix product, A[N,P]·B[P,M], and now the result C is a whole grid: N rows × M cols, count them yourself: that's 2NPM FLOPs total (N×M cells, each one a 2P-FLOP dot product). Tap any cell of C and watch one row of A and one column of B light up; those are the only two things that determined it.
The general version of this rule is what actually matters for reading the accounting tables later, and it's worth being precise about what contracting and batching axes actually change, because the natural assumption is that they must affect the FLOP count differently. They don't.
Start from the matmul above: A[N,P]·B[P,M]→C[N,M], costing 2NPM. N, P, and M each appear exactly once in that formula: none of them is "extra expensive" or "extra cheap." P happens to get summed away and disappear from the output C; N and M happen to survive into it. That's the only difference between them, not how much they cost.
Now run G independent copies of that exact matmul side by side, G attention heads, say, each with its own A and B:
G shows up exactly once too, multiplying the total exactly the way P did: batching and contracting axes cost identically, per axis. The one real difference: G survives into the output (C is now shaped [G,N,M] instead of [N,M]), while P still doesn't. Whether an axis disappears from the output or not is the entire distinction between them; it has no bearing at all on how the FLOP count scales.
So the fully general rule really is as simple as it sounds: FLOPs = 2 × every axis that appears in either input, each counted once, full stop, no exception for which kind of axis it is. The categories only tell you which axes survive into the output shape:
One consequence is easy to miss and explains a lot about why Transformers scale so well: for a square matmul, compute grows as N³ while the data you have to move only grows as N². Bigger matmuls become relatively cheaper to feed, which is exactly why architectures built almost entirely out of matmuls are the ones that scale.
During training you don't just run the matmul forward; you also need its gradient, and that gradient is itself two matmuls. If C = A·B, the chain rule gives you ∂L/∂B = Aᵀ·(∂L/∂C) and ∂L/∂A = (∂L/∂C)·Bᵀ, each one another 2NPM-FLOP contraction.
That's the origin of the shorthand you'll see everywhere in scaling-law papers: training FLOPs ≈ 6 × parameters × tokens. It's not a rule about Transformers specifically; it falls straight out of how backprop works for any stack of matmuls.
This section walks through each block's FLOP and parameter accounting right alongside its diagram, with batch B, sequence length T, model width D, MLP width F, N query heads, K key/value heads, and head dimension H doing the bookkeeping throughout.
In its traditional form this block is about as simple as a neural-net layer gets: project up from D to a wider hidden size F, apply a nonlinearity, then project back down to D.
Most modern implementations (LLaMA, DeepSeek, Gemma, and most other current open models) extend this with a gating einsum: a second, parallel up-projection that gets multiplied elementwise with the first before the same down-projection runs. That's 3 weight matrices of size D·F instead of the 2 above: the extra matrix buys a learned gate at the cost of 50% more MLP parameters.
Putting numbers on the gating version above (the 3-matrix one, which is what's actually running in most current models):
| operation | train FLOPs | params |
|---|---|---|
| x·W_in1 (D→F) | 6BTDF | DF |
| x·W_in2 (D→F) | 6BTDF | DF |
| g·W_out (F→D) | 6BTDF | DF |
| total | ≈18BTDF | 3DF |
This block runs on the same D-wide residual stream as the MLP above. Attention's whole trick is splitting that D-wide vector into several smaller "heads" it can attend with independently, then stitching them back together. That introduces a few new letters: N is the number of query heads, K is the number of key/value heads, and H is the size of each head, so N·H = D. Two more track position counts: T is the number of query positions and S is the number of key positions (usually the same sequence, so T = S).
Four weight matrices move x in and out of that per-head shape: Q, K, and V each project the D-wide input down into several H-wide heads (N of them for Q, K of them for K and V), and the output projection O puts all the heads back together into D.
It's easy to read "projections" and "dot‑product attention" as two names for the same thing: both are matmuls, and §01 already showed every matmul is nothing but dot products underneath. The real difference is what's on each side of the multiply, and it's worth being clear about before going further, because the accounting tables below treat these two as separate rows for exactly this reason:
Why does having two sequence-length sides turn linear scaling into quadratic? Because, as §01 already established, a matmul's cost is just (rows × cols) output cells, each one a fixed-size dot product. In a projection, only the row count (number of tokens) depends on sequence length; the column count is a fixed hyperparameter (N·H or D). In attention, both the row count (T query positions) and the column count (S key positions) depend on sequence length, so doubling the sequence length doesn't just double one side of the output grid, it doubles both sides at once, and the grid's area goes up 4×, not 2×. Scale tokens by 10× and a projection costs 10× more; attention costs 100× more.
That per-head computation, the head_i line in the equation above, is exactly what the diagram below walks through:
Once every head has produced its own [T,H] output, all of them get concatenated back into [T, N·H] and pass through one more projection, W_O, back down to D: the fourth matmul in the accounting table below. The dot‑product attention step shown above is the one that scales with sequence length squared rather than with model width, which is exactly why the "when does attention actually matter?" section below treats it separately from the four QKVO projection matmuls.
Putting numbers on the four projections first:
| operation | train FLOPs | params |
|---|---|---|
| x·W_Q (D→N·H) | 6BTDNH | DNH |
| x·W_K (D→K·H) | 6BTDKH | DKH |
| x·W_V (D→K·H) | 6BTDKH | DKH |
| concat·W_O (N·H→D) | 6BTDNH | DNH |
| total | 12BTD(N+K)H | 2D(N+K)H |
And now the dot‑product attention step itself: the one with no weight matrix, from the diagram above:
| operation | train FLOPs |
|---|---|
| Q·Kᵀ | 6BT²NH |
| softmax·V | 6BT²NH |
| total | ≈12BT²NH |
Everything above holds for any split between N query heads and K key/value heads: that split is exactly what separates the three attention variants named everywhere.
Nothing about the computation itself changes: same equation, same per-head softmax(Q·Kᵀ/√H)·V from above, only how many distinct K and V arrays feed it. Tap through the three modes below to see the actual head-to-head assignment, not just the head counts:
Whichever mode you're in, every query head still runs its own attention computation. GQA and MQA don't reduce how many times the QKᵀ·V sequence runs, only how many distinct K and V arrays feed it, which is exactly what the accounting above already priced in via K. The real payoff shows up later, in how much has to be cached at inference time: see §08, KV caching, below.
You've already seen every number in this section. The MLP total (≈18BTDF) and the attention block's two totals (12BTD(N+K)H for the projections, ≈12BT²NH for dot‑product attention itself) were worked out alongside their diagrams in §03. All that's left is adding them up:
Add the MLP and the QKVO projections together (leaving dot‑product attention out for a moment) and the per-layer, per-token FLOPs collapse to something clean:
Notice what actually did the work in step 2: it wasn't anything specific to the MLP or to attention. It's §02's fact that every projection matmul (activations times a fixed, learned weight matrix) costs 6× its own tokens×params. Add up FLOPs from any collection of projection matmuls and you're automatically adding up 6×tokens×(their combined params), regardless of which projections they are. That's the whole reason FLOPs ≈ 6·N_params·N_tokens works as a rule of thumb across entire models, not just single layers.
This is the whole justification for the rule of thumb you'll see quoted everywhere: FLOPs ≈ 6 · N_params · N_tokens. It's an approximation that quietly assumes the T²-scaling attention term is small next to everything else, which is true until it isn't.
Putting the two FLOP terms side by side (attention itself vs. everything else) and simplifying with the typical assumptions F=4D and D=NH gives a strikingly simple crossover:
So dot‑product attention only starts to dominate once sequence length exceeds roughly 8× the model width. For a D≈8k model that's around 64K tokens of context; for a smaller model like a 4.6k-width Gemma variant, the crossover arrives much earlier, near 37K tokens. This is why the "quadratic attention is a scaling disaster" intuition is misleading for large models but genuinely bites for small ones at long context. It's also why techniques like Flash Attention and local/sliding-window attention exist to push that crossover further out.
An MoE layer is, to a rough first approximation, an ordinary dense layer duplicated into E independent MLP "experts," with a router sending each token to only k of them. The ratio E/k is the sparsity, typically 8–64× (DeepSeek-V3, for instance, routes each token to 8 of 256 experts).
Nothing about the block itself is new: an MoE layer really is just E copies of the exact 18BTDF / 3DF block from §03, sitting side by side. Compute scales with k because that's how many of those copies any single token actually runs through; parameters scale with E because all of them have to be stored whether or not a given token uses them. Divide one scaling by the other and you get E/k back: DeepSeek-V3's E=256, k=8 means 256× the parameter capacity of a single dense MLP for only 8× its compute.
That buys you a much larger parameter count for a nearly fixed compute budget per token. The bill comes due in communication instead of FLOPs: routing tokens to their assigned experts and back requires an extra pair of all-to-all exchanges per MoE layer that a dense model never needs.
Backprop needs the intermediate activations from the forward pass to compute gradients: every input to every nonlinearity, in principle. Saving all of them is what turns backprop's compute cost from quadratic-in-layers into linear-in-layers, but the memory bill is brutal: a model with 4M tokens per batch, 64 layers, and D=8192 would need on the order of ~84TB of activations in bf16 if it saved everything.
Two policies bracket the tradeoff space:
Block remat: save only each layer's input, nothing else. Most aggressive, cuts memory to roughly 1/20th, but the backward pass has to redo essentially the entire forward computation to reconstruct everything else, pushing training FLOPs from the usual 6·params·tokens up to roughly 8·params·tokens. That's a real compute tax paid specifically for the memory you saved.
Big-matmuls-only: save just the outputs of the large matmuls (the ones that are actually expensive to recompute) and let cheaper activation functions and parts of attention be recomputed. This lands around 7 saved tensors per layer instead of 20, at a much smaller compute penalty than block remat, because what's cheap to recompute (elementwise ops) gets recomputed and what's expensive (matmuls) gets kept.
Inference splits into two phases with very different cost profiles: prefill, which processes the prompt in one shot and stores the resulting key/value projections, and generation, which reuses that cache to sample one token at a time without recomputing attention over the whole prefix.
Plug in realistic numbers (8K context, 64 layers, K·H=D=8192) and you land around 8 GiB per sequence. That single formula is the entire reason grouped-query attention exists: K (the number of KV heads) is the one lever in this equation that shrinks the cache without touching model width D, which is why production inference systems push K far below N.