Section 5 of 5 · LLM Foundations

KV Cache, Prefill & Decode

Autoregressive generation has two distinct execution regimes and a stateful cache. This final section proves cache correctness first, then explains capacity and latency without importing an entire serving stack.

7 units 26 topics 0/7

Units

7 in this section

5.1 Why cache: recomputation in autoregressive decoding
explain · derive
2 h
5.2 Building a KV cache
implement
5 h
5.3 Proving the cache correct
implement
4 h
5.4 KV memory arithmetic
derive
2 h
5.5 Two regimes: prefill and decode
explain · derive
3 h
5.6 Measuring time to first token and inter-token latency
measure
4 h
5.7 Beyond one cache: prefix sharing, paging and speculation
explain
2 h

Before moving on

Checkpoint

Prove full-prefix and cached logits or tokens agree across adversarial MHA and GQA shapes, offsets and variable lengths. Derive KV capacity and measure prefill and decode separately across at least two context and batch regimes with trustworthy timing.

Deliberately out of scope

Leave for later

FlashAttention, Triton and custom GPU kernels
Paged attention, continuous batching and production schedulers
Tensor parallelism, disaggregation and serving-fleet design

Resources for this section

Shared by every unit above

LLMs from Scratch companion - KV cache and GQA
Sebastian Raschkaimplementation

primary

Use ch04/03_kv-cache and ch04/04_gqa with the chapter model code. Add explicit per-layer ownership, position offsets, variable lengths, EOS compaction and MHA-versus-GQA cache arithmetic to the runnable comparison.
Modern LLM Notebook - KV cache and inference acceleration
Walking Labsimplementation

required

Use the KV-cache portions of English notebooks 11-mla-kv-cache and 21-inference-acceleration. Skip MLA and advanced acceleration as core work. Recalculate cache sizes and verify timing after warmup and explicit synchronization.
Stanford CS336 - Lecture 10 Inference
Stanford Universitytheory

required

Use the Spring 2026 inference lecture to connect KV arithmetic, phase behavior and batching vocabulary. Keep kernels, distributed serving and expensive performance experiments for the later tracks.
How to Scale Your Model - inference chapter
JAX MLtheory

depth

After cache equivalence and phase measurement pass, read the inference chapter for arithmetic intensity, parallelism and production-scale reasoning. Treat it as the handoff to inference and performance tracks, not required foundation work.
AI Performance Engineering - inference-engine bridge
Waferreference

depth

Use Sections 4 and 5 only after this gate to choose the next inference, KV-cache, scheduling, quantization or distributed-serving resource. Do not treat its GPU and kernel material as unfinished LLM Foundations work.
No Magic - micro KV or paged-cache simulation
No Magicpractice

optional

Use one small CPU simulation such as microkv.py or micropaged.py only to visualize state ownership. Do not present Python timing as GPU or serving-performance evidence.

Apply the section

Mini-project · Cached generation, proved and measured (10 to 14 h)

Add a KV cache to the decoder and prove that it changes nothing but cost. Then measure prefill and decode separately, and explain the numbers with the arithmetic from section 4.

Steps

Add a per-layer cache of unexpanded \(K\) and \(V\) to your model, with absolute positions for RoPE, and a forward that takes new tokens, the cache and a start position.
Prove equivalence: greedy tokens match and per-step logits agree within tolerance with the uncached model for MHA, GQA and MQA, chunked prefill, a one-token prompt, a variable-length batch with padding, a sequence that stops early at EOS, and a reused cache slot.
Derive the KV-cache bytes for your config and for one published model config, and check the formula against the real tensor sizes of your cache.
Measure time to first token and inter-token latency on a CPU or a free GPU across two context lengths and two batch sizes, with warm-up, device synchronization and repeated trials.
Explain the measurements: which phase is compute-bound, which is memory-bound, and how the cache changes the cost of each decode step.
Add a test that fails if a decode step recomputes the prefix, for example by counting the tokens that reach the \(K\) and \(V\) projections at each step.

Deliverables

The cached model and its equivalence test suite.
A KV memory sheet: the formula, your config, one published config and the measured cache size.
A latency report: time to first token and inter-token latency per regime, the timing method, and an explanation of each trend.

Checks

rubric · Self-grade with the rubric below

Rubric

Criterion	Pass	Strong
Equivalence	Cached and uncached greedy tokens match on every case in your equivalence tests.	Per-step logits also match within `1e-5` on your own adversarial cases, and a deliberately wrong RoPE offset makes your tests fail.
Memory	The KV formula matches the measured cache size.	The sheet scales the formula to a published model at long context and large batch, and names what dominates memory there.
Timing method	Prefill and decode are timed separately with warm-up and repeated trials.	The report gives medians with spread, synchronizes the device and states what the timings exclude.
Explanation	Each trend is explained in a sentence.	The explanation uses arithmetic intensity from section 4 and predicts one result before measuring it.

Feeds into: Capstone artifact 3: the cached-generation half of the dossier and the full-prefix cache equivalence that the stage milestone asks for.