Built up from scratch, in order: what a transformer block actually does, the intuition and exact math of self-attention, how prefill and decode differ at inference time, why that forces you to keep a KV cache, and finally why almost every serving-grade model today replaces plain multi-head attention with Grouped-Query Attention (GQA) to make that cache affordable.
Self-attention lets every token pull context from every other token: it asks a question (query), the others answer with labels (keys), and it absorbs a weighted blend of their content (values). At inference, a prompt is processed once in parallel (prefill), then tokens are generated one at a time (decode) — and because decode would otherwise recompute every past key and value on every single step, you cache them instead: the KV cache. That cache is one of the two matrices per layer, per head — and it's expensive: it grows with context length and must be re-read from memory on every decode step. Grouped-Query Attention shrinks it by having several query heads share one key/value head, cutting cache size and memory-bandwidth cost by the group size, at a small, controllable cost in quality.
Before the mechanism, the shape it lives in. A transformer takes a sequence of T tokens and turns
them into a matrix — T rows, one per token, each row a vector of length dmodel
(the embedding dimension). Everything that follows just transforms that (T, dmodel)
matrix, over and over, without ever changing its shape.
X ∈ (T, dmodel).x ← x + SelfAttention(LayerNorm(x)), then x ← x + FFN(LayerNorm(x)). Pre-norm and residual connections everywhere, so gradients and information have a direct path through the whole depth.dmodel vector to vocabulary size, softmax → a probability over the next token.
A note on the two norm placements above: the figure shows post-norm — LayerNorm
after the residual add, exactly as in the original 2017 paper: x ← LayerNorm(x + Sublayer(x)).
Most modern decoder LLMs (GPT, Llama, Qwen) instead use pre-norm — normalizing before
the sub-layer, as written in step 3 above and used for the rest of this piece — because it trains more stably
at depth. Both wrap the exact same residual highway; only the order of Norm relative to the + differs.
Here's the picture worth holding onto: we have a set of tokens, and after embedding and adding positional information, each token is a vector. For every token, self-attention builds three separate projections of that vector — a query, a key, and a value.
The query is the question a token asks: "what context do I need?" The key is how every other token
advertises what it has to offer. So each token's query gets compared against every other token's key —
q · k — and that comparison becomes a weight for how much attention to pay. Then, using those
weights, the token takes a weighted sum of everyone's values and adds that blend back onto
its own embedding as a residual. That weighted blend is the extra context the token picked up by looking at
the rest of the sequence.
Query asks, keys answer with a match score, softmax turns scores into weights, and the token absorbs a weighted sum of everyone's values as its update.
Now the same story, formalized term by term. Dimensions used below, matching a small real config:
T = sequence length · dmodel = 1024 ·
hq = 8 query heads · dk = dv = 128
(so hq · dk = dmodel). We'll do one head first,
then stack heads in §4.
1 · Three projections per token. Learned weight matrices turn each token's embedding into a query, key, and value:
Row t of Q is token t's question; row t of K is what token t advertises; row t of V is what token t actually hands over once picked.
2 · Every query against every key, at once. A single matrix multiply computes all pairwise qt · ki dot products in one shot:
This is the T×T object that makes attention O(T²): row t holds token t's query compared against every key. Entry (t, i) is literally qt · ki, exactly the "query asks, key answers" comparison from the intuition above.
3 · Scale by √dk. If each component of q and k is roughly unit-variance and independent, the dot product q · k (a sum of dk such terms) has variance proportional to dk. Bigger heads → bigger raw scores → softmax saturates into a near one-hot distribution and gradients vanish. Dividing by √dk cancels that growth and keeps the scale roughly constant regardless of head size:
There's an interactive for exactly this effect below — it's not a minor detail, it's the difference between attention training stably or not.
4 · Causal mask, then softmax per row. For language modeling, token t may only look at tokens ≤ t (it can't peek at the future it's trying to predict), so entries above the diagonal are set to −∞ before the softmax zeroes them out. Softmax is applied row-wise so each token's weights sum to 1:
5 · Weighted sum of values — the payoff. Multiply the weights by V and every token's output is exactly the weighted blend the intuition promised:
Put together, this is the equation of attention:
Finally the output is projected back and added to the residual stream: x ← x + O·WO, with WO ∈ (128, 1024) mapping the head back to model width (this matters more once there are several heads to recombine, next section).
Fix five keys with a spread of similarity to one query (from very aligned to very opposed), and grow the head
dimension dk. Raw dot-product scores grow with √dk; scaling by
1/√dk should hold the softmax steady no matter how wide the head is. Drag the slider.
One head can only learn one notion of "relevant." Multi-head attention runs hq of these
in parallel, each with its own WQ, WK, WV, each looking at a
different dk = dmodel/hq-wide slice of representation space, then
concatenates the hq outputs and projects back down:
Every head has its own keys and values here — that's Multi-Head Attention (MHA),
the original 2017 design and still the highest-quality option. Before asking what it costs, there's a detail
the equation hides completely: X isn't always the same shape. At inference time it's either a
whole prompt or a single new token — and that distinction changes everything about how expensive this equation
is to run.
Training sees a whole sequence at once and predicts every next-token in parallel. Inference is different: you're generating text you don't have yet, one token at a time. That splits a request into two very different regimes.
Toggle between the two phases below. Same equation, same weights — but the steps that make it up, and which
ones dominate the wall-clock time, are different. Each step below is tagged with its actual latency: GPT-2
Medium (dmodel=1024, 24 layers, 16 heads of 64 dims) on an A100 80GB SXM, prefilling
T=1,024 tokens or decoding one token at context length t=1,024 — memory-tagged
steps read weights off HBM, compute-tagged steps are matmul arithmetic
(derivations in the appendix; numbers are per layer — ×24 for the full model):
X ∈ (T, d_model) in one batch.(dmodel×dmodel) weight matrices — ~2.1 MB each in bf16 — stream in once, shared by every one of the T tokens about to use them.2·T·dmodel² FLOPs each — GPU-friendly, fully parallel across tokens. The more tokens T share that one weight read, the more it pays for itself.O(T²·dk) FLOPs per head. Every token attends to everything ≤ its position, all computed together. Pure activation arithmetic, no weight traffic.O(T²·dk) matmul, same story as step 4 — activations in, activations out.
Sum the modeled steps per layer and the split is telling: prefill spends
≈45.4 µs/layer, almost all of it in the matmuls (steps 3–6); a vanilla decode spends
≈18 µs/layer, and — strikingly — almost all of that is step 3, re-projecting K, V
for the whole context, the same work redone every step. The two weight reads (steps 2 and 5, ≈4.1 µs together)
and the new-token math are tiny beside it. Steps 2 and 6 in prefill (2 and 5 in decode) are
memory operations — reading weights off HBM; everything else is arithmetic.
Whether memory or arithmetic dominates comes down to one ratio — arithmetic intensity, FLOPs
done per byte read. For a single (dmodel×dmodel) projection over
T tokens:
A GPU has its own fixed ratio — peak FLOPs/s divided by HBM bandwidth — sometimes called its ridge point.
Roughly ~150 FLOPs/byte on an A100, ~300 on an H100 (bf16). Compare that constant to T above:
During prefill, arithmetic intensity blows past the ridge point — the matmuls (steps 3, 5, 6) are what take the
time, and the GPU is busy doing math, not waiting on HBM. Decode should be the opposite: only one
genuinely new token enters per step, so its own work is AI ≈ 1, far below the ridge point —
memory-bound, with the weight reads (steps 2, 5) as the real cost and the new-token math finishing instantly.
But vanilla decode never gets there, because step 3 drags all t old tokens back through
WK, WV every step — an AI ≈ t recompute that puts decode right
back in the compute-bound regime, doing prefill-scale math to reproduce values that never change.
(Real numbers for a concrete model and GPU are in the appendix.)
So the open problem isn't step 4's arithmetic — it's step 3's recompute: rebuilding all t keys and
values from scratch every step, even though nothing about the old tokens changed. Storing them instead of
recomputing — and what that storage costs to keep and to read — is next.
Decode step 4 needs K, V for every one of the t tokens generated so far, not just the
new one. There are exactly two ways to get them: recompute them from scratch every step, or
keep them around from when they were first computed. The difference between those two options
is the entire reason the KV cache exists.
Recomputing is exactly what the vanilla decode in the last section does: step 3 re-projects K,V for all
t tokens at every single decode step — effectively a fresh mini-prefill of length t,
every step. That's
2·t·dmodel² FLOPs for the K,V projections alone at step t, and summing
that over an N-token generation costs O(N²) instead of O(N). Worse,
decode is already memory-bandwidth-bound (previous section) — burning compute on work whose answer
never changes is pure waste layered on top of an already slow step. Concretely, on the same GPT-2
Medium / A100 setup, re-projecting K,V for t=1,024 prior tokens costs ≈13.8 µs/layer —
every single decode step, forever.
Option B is what every real system does: keep the K and V matrices from previous steps around instead of recomputing them — the KV cache. Concretely, per layer, per head, the cache is two growing matrices:
Two things happen. First, step 3 collapses: with the old K,V already saved, there's nothing to
re-project, so it drops from the t-token recompute (≈13.8 µs) back to just the one new row
(≈0.02 µs). Second, decode step 4 splits into two cheap moves instead of one expensive recompute (same setup,
still t=1,024):
qnew · Kᵀ against every cached key, softmax, weighted sum of every cached value — no projection recomputation at all, just a read of t rows off HBM (≈2.06 µs) plus the same tiny attention arithmetic from step 4 (≈0.01 µs).
That gap isn't fixed — it widens every subsequent token, since recompute's cost grows linearly with t
while the cache's new-token projection (step 3) stays flat at O(1); only its cheap read (step 4b) grows with
t. That's the O(N²)-vs-O(N) difference from a moment ago, in real microseconds (full derivation in
the appendix). Across the whole step, that drops a decode layer from the vanilla ≈18 µs to ≈6.2 µs,
and — just as importantly — moves the bottleneck off the recompute matmul and onto memory traffic, restoring the
AI ≈ 1, memory-bound decode the previous section said we should have. Prefill changes too,
in the mirror-image way: its step 3 (project Q,K,V for all T positions) now
also writes all T of those K,V rows into the cache in one shot — that's the entire point of
prefill from the cache's perspective, seeding it before decode ever runs.
Total cache memory, summed over every layer and every head, in bytes:
Against the O(N²) recompute cost above, reading this cache is O(N) over a generation
— each step reads, at most, t rows once, never redoing earlier work. With plain MHA,
hkv = hq — every query head drags its own full K,V along — and that
O(N) win comes with two costs of its own:
The cache-size formula above has exactly one term you can shrink without touching sequence length or model
width: hkv. The question is how many query heads should share one K/V head — and that
single design choice spans a whole family of attention variants.
hkv = hq. Every query head gets its own key/value head. Best quality, largest cache.hkv = 1. Every query head shares a single key/value head. Smallest possible cache, but one K/V representation has to serve all query heads at once, which can hurt quality (Shazeer, 2019).1 < hkv < hq. Query heads are split into hkv groups of hq/hkv heads each; every head in a group reads the same K, V. A tunable interpolation — e.g. Llama 2 70B uses 8 KV heads for 64 query heads (8:1); the Qwen3.5 config discussed in the companion piece uses 4:1 (16 query heads, 4 KV heads).The math is a one-line change to multi-head attention: query head i no longer uses its own Ki, Vi — it looks up which group it's in and uses that group's shared K, V:
Nothing about how many query heads exist changes — GQA never touches representational capacity on the
query side, which is a big part of why it recovers so much of MHA's quality. What changes is narrower — same
GPT-2 Medium / A100 setup as before, now with hkv=4 (a 4:1 ratio, matching the
Qwen3.5 example earlier):
WK and WV shrink from (dmodel×dmodel) to (dmodel×hkvdk): 3.1 µs → 1.5 µs per layer. WQ is untouched, so the read doesn't shrink by the full 4× — only the K,V portion does.hkv physical rows to write and read instead of hq: 2.1 µs → 0.53 µs per layer, close to the full 4× since this step is almost entirely the cache read.
Everything else — the scores, the scaling, the causal mask, the softmax, the weighted sum (§3) — is
identical inside each head, and costs exactly what it did under MHA (steps 4/5 in prefill, the
attend-arithmetic half of step 4 in decode): those still scan all hq=16 query heads,
GQA or not.
Both weight and cache savings scale with the same ratio, hq/hkv — but neither
hits that full ratio in the per-token total, because steps that don't touch K,V (Q's projection, the attention
arithmetic, WO) don't shrink at all:
Drag the slider below to see the actual head wiring (full arithmetic for every number above is in the appendix):
Sharing K,V across a group means those query heads can no longer attend to independently chosen positions within the group — they're all scored against the same keys. In practice this costs surprisingly little: with a moderate group size (4–8), GQA lands within a fraction of a point of MHA on downstream benchmarks while cutting the cache by the same factor, which is why it — or MQA outright — is the default in nearly every serving-oriented model released since 2023.
Optional aside — none of the reasoning in §4 depends on this, but it's satisfying to see real numbers attached
to the ridge-point argument. GPT-2 Medium has dmodel=1024, 24 layers, 16 heads of 64
dims each, ~355M parameters — same dmodel as the running example throughout this piece,
so the weight-matrix math carries over unchanged. Run it on an NVIDIA A100 80GB SXM: 312 TFLOPS bf16
tensor-core compute, 2,039 GB/s HBM bandwidth (both published spec-sheet numbers) — a ridge point of
312e12 / 2.039e12 ≈ 153 FLOPs/byte, right where the earlier ~150 estimate landed.
Two regimes, same weights, same GPU, wildly different bottleneck — the QKVO projection matmuls alone, per layer:
Sum those over all 24 layers and the split gets stark: decoding one token spends ≈99 µs just streaming Q/K/V/O weights off HBM, while the matmul arithmetic for that same step finishes in well under a microsecond. Prefilling a 1,024-token prompt spends ≈661 µs on QKVO matmul arithmetic across those 24 layers, versus only ≈99 µs reading the identical weights once for the whole batch. Same model, same weights, same GPU — the bottleneck simply flips depending on how many tokens ride along with each weight read. That's the QKVO projection matmuls only (steps 2, 3, 6 in §4's prefill list); the rest of this appendix fills in the remaining steps — attention itself, the KV cache, and GQA — to get the full per-step numbers shown throughout the piece.
Unlike the QKVO projections, these two steps touch no weight matrices at all — just activations already in
SRAM — so their cost is pure FLOPs, summed across all hq=16 heads:
Add those two steps to the QKVO projection numbers above and the full per-layer prefill cost (steps
2–6) is 3.1 + 20.7 + 6.9 + 6.9 + 7.9 ≈ 45.4 µs, ×24 layers ≈ 1.09 ms for the entire
1,024-token prompt — this is the number the piece calls the full prefill total, larger than the ≈661 µs
projection-only figure above because it also counts the T² attention arithmetic.
At T=1 the QKᵀ / weighted-sum arithmetic above shrinks to almost nothing —
4 hq t dk FLOPs (one query row against t keys, both directions),
which at t=1,024 is only 4·16·1024·64 ≈ 4.19M FLOPs ≈ 0.013 µs per layer. What
actually costs time is fetching the t rows of K,V those FLOPs run against:
Compare that to recomputing K,V for all 1,024 prior tokens instead of reading them back:
Full decode total (steps 2, 3, 4b, 5) with a cache: 3.1 + 0.02 + 2.1 + 1.0 ≈ 6.2 µs/layer, ×24 ≈
149 µs per generated token at this context length — versus 3.1 + 0.02 + 13.8 + 1.0 ≈ 17.9
µs/layer, ×24 ≈ 430 µs, if every step recomputed K,V from scratch instead.
Only WK, WV, and the cache shrink — they go from
hq=16 physical heads to hkv=4, a 4× reduction.
WQ, WO, and the attention arithmetic itself (steps 4/5 in
prefill, the compute half of step 4 in decode) are untouched, since the query side and the score/softmax/weighted-sum
math still run over all 16 query heads:
Step 3 (project Q,K,V) also drops a little, since it's now projecting into a narrower K,V output:
FLOPs = 2T dmodel(dmodel + 2hkvdk) instead of
2T · 3dmodel², which at T=1,024 works out to ≈10.3 µs (prefill) instead
of 20.7 µs. Rolling all of that up:
| Step | MHA | GQA (4:1) |
|---|---|---|
| Prefill, per layer (steps 2–6) | 45.4 µs | 33.5 µs |
| Prefill, full 24 layers | ≈1.09 ms | ≈0.80 ms |
| Decode, per layer (steps 2,3,4b,5) | 6.2 µs | 3.1 µs |
| Decode, full 24 layers | ≈149 µs | ≈75 µs |
The full-model ratios (~1.4× prefill, ~2.0× decode) land well under the raw hq/hkv=4×
cache-size reduction from §5 — because WQ, WO, and the attention
arithmetic dilute the win. The steps that scale purely with hkv (the cache read, step
4b) get close to the full 4×; the steps that don't (everything touching Q or O) don't move at all.