Reference — the answer first
What is the KV cache?
Direct answer
A Transformer language model generates text one token at a time. At each step it keeps a key and a value for every earlier token, in every attention layer and every key-value head. That store is the KV cache. It exists so the model does not recompute the keys and values of earlier tokens at every step. Its size grows linearly with the number of tokens stored and with the number of key-value heads each layer keeps. Grouped-query attention (GQA) and multi-query attention (MQA) shrink it by letting several query heads share the same stored key-value heads. With long contexts and large batches, the cache can take more memory than the model weights themselves.
The memory formula
Plain text: Mem_KV = B * N_layers * T * H_kv * d_head * 2 * bytes
- B
- batch — sequences decoded together
- N_layers
- model layers that each store keys and values
- T
- earlier token positions kept in the cache
- H_kv
- stored key-value heads per layer — the term head sharing changes
- d_head
- dimensions per head
- 2
- one tensor for keys, one for values
- bytes
- bytes per stored number, such as 2 for fp16/bf16
This counts the bytes of the keys and values themselves. A real serving system also spends memory on allocation overhead, paging metadata and padding, and changing the attention architecture can change model quality; the formula covers neither.
A worked example
Take one configuration, the same defaults as the KV memory step in the Lab: batch 1, 32 layers, 4,096 stored positions, 32 query heads, head dimension 128, fp16/bf16 (2 bytes per number). All results are in decimal GB (10⁹ bytes), computed with the formula above.
- Full multi-head attention stores H_kv = 32 key-value heads: 2.15 GB.
- GQA with H_kv = 8 at the same context: 0.54 GB — 4.0× smaller, because 8/32 of the heads remain.
- The same GQA cache at 32,768 positions: 4.29 GB, 2.0× the full multi-head cache at 4,096 positions, because multiplying the context by 8 outweighs dividing the stored heads by 4.
With B, N_layers, d_head and precision fixed, KV-cache memory is proportional to T × H_kv. Sharing heads makes the cache narrower; a longer context makes it longer.
MHA vs GQA vs MQA
| Variant | Stored K/V heads (H_kv) | Cache memory | vs MHA |
|---|---|---|---|
| MHA — no sharing | 32 | 2.15 GB | 1.0× |
| GQA — grouped sharing | 8 | 0.54 GB | 4.0× smaller |
| MQA — one shared head | 1 | 0.07 GB | 32.0× smaller |
The formula gives cache size only. Quality and speed have to be measured: MQA can lower quality, and Ainslie et al. (2023) report that GQA models converted from multi-head checkpoints come close to multi-head quality while decoding nearly as fast as MQA.
Three common misconceptions
- “Sharing shrinks the query heads.” GQA keeps every query head — 32 stay 32 in the worked model. What shrinks is the number of stored key-value heads (H_kv), from 32 to 8. Query-head count and stored-head count are different quantities.
- “Sharing forgets earlier words.” The same earlier positions stay in the cache — T does not change. Each position simply stores fewer distinct key-value sets for the heads to read.
- “A smaller cache proves faster, better answers.” The formula estimates stored bytes only. Allocator behavior, paging, batching, scheduling, and answer quality sit outside it; speed and quality claims need separate evidence.
Equation to code
The same calculation in Python, term for term, as used in the efficient-attention notebook:
def kv_cache_gb(
batch=1,
layers=32,
tokens=4096,
kv_heads=32,
head_dim=128,
bytes_per_value=2,
):
scalars = batch * layers * tokens * kv_heads * head_dim * 2
return scalars * bytes_per_value / 1e9| Equation factor | Code term | Meaning |
|---|---|---|
B | batch | sequences decoded together |
N_layers | layers | layers that store keys and values |
T | tokens | earlier positions kept in the cache |
H_kv | kv_heads | stored key-value heads per layer |
d_head | head_dim | dimensions per head |
2 | literal 2 | one key tensor plus one value tensor |
bytes | bytes_per_value | bytes per stored number (2 for fp16/bf16) |
÷ 1e9 | / 1e9 | decimal gigabytes, not GiB |
Sources
- Attention Is All You Need (Vaswani et al., NeurIPS 2017)
The attention mechanism the cache serves: each new query attends over stored keys and values.
- Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, arXiv 2019)
Multi-query attention: all query heads share one key/value head, which shrinks the cache and the memory read at each decoding step.
- Efficiently Scaling Transformer Inference (Pope et al., MLSys 2023)
Why the cache matters at serving time: for a 500B+ parameter model with multi-head attention, batch 512 and context 2,048, the cache totals 3 TB, three times the size of the weights, and it is read from memory for every generated token.
- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., EMNLP 2023)
Grouped-query attention, between MHA and MQA: multi-head checkpoints uptrained to GQA come close to multi-head quality while decoding at nearly MQA speed.
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., NeurIPS 2022)
A separate idea: exact attention made faster by moving less data between GPU memory levels. It does not change the size of the KV cache.
- Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., SOSP 2023)
Why the formula’s total is not the memory a serving system can actually use: how cache memory is allocated, fragmented and shared lies between the two.
- Speech and Language Processing, 3rd ed. draft, Chapter 7: Transformers and Pretraining (Jurafsky & Martin, draft of August 2026)
A textbook introduction to the KV cache: why keys and values of earlier tokens are stored rather than recomputed at each generation step.
Test your understanding
To test yourself, answer the two questions in the KV-cache challenge, or open the KV memory step in the Lab to change one input at a time and watch this formula respond. Further reading: grouped-query attention and efficient attention at scale.