GPT-2, and the price of perfect recall

Softmax attention has a property nothing later in this series can match: it never forgets, never overwrites, and never confuses two facts. It gets that by keeping one memory slot per token — so the memory grows for as long as you keep talking. Everything after this chapter is an attempt to keep the property and lose the bill.

The entire model, in nine lines

GPT-2 is decoder-only: one stack of identical blocks, each token allowed to look only at itself and everything before it. The top-level forward pass is genuinely this short.

# idx: (B, T) token ids.  pos: (T,) positions 0..T-1
tok_emb = self.transformer.wte(idx)   # (B, T, n_embd) — what the token is
pos_emb = self.transformer.wpe(pos)   # (T, n_embd)    — where it sits
x = self.transformer.drop(tok_emb + pos_emb)
for block in self.transformer.h:      # 12 identical blocks
    x = block(x)
x = self.transformer.ln_f(x)
logits = self.lm_head(x)              # (B, T, vocab_size)
return logits

Two embedding tables, added together. GPT-2 learns its positional embeddings as a lookup table indexed by absolute position — a design decision that later models abandon, but it is not what this series is about. The interesting part is what happens inside block.

class Block(nn.Module):
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # mix across tokens
        x = x + self.mlp(self.ln_2(x))    # mix across channels
        return x

Two lines, and they do genuinely different jobs. The attention line is the only place in the whole model where information moves between token positions. The MLP line processes each position independently. Every architecture in this series changes the first line and leaves the second alone until chapter 7.

Both lines are written as x = x + f(x). That plus sign is the residual stream: a running sum that every layer reads from and adds to. It looks like plumbing. In chapter 7 it becomes the target of an architectural change of its own.

What attention actually computes

Strip away the batching and the heads and attention is a weighted average of value vectors, where the weights come from how well each key matches the current query:

q, k, v are the three linear projections of the same hidden state; i ≤ t is the causal mask; √d keeps the dot products from saturating the exponential at large head width.

The softmax does exactly two things: it forces the weights to be non-negative, and it makes them sum to one. That is what makes \(o_t\) an average rather than an arbitrary linear combination. Hold onto those two requirements — chapter 2 satisfies them a completely different way.

And notice where the exponential sits. It is applied to the dot product, after \(q_t\) and \(k_i\) have already met. This is the single most consequential structural fact in the whole series, for two opposite reasons.

Why the KV cache exists, and what it buys

Autoregressive decoding produces one token at a time. Feed in the prompt, take the logits at the last position, sample a token, append it, and go again. The observation that makes this affordable is small and slightly boring: appending a token does not change the key and value vectors of any earlier token. Each \(k_i\) and \(v_i\) depends only on hidden state \(i\), which depends only on tokens \(1..i\). Adding token \(t+1\) at the end leaves all of them untouched.

So keep them. That store is the KV cache: for every layer, for every head, the \(k\) and \(v\) vectors of every token seen so far.

The saving is larger than it looks. Run the numbers on generating \(N\) tokens:

The cache turns cubic decoding into quadratic decoding. It is not an optimisation, it is the thing that makes generation possible at all. And in exchange you now carry a block of memory that grows by a fixed amount with every token you emit, forever.

How big does it actually get?

GPT-2 small: 12 layers, 12 heads, \(n_{\text{embd}} = 768\), so \(d = 64\) per head and \(12 \times 64 = 768\) per layer. Two tensors (K and V), two bytes each in fp16:

36 KiB per token, for a model whose entire weight set is 248 MB in fp16. Which means:

Scale that up and it stops being a curiosity. A 70B-class model with plain multi-head attention — 80 layers, 8192 hidden — carries 2.62 MB of cache per token. At 100k context that is 262 GB, for a single sequence, on top of the 140 GB of weights. This is the pressure that produced grouped-query attention, then multi-head latent attention (chapter 7), and eventually the entire linear-attention branch of this series.

The bottleneck is bandwidth, not arithmetic

It is tempting to read \(O(N)\) work per step as "attention decode is slow because of the FLOPs." It isn't. Look at the ratio of arithmetic to memory traffic.

To attend over one cached position you read its \(k\) and \(v\) — \(2d\) numbers, so \(4d\) bytes in fp16 — and you do one dot product plus one scaled accumulation, about \(4d\) floating-point operations. The arithmetic intensity is therefore

A modern accelerator wants somewhere in the region of a few hundred FLOPs per byte to keep its matrix units busy — an H100 pairs roughly 500–1000 TFLOP/s of tensor throughput with about 3.35 TB/s of HBM bandwidth. At an intensity of 1, the arithmetic units are idle essentially all the time. Every cached byte is read, used once, and thrown away.

The property worth preserving

Before dismantling it, be precise about what softmax attention gives you, because the next five chapters are all measured against it.

One slot per token

Attention never has to decide what to evict, because nothing ever shares a slot. Token 4,000 and token 40,000 sit in separate rows of the cache and cannot corrupt each other. There is no capacity limit, so there is no eviction policy.

Sharp, content-addressed retrieval

The unbounded exponential lets one score dominate all the others. The model can genuinely look up one earlier token. Every approximation in this series trades some of this away.

Position-independent addressing

Retrieval quality does not degrade with distance. A key 100,000 tokens back competes on equal terms with one from the previous word — assuming positional encoding does not interfere.

The baseline config, for reference:

vocab_size: int = 50304   # GPT-2's 50257, padded to a multiple of 64
n_layer:    int = 12
n_head:     int = 12
n_embd:     int = 768      # → 124M parameters, → 22,580 of these fit in Kimi K3