Transformers, from attention to KV cache

The problem attention solves

Before transformers, sequence models processed tokens one at a time, carrying a hidden state forward. That design has two defects that turned out to be fatal.

It is inherently sequential. You cannot compute step 50 until you have computed step 49, so training cannot be parallelised across the sequence. On modern accelerators, which are enormous parallel arithmetic engines, this leaves most of the hardware idle.

Information has to survive the journey. For token 200 to use something from token 3, that information must be carried through 197 successive state updates, each of which can overwrite it. Gated architectures mitigated this and did not solve it.

Attention's answer is direct: let every position look at every other position, in one operation, in parallel. Token 200 reads token 3 directly. Nothing has to survive anything.

The price is that "every position looks at every other" means n² comparisons, and that quadratic term is the source of most of what follows in this article.

Scaled dot-product attention

Each token produces three vectors by multiplying its embedding by three learned matrices:

Query — what this token is looking for. Key — what this token offers. Value — what this token contributes if selected.

The database analogy is worth keeping, because it is genuinely accurate: you match queries against keys to decide how much of each value to retrieve. The difference is that the match is soft — every key matches to some degree — which is precisely what makes it differentiable and therefore learnable.

# x: [seq, d_model]
Q = x @ W_q          # [seq, d_k]
K = x @ W_k          # [seq, d_k]
V = x @ W_v          # [seq, d_v]

scores  = Q @ K.T / sqrt(d_k)        # [seq, seq]
scores  = scores + causal_mask       # -inf above the diagonal
weights = softmax(scores, axis=-1)   # rows sum to 1
out     = weights @ V                # [seq, d_v]

Four lines. That is the mechanism in its entirety.

Why divide by the square root of d_k

This is the detail interviews ask about, and the answer is a variance argument rather than a convention.

Take two vectors whose components are independent with mean 0 and variance 1. Their dot product is a sum of d_k such products, so it has variance d_k and typical magnitude proportional to √d_k. With a head dimension of 64, scores routinely land around ±8 before scaling.

Feed scores of that magnitude into softmax and it saturates: one entry goes to nearly 1, the rest to nearly 0. The gradient of softmax in that regime is approximately zero, so nothing learns. Dividing by √d_k restores unit variance regardless of head dimension and keeps softmax in a region where gradients flow.

The causal mask

For a language model, token 5 must not see tokens 6 onward — otherwise the training objective is trivial, since the answer is in the input. The mask sets the upper triangle of the score matrix to negative infinity before softmax, which makes those weights exactly zero afterwards.

Using negative infinity rather than deleting entries is what keeps the operation a single dense matrix multiply, which is the whole point on parallel hardware. It is also why the compute cost is n² even though half the matrix is discarded.

Multi-head attention

One attention operation produces one weighted average per token, which forces a single notion of relevance. But the relationships in language are plural: a token relates to its syntactic head, to the entity it refers to, to the topic of the paragraph, and to the token that immediately precedes it, all at once. One averaging operation cannot represent all of those.

Multi-head attention runs h attention operations in parallel with separate learned projections, then concatenates the results and mixes them with a final output projection.

# d_model = 512, h = 8  ->  d_head = 64
# Project once, then reshape into heads — one matmul, not eight.
Q = (x @ W_q).reshape(seq, h, d_head).transpose(1, 0, 2)
K = (x @ W_k).reshape(seq, h, d_head).transpose(1, 0, 2)
V = (x @ W_v).reshape(seq, h, d_head).transpose(1, 0, 2)

# attention per head, batched: [h, seq, seq] @ [h, seq, d_head]
heads = attention(Q, K, V)                     # [h, seq, d_head]
out   = heads.transpose(1, 0, 2).reshape(seq, d_model) @ W_o

The critical detail is that d_head = d_model / h. Eight heads of 64 dimensions cost the same as one head of 512 — multi-head attention is not more expensive, it is the same compute partitioned differently. That is why it was free to adopt.

The output projection W_o is not decoration. Without it the concatenated heads are just stacked side by side with no mixing, and each output dimension would depend on exactly one head.

The block, end to end

A transformer block is attention plus a feed-forward network, each wrapped in a residual connection and a normalisation layer. Modern models place the norm before each sublayer rather than after:

# Pre-norm block. The residual stream is never normalised —
# it runs clean from input to output, which is what lets
# gradients reach early layers in a very deep stack.
x = x + attention(norm(x))
x = x + feed_forward(norm(x))

The feed-forward network is two linear layers with a non-linearity between them, expanding to roughly four times d_model and back:

def feed_forward(x):                # d_ff is typically 4 * d_model
    return (gelu(x @ W_1 + b_1)) @ W_2 + b_2

It is worth being clear about the division of labour, because it is frequently misstated. Attention moves information between positions. The feed-forward network processes each position independently. The feed-forward network sees one token at a time and has no idea the others exist. All cross-token communication happens in attention, and nowhere else.

Why pre-norm rather than post-norm

The original architecture applied normalisation after the residual addition. That works, but it requires a carefully tuned learning-rate warmup to train deep models without diverging — the residual path passes through a normalisation layer, so the gradient gets rescaled at every block.

Pre-norm normalises the input to each sublayer and leaves the residual stream untouched. The gradient now has a clean identity path from the output all the way to the first layer, and training becomes markedly more stable at depth. The trade is slightly worse final quality at equal depth in some settings, which is why the architecture is universal at large scale and not universal at small.

Most recent large models also use RMSNorm rather than LayerNorm — it drops the mean-centring and keeps only the scaling, which is cheaper and empirically works as well.

Position, and why RoPE won

Attention as described is permutation-invariant. Shuffle the input tokens and the set of outputs is the same set, reordered. "Dog bites man" and "man bites dog" are indistinguishable. Something has to inject order.

Absolute learned embeddings. Add a learned vector per position. Simple, effective, and fundamentally limited: position 5000 has no embedding if you trained to 2048, so the model cannot extrapolate at all.

Sinusoidal encodings. Add fixed sinusoids of varying frequency. Deterministic and defined at any position, so they extrapolate in principle. In practice quality still degrades well beyond the training length.

Rotary position embeddings (RoPE). Do not add anything. Instead, rotate the query and key vectors by an angle proportional to their position, treating consecutive pairs of dimensions as coordinates in a plane.

The reason this is the right idea is a small piece of algebra. Attention depends on q·k. If you rotate q by angle mθ and k by angle nθ, then the dot product of the rotated vectors depends only on (m − n) — the difference in position. Absolute position vanishes and relative position remains.

That is exactly the property language needs. The relationship between adjacent tokens is the same relationship whether they sit at positions 5 and 6 or 5005 and 5006, and RoPE encodes it identically in both cases. Different dimension pairs use different frequencies, so some encode fine local structure and others encode coarse long-range structure.

It also degrades more gracefully past the trained length, and — more usefully — its structure is explicit enough to manipulate. Position interpolation rescales the rotation frequencies so a longer sequence maps back into the trained angular range, and frequency-aware variants scale the high- and low-frequency bands differently on the reasoning that local detail should be preserved while long-range resolution is stretched. Both usually need a short fine-tune to recover quality, but they work, and they are the reason context windows have been extended after training rather than only by retraining.

The KV cache

Generation is autoregressive: produce a token, append it, produce the next. Done naively, each step re-runs the whole model over the whole sequence, so generating n tokens costs O(n²) forward passes worth of work — and in practice, an unusable amount.

The saving observation is that the keys and values of previous tokens never change. Token 3's key was computed from token 3's embedding and the weights; nothing about appending token 47 alters it. So compute them once and keep them.

# Prefill: process the whole prompt once, keep K and V.
K_cache, V_cache = project_kv(prompt)     # [layers, heads, n_prompt, d_head]

# Decode: each new token contributes one row to the cache
# and attends over everything accumulated so far.
for step in range(max_new_tokens):
    q, k, v = project_qkv(last_token)     # one position only
    K_cache = concat(K_cache, k)
    V_cache = concat(V_cache, v)
    logits  = attend(q, K_cache, V_cache)
    last_token = sample(logits)

This turns generation from quadratic into linear, and it is why interactive language models are possible at all.

It also relocates the bottleneck

The cache is large, and its size is the constraint that governs serving economics:

kv_bytes = 2                 # K and V
         * n_layers
         * n_kv_heads
         * d_head
         * seq_len
         * batch_size
         * bytes_per_element

# Two things follow immediately:
#   the cache grows linearly with BOTH context length and batch size,
#   so long contexts directly reduce how many users you can serve.

Put concrete numbers from your own model into that and you get the batch size at which you run out of memory — which is to say, your maximum throughput, which is to say, your cost per request. This calculation deserves to be done before deployment rather than discovered under load.

Two optimisations exist specifically because of it. Grouped-query attention lets several query heads share one key/value head, cutting n_kv_heads and therefore the cache by the group factor, at a modest quality cost. Paged attention allocates the cache in fixed blocks rather than contiguous per-sequence buffers, which removes the fragmentation waste from reserving worst-case length for every request.

What context length actually costs

"Long context is expensive" is true and too vague to act on. The components scale differently and the dominant one depends on which phase you are in.

Two consequences people find counterintuitive.

The feed-forward network often dominates prefill, not attention. Attention is quadratic in n and the feed-forward is quadratic in d. For a model with d = 4096 and a prompt of 1,000 tokens, d ≫ n and the feed-forward is the larger term. Attention only takes over once n approaches d. So "attention is the bottleneck" is a statement about very long contexts, not about all contexts.

A long conversation gets progressively more expensive per token. Each decode step attends over the accumulated cache, so token 4000 costs measurably more than token 40. Users experience this as a chat that slows down as it goes, and it is not their imagination.

Memory-efficient attention implementations change the constant substantially — by keeping tiles of the computation in fast on-chip memory they avoid ever materialising the full n×n score matrix, which removes the quadratic memory cost — but the arithmetic remains quadratic. They make long context practical; they do not make it cheap.

Where the parameters live

A useful thing to be able to work out on a whiteboard, because it grounds every other intuition about these models.

per block:
  attention   4 * d_model²        # W_q, W_k, W_v, W_o
  feed-forward 2 * d_model * d_ff  # = 8 * d_model²  when d_ff = 4d
  ------------------------------------------------
  total       12 * d_model² per block

whole model:  n_layers * 12 * d_model²  +  vocab * d_model

Two facts fall out of this that are worth carrying.

Roughly two thirds of the parameters are in the feed-forward layers, not in attention. Attention gets the attention, but the feed-forward network is where most of the model's capacity sits — which is consistent with the evidence that a great deal of factual knowledge is stored there.

The embedding matrix is significant for small models and negligible for large ones. With a 50,000-token vocabulary and d = 768, embeddings are about 38M parameters — a large fraction of a 120M model. At d = 8192 the same vocabulary is 400M parameters against tens of billions elsewhere, which is a rounding error. This is why vocabulary size is a real design decision for small models and mostly is not for large ones.