Build your own KV cache

What you'll build

You'll turn the GPT-2 124M forward pass from llm-inference — which recomputes K and V for the whole sequence on every token — into a real autoregressive engine with a key/value cache. Each decode step computes K/V for only the new token, appends them, and attends over the cached history; a prefill pass populates the cache from the prompt; a prefix cache reuses it across requests. The bar is the one production lives by: cached decode must reproduce the slow recompute path's logits, within tolerance, at every step.

• A per-layer KV cache and the single-token cached attention that appends one K/V and attends over all of cached history
• Incremental prefill plus a full 12-layer cached decode step — logit-identical to full recompute at every position
• A greedy generate loop and a prefix cache that reuses a saved prompt cache across requests, continuing absolute positions
• A deterministic proof of the win — per-step op counts (O(n) cached vs O(n²) recompute) and the exact KV-cache memory bill

What you'll learn

• Why the causal mask makes caching correct, not just fast — past K/V are frozen, so storing and reusing them changes nothing
• The absolute-position rule that trips everyone: a new token's position is the cache length before you append
• Why 'identical' means allclose, not bit-identical — float reassociation, and how to grade an optimization honestly
• Time-to-first-token (prefill) vs inter-token latency (decode), and why context length blows up memory

Prerequisites

• Completed llm-inference — a working GPT-2 attention + forward pass, which you refactor into a cached form here
• Comfortable reading and running TypeScript
• Basic linear algebra — vectors and matrix multiplication