GPT-2 from scratch

A GPT-2 style decoder built from nothing but torch.nn.Linear — causal self-attention, multi-head attention, layer norm, transformer blocks, training, and KV-cached inference.

model/      attention → architecture → pretraining → training
kv_cache/   the full model, and the same model with a KV cache
video/      manim scenes explaining prefill and KV-cache redundancy
text/       training corpus

What a KV cache is

Generation is autoregressive: the model emits one token, appends it, and runs again. Without a cache, every step re-runs attention over the entire sequence — so producing token 500 recomputes the keys and values for tokens 1–499 that were already computed at every previous step. The work per token grows with the sequence, and total work is quadratic in the number of tokens generated.

The keys and values for past tokens never change, though. Causal masking means token i only ever attends backwards, so nothing a later token does can alter an earlier token's K or V. So keep them. A KV cache stores each layer's keys and values and appends one row per step, which splits generation into two phases:

• prefill — run the prompt through once, filling the cache
• decode — feed only the newest token, project its single K/V, append to the cache, attend against everything stored

Decode goes from re-reading the whole sequence to appending one row, so cost per token becomes constant instead of growing.

How much faster

Measured on an RTX 5080, model/ vs kv_cache/, greedy decoding from a 4-token prompt. Both models share identical weights and produce byte-identical output sequences — this is a speed change, not a behaviour change.

tokens generated	no cache	with cache	speedup
32	66.7 ms — 480 tok/s	67.7 ms — 473 tok/s	0.99×
64	152.6 ms — 419 tok/s	135.0 ms — 474 tok/s	1.13×
128	392.4 ms — 326 tok/s	269.0 ms — 476 tok/s	1.46×
256	1008.9 ms — 254 tok/s	539.3 ms — 475 tok/s	1.87×
512	2819.1 ms — 182 tok/s	1081.1 ms — 474 tok/s	2.61×
1024	9260.4 ms — 111 tok/s	2161.9 ms — 474 tok/s	4.28×

The speedup number is the least interesting column. What matters is that cached throughput is flat at ~474 tok/s at every length, while uncached throughput decays from 480 to 111 tok/s as the sequence grows. That flat line is the whole point: the cache turns per-token cost from O(sequence) into O(1). The speedup keeps widening with length — at 32 tokens the cache is marginally slower, because bookkeeping costs something and there is almost no history to save yet.

Benchmarks use randomly initialised weights — timing depends on shapes, not values. Best of 3 runs after a warm-up, with torch.cuda.synchronize() around each.