Unit 5.3 · about 4 h
320 / 324Proving the cache correct
A cache that is fast but subtly wrong poisons every measurement after it. Equivalence across shapes, offsets and messy batches is what lets you trust the rest of the section.
- implement
You can now
- Write equivalence tests comparing full-sequence logits with token-by-token cached logits across MHA and GQA, batch sizes and position offsets.
- Handle variable lengths with padding masks, EOS compaction and cache cleanup without breaking equivalence.
- Diagnose truncation, fork and stale-cache failures from a failing equivalence test.
Covers
- Full-prefix versus cached-logit equivalence
- Variable-length batches and padding state
- Finished-sequence compaction and EOS handling
- Cancellation cleanup, truncation, forks and stale-cache hazards
Study, in this order
3 steps
- Raschka `tests.py` and optimized cache: equivalence, chunking and overflowLLMs from Scratch companion - KV cache and GQA
test_gpt_model_equivalence_cachedandnot_cached(30 tokens withtorch.equal),test_prefill_chunking_basic(kv_window_size4 over 10 tokens in 3 chunks, 12 total with input preserved) andtest_context_overflow_bug(context_length10 with 5 plus 10 tokens hitting the position-embedding assert) - CS336 Lecture 10: selective batching and paged sharingStanford CS336 - Lecture 10 Inference
continuous_batching(ragged lengths[3,9,5]with attention per sequence and non-attention concatenated) andpaged_attention(blocks, sharing the system prompt and parallel samples, copy-on-write at block level)Depth · Read the two functions only. Keep kernels, distributed serving and expensive performance experiments for the later tracks.
- JAX scaling book inference: ragged reads and the EOS loopHow to Scale Your Model - inference chapter
Tricks section (ragged HBM reads and Paged Attention that read only the non-padding part) and The Basics (Generation repeats until EOS or a maximum length)
Focus on
- Generated tokens must match exactly at a fixed seed; logits must match within a small tolerance, such as
1e-5infloat32, because different matmul shapes round differently. Check both before trusting any timing. - Variable lengths need per-row lengths and padding masks. Reading the padded tail as real tokens breaks both equivalence and the batch.
- Finished sequences need compaction: free or clear their rows, compact the batch, and stop them at EOS plus a length cap. Reused slots must be cleared first.
- Truncation, forks and reuse without reset are stale-cache hazards. Chunked prefill and windowed caches must preserve order and positions.
Reproduce closed-book
Traps
- Testing equivalence on one prompt length and head count. Offset, GQA grouping and batch bugs hide until adversarial shapes and longer contexts.
- Freeing a finished row without clearing it. The next request inherits old \(K\) and \(V\) and attends to another sequence, which looks like a model bug.
- Truncating the cache but keeping old position ids. Tokens shift left while positions stay high, so RoPE and absolute embeddings disagree with the full prefix.
Practice
Raschka `ch04/03_kv-cache` tests: chunking and overflow
test_prefill_chunking_basic and test_context_overflow_bug (windowed chunking with input preservation, plus the position-embedding overflow assert)
Then extend in your own repository to variable-length batches with padding masks, EOS compaction and a stale-cache reuse test. The mini-project checks the full matrix across MHA and GQA.
Self-test
-
Cached and full-prefix logits differ from position 0 onward. What class of bug is it?
A systematic one: wrong weights, a different prompt, a stale cache from a prior run, or a wrong mask or position base. Late-only divergence would point at windows or truncation instead. -
With
kv_window_size4, a 10-token prompt prefills in 3 chunks and generates 2 tokens. What total length and input check must hold?12 total tokens. The first 10 must equal the input, as the chunking test asserts with shape 12 and input preservation. -
Batch lengths are
[3,9,5]. Why does CS336 run attention per sequence but concatenate to[17,H]for the MLP?MLP weights are shared, so concatenated tokens batch well. Attention KV differs per sequence, so scores must stay per sequence; padding all to 9 wastes compute. -
A finished sequence is freed without clearing and a new request reuses the slot. What hazard and what fix?
A stale-cache hazard: the new request attends to old \(K\) and \(V\), possibly leaking another sequence. Clear or reset slots on finish, compact the batch, and test slot reuse against the full prefix.
Completion criterion
Stop when
Your harness shows identical tokens and per-step logits within tolerance on adversarial MHA and GQA shapes with offsets, and your chunking, overflow and slot-reuse tests all pass.
Mark passed, then open 5.4 Only mark this when you can do it without notes.Leads to
- 5.4 KV memory arithmetic Derive capacity for the cases you proved.
- 5.6 Measuring time to first token and inter-token latency Time prefill and decode separately once equivalence holds.
Comes after
- 2.7 Assembling and verifying a decoderReuse prefix equivalence when validating incremental execution.
- 5.2 Building a KV cacheProve the cache matches full-prefix logits.
Resources for this unit
- LLMs from Scratch companion - KV cache and GQA 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.
- Stanford CS336 - Lecture 10 Inference 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 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.
Also in this section
- Modern LLM Notebook - KV cache and inference acceleration 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.
- AI Performance Engineering - inference-engine bridge 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 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.