Reference — a direct answer, nothing sealed
What is the KV cache?
Direct answer
When a Transformer chatbot writes one new word at a time, it keeps a small reusable record — a key and a value — for each earlier position, in every attention layer. That store is the KV cache. It exists so the model does not redo the same work for every earlier word on every step. Its size grows linearly with how many positions are stored and with how many 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. At long context, the cache — not the model weights — can become the dominant memory cost of serving.
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 is a formula estimate. It does not include allocator fragmentation, paging metadata, scheduler behavior, prefix sharing, tensor layout padding, or model-quality effects from changing the attention architecture.
A worked example
Take a concrete configuration — the same defaults the interactive investigation uses: batch 1, 32 layers, 4,096 stored positions, 32 query heads, head dimension 128, fp16/bf16 (2 bytes per number). All results are decimal GB (10⁹ bytes), computed by the same formula shown 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 — head sharing narrowed the cache, and longer context grew it right back.
The invariant: holding B, N_layers, d_head, and precision fixed, KV-cache memory scales linearly with T × H_kv. Sharing reduces cache width; long context grows cached length.
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 witnesses cache width only. Model quality under sharing, uptraining requirements, and serving latency are separate claims that need their own evidence — the GQA paper places grouped sharing between MHA and MQA precisely because the endpoints trade differently.
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 formula maps term-for-term onto the code witness 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, 2019 (MQA))
Sharing one key/value head across query heads reduces stored KV width and repeated decoder K/V reads.
- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., EMNLP 2023)
GQA groups query heads over fewer key/value heads as an intermediate point between MHA and MQA.
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., NeurIPS 2022)
An adjacent claim kept separate here: exact attention can get faster by moving less data, without changing cache width.
- Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023 (vLLM))
Why formula bytes are not usable serving capacity: allocation, fragmentation, and sharing sit between them.
Do not read this page as a claim that GQA always preserves quality, that FlashAttention reduces the KV cache, or that the formula predicts serving capacity exactly. Those claims need their own evidence.
Test your understanding in the investigation
This page answers directly and seals nothing. If you would rather test your own model of the mechanism first: start with one prediction at the door, or open the KV memory station to change one input at a time and watch this same formula respond. Deeper reading: grouped-query attention and efficient attention at scale.