Unit 5.4 · about 2 h
321 / 324KV memory arithmetic
The cache can outweigh the weights at long context, so every serving decision starts with its size. One formula with head count and dtype predicts whether a workload fits before anything runs.
- derive
You can now
- Derive KV bytes per token, sequence, layer and batch from layers, KV heads, head dim, dtype and batch.
- Predict cache capacity across dtypes including INT8 and INT4 and state GB versus GiB units.
- Explain the quantized-cache trade-off: a smaller cache for a small quality risk.
Covers
- KV bytes per token, sequence, layer and batch
- Cache dtype and quantized-cache vocabulary
- Cache-capacity sensitivity across dtypes
Study, in this order
3 steps
- Modern LLM Notebook 21 Section 3 and Exercise 1: KV sizesModern LLM Notebook - KV cache and inference acceleration
Section 3 (
kv_cache_sizeas \(2 \times \text{layers} \times \text{heads} \times \text{head\_dim} \times \text{seq\_len} \times \text{bytes}\), LLaMA-2 7B, 13B and 70B from 2K to 128K, FP16/INT8/INT4 plot) and Exercise 1 (32 layers, 32 heads, 128 dim, 4096 tokens FP16 with the 2048 MB assert) - Raschka GQA estimator and README numbersLLMs from Scratch companion - KV cache and GQA
calc_kv_bytes_total(per_layeras \(\text{batch} \times \text{context} \times \text{head\_dim} \times \text{n\_kv\_heads} \times 2 \times \text{bytes}\), times \(\text{layers}\);DTYPE_BYTESwithfp324,bf16andfp162,fp8andint81) and the README example (4096 emb, 32 heads, 32 layers, 32768 context, 4 groups,bf16: MHA 17.18 GB versus GQA 4.29 GB at 4.00x and 75 percent) - Modern LLM Notebook 11 Sections 1 and 2: back-of-envelope and head schemesModern LLM Notebook - KV cache and inference acceleration
Section 1 (Llama-7B 32K FP16 back-of-envelope with the cache above the 13 GB weights) and Section 2 (MHA, GQA-8, GQA-4 and MQA table at 32768 tokens with the
kv_cache_elementsfunction)
Focus on
- Bytes equal \(\text{batch} \times \text{seq\_len} \times \text{layers} \times \text{n\_kv\_heads} \times \text{head\_dim} \times 2 \text{ for } K \text{ and } V \times \text{bytes per element}\). Every term is load-bearing.
- MHA stores \(\text{n\_heads}\) KV heads, MQA stores one, and GQA stores something in between, so the saving is \(\text{n\_heads}/\text{n\_kv\_heads}\). In the Raschka estimator,
--n_kv_groupsis the number of query heads per KV head: 32 heads with--n_kv_groups 4store 8 KV heads. - Dtype scales the total directly: FP16 2 bytes, INT8 1 byte, INT4 half a byte. Quantized cache is lossy while head sharing is structural.
- The estimator divides by \(1000^3\) (GB) and the notebooks by \(1024^3\) (GiB). Same bytes, different labels.
Reproduce closed-book
Traps
- Forgetting one factor of 2. \(K\) and \(V\) each store a full vector, so dropping either 2 halves or doubles the answer.
- Counting query heads instead of KV heads under GQA. Only KV heads occupy cache; 32 \(Q\) heads with 8 KV heads store 8, not 32.
- Comparing GB and GiB numbers as if they disagree. 17179869184 bytes is both 17.18 GB and 16.0 GiB.
Practice
Modern LLM Notebook 21 Exercise 1: KV cache memory
Exercise 1: compute kv_bytes for 32 layers, 32 heads, 128 dim and 4096 tokens FP16; the cell asserts 2048 MB
Solve by hand from the head formula first, then confirm with the cell.
Modern LLM Notebook 11 Exercise 1: KV footprint at 16K
Exercise 1: fill elements_per_token_per_layer, total_elements and total_gb for 4096 hidden, 32 layers, 32 heads and 16384 tokens FP16; the cell asserts 8.00 GB
Report the answer in both GB and GiB and name which divisor each uses.
Self-test
-
With 32 heads and head dim 128 in FP16, what are the \(K\) plus \(V\) elements and bytes per token per layer?
8192 elements (\(32 \times 128 \times 2\)) and 16384 bytes (\(\times 2 \text{ bytes}\)), which is 16 KiB. -
Same model with 32 layers at 16384 tokens FP16. What is the total cache?
\(4294967296 \text{ elements} \times 2 \text{ bytes} = 8589934592 \text{ bytes}\), which is 8.0 GiB in \(1024^3\) units. -
The estimator says 17.18 GB and notebook math says 16.0 for Llama-7B at 32K. Which is right?
Both. The bytes are 17179869184 either way; the estimator divides by 1000 cubed (GB) and the notebook by 1024 cubed (GiB). State the units. -
32 heads at
bf16with the estimator's--n_kv_groups 4, which stores 8 KV heads, then an INT8 cache. What are the savings?Head sharing gives 32/8 = 4x (MHA 17.18 GB to GQA 4.29 GB decimal, as run). INT8 halves it again to about 2.15 GB, for 8x versus FP16 MHA at a small quality risk.
Completion criterion
Stop when
You derive the 2048 MB and 8.00 GB answers by hand and both exercise asserts pass with units stated.
Mark passed, then open 5.5 Only mark this when you can do it without notes.Leads to
- 5.5 Two regimes: prefill and decode Use the sizes to predict compute versus bandwidth limits.
- 5.6 Measuring time to first token and inter-token latency Check the arithmetic against measured memory and latency.
Comes after
- 2.4 MQA and GQA: sharing keys and valuesExtend the per-layer KV calculation across layers and sequences.
- 4.6 Counting parameters and FLOPsCarry KV and sequence terms into cache capacity math.
- 4.7 Memory, arithmetic intensity and reconciling estimatesApply KV math to cache capacity and layout design.
- 5.2 Building a KV cacheSize the cache you built across dtypes and head schemes.
- 5.3 Proving the cache correctDerive capacity for the cases you proved.
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.