Unit 5.2 · about 5 h
319 / 324Building a KV cache
A cache that shares layers, mixes sequences or misplaces positions returns fluent garbage with matching shapes. Ownership, layout and offsets are what make it correct.
- implement
You can now
- Implement per-layer KV allocation, append and single-token cached decode with absolute position offsets for the tiny transformer.
- Explain MHA, MQA and GQA head mappings and the
repeat_interleaverule for shared KV heads. - Explain KV tensor shapes and why each layer owns its cache for its own sequence.
Covers
- Per-layer KV ownership and sequence identity
- KV tensor shapes, layout, capacity and append state
- Position offsets and RoPE consistency under cached decode
- MHA, MQA and GQA KV head mapping
Study, in this order
3 steps
- Raschka `gpt_with_kv_cache.py`: per-layer cache with positionsLLMs from Scratch companion - KV cache and GQA
MultiHeadAttention(cache_kandcache_vbuffers,use_cacheappend bytorch.caton dim 1,ptr_current_posmask slice,reset_cache),GPTModel(current_poswithpos_idsarange,ModuleListloop,reset_kv_cache) andgenerate_text_simple_cached(full-prompt prefill then single-token loop) - Raschka `ch04/04_gqa`: GQA head mapping and memoryLLMs from Scratch companion - KV cache and GQA
GroupedQueryAttention(num_kv_groupswithgroup_size, \(K\) and \(V\) cache on dim 2,repeat_interleavemapping with therepeatversusrepeat_interleavecomment,q_positionsversusk_positionsmask) plus the README bytes formula withn_kv_headsand the 4096/32/32/32768 example - Modern LLM Notebook 21 Sections 2 and 4: shapes and head sharingModern LLM Notebook - KV cache and inference acceleration
Section 2 shape evolution (
[batch, num_heads, cached_seq_len, d_k]from[1,4,1,8]to[1,4,5,8]) and Section 4 MHA/MQA/GQA (32 \(Q\) heads with 32, 8 and 1 KV heads at 8K, the 8 \(Q\) with 2 KVrepeat_interleavedemo, savings ratio 32/8 = 4x)Depth · Skip Sections 5 to 8 and 10 (quantization, GGUF, PagedAttention, FlashAttention, batching) until their own units.
Focus on
- Each layer owns its \(K\) and \(V\) tensors and each batch row owns its positions. Sharing a cache across layers or sequences mixes unrelated spaces.
- Append grows the sequence dim by the new chunk only. Prefill appends the whole prompt at once, decode appends one token per step.
- Positions are absolute.
pos_idsand the mask pointer advance by the chunk length, so decode step \(k\) uses positioncurrent_posplus \(k\). - GQA stores fewer KV heads and expands them with
repeat_interleave. With 8 \(Q\) heads and 2 KV heads, heads 0 to 3 share KV head 0.
Reproduce closed-book
pos_ids, reset, and the GQA repeat_interleave mapping for 8 \(Q\) heads with 2 KV heads.Traps
- Forgetting
reset_kv_cachebetween independent generations. Stale \(K\), \(V\) and pointers leak the old sequence into the new one. - Using
repeatinstead ofrepeat_interleavefor grouped heads. Tiling maps the wrong KV head to each query group while shapes still match. - Restarting positions at 0 on every decode step. Absolute embeddings repeat and RoPE rotations reuse early angles, so logits diverge from the full prefix.
Practice
Raschka `ch04/03_kv-cache` tests: cached equivalence
test_gpt_model_equivalence_cached and test_gpt_model_equivalence_not_cached (30 new tokens from Hello, I am across base, cached and optimized models with torch.equal)
Run against the reference files, then re-run against your own per-layer cache in your own repository.
Self-test
-
After a 4-token prefill, the first decode step feeds 1 token. What are the \(Q\) and \(K\) lengths and which mask row applies?
\(Q\) has 1 token and \(K\) has 5 (4 cached plus 1 new). The pointer is 4, so the slice is row 4 over 5 columns, which lets the new query attend to all 5 positions. -
With 32 \(Q\) heads and 8 KV heads, which \(Q\) heads share KV head 0, and why
repeat_interleave?\(Q\) heads 0 to 3 share KV head 0 in contiguous groups of 4.repeat_interleavegives[K1,K1,K2,K2], whilerepeatwould tile as[K1,K2,K1,K2]and map the wrong heads. -
Why must decode
pos_idsusecurrent_pos, and what breaks for RoPE if they restart at 0?Positions are absolute. Restarting at 0 reuses early embeddings and early RoPE angles for late tokens, so relative distances are wrong and logits diverge from the full prefix. -
Two layers share one cache object to save memory. What fails?
Each layer has its own projections, so its \(K\) and \(V\) live in different spaces. Sharing mixes them and equivalence fails; each layer needs its own cache just as each sequence needs its own rows.
Completion criterion
Stop when
Your own cache passes both equivalence tests with torch.equal and you can state the cache shape after each step.
Leads to
- 5.3 Proving the cache correct Prove the cache matches full-prefix logits.
- 5.4 KV memory arithmetic Size the cache you built across dtypes and head schemes.
Comes after
- 2.5 RoPE: positions as rotationsPreserve absolute positions when appending rotated keys.
- 5.1 Why cache: recomputation in autoregressive decodingBuild the cache you just motivated.
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.
- 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.
Also in this section
- 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.
- 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.