2 episodes
Paper: Multi-head Latent Attention (DeepSeek-V2)
The KV cache as the bottleneck of decoding and how three designs shrink it. Episode 1 derives the cache formula, explains why small-batch decoding is memory-bound, introduces MQA and GQA (shared key/value heads), implements grouping by broadcasting rather than copying, and measures the cost of mean-pool checkpoint conversion. Episode 2 derives Multi-head Latent Attention: a cached latent per token, up-projections to keys and values, absorption of W_UK and W_UV by associativity, why RoPE on content keys breaks absorption, and the decoupled RoPE key. Labs build a KV calculator for real model shapes, NumPy GQA and MLA with equivalence tests, and torch modules for a TinyStories comparison of MHA, GQA, MQA and MLA.
Objectives
Language
- Explain why long contexts and many concurrent users make the KV cache the limiting resource in serving language models
- Read efficiency claims with their baselines (93.3% vs DeepSeek 67B; 98.2% vs a same-shape MHA counterfactual)
Math & data structures
- KV-cache size formula; arithmetic intensity and memory-bound decoding
- Low-rank joint compression; associativity and transposes for absorption; non-commutativity of rotations
- Grouped broadcasting (einsum with a group axis) as a data-layout choice
Code
- kv_accounting: formulas, presets, tables and charts; GQA without copies plus a G-head KV cache; mean-pool conversion and its error
- MLA naive and absorbed forms, a latent-only cache, the RoPE counterexample, decode-cost accounting
- Torch attention 'gqa' and 'mla' registered in nlpcore; TinyStories comparison with measured cache bytes
Episodes
- The KV Cache Problem: MHA, MQA, GQA15 min · 22 cues
KV memory math, MQA/GQA trade-offs.
- MLA: Compressing Keys and Values into a Latent16 min · 24 cues
Low-rank latent c_t, up-projections, absorption, decoupled RoPE key (arXiv 2405.04434).
Exercises
| Exercise | Tier | Ladder | Time |
|---|---|---|---|
| A KV-cache calculator: MHA, MQA, GQA and MLA, from formula to chart | browser | Predict · Scratch | 40 min |
| GQA without copies, a KV cache that stores G heads, and breaking the MHA-to-MQA conversion | browser | Scratch · Harden · Break it | 70 min |
| MLA from scratch: naive vs absorbed, a latent-only cache, and why RoPE must be decoupled | browser | Predict · Scratch · Harden · Probe a claim | 90 min |
| GQA and MLA in PyTorch, cache bytes measured, and a TinyStories quality comparison | gpu | Port · Compose · Probe a claim | 180 min |
Labs
- The KV Cache Problem: MHA, MQA, GQA browser F16.01 · 75 min
- MLA: Compressing Keys and Values into a Latent browser F16.02 · 80 min
Library contributions
Files in nlpcore this module teaches you to write.
Sources
- Ainslie, Lee-Thorp, de Jong, Zemlyanskiy, Lebrón & Sanghai (Google Research), GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, arXiv 2305.13245 (v3 Dec 2023) · arXiv:2305.13245
- DeepSeek-AI, DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, arXiv 2405.04434 (v1 7 May 2024; v5 19 Jun 2024) · arXiv:2405.04434
- DeepSeek-AI, DeepSeek-V3 Technical Report, arXiv 2412.19437 · arXiv:2412.19437
- Paper: Multi-head Latent Attention (DeepSeek-V2) labs (this module): NumPy measurements with random weights
- Hugging Face Transformers documentation, Llama 2 model page and LlamaConfig · link
- Noam Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, arXiv 1911.02150 (6 Nov 2019) · arXiv:1911.02150
- Jurafsky & Martin, Speech and Language Processing, 3rd ed. draft (Aug 19, 2026), Ch. 7 (7.1 Attention; 7.6 Decoding) · link · §7.1, 7.6
- Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding, arXiv 2104.09864 · arXiv:2104.09864