KV Cache and Inference Memory
The redundancy problem in autoregressive generation
A decoder-only LLM generates text one token at a time: predict the next token, append it to the sequence, repeat. Naively, you'd recompute attention over the entire sequence from scratch at every step — but each of the previous tokens' key and value projections ($K$ and $V$) never change once computed, because causal masking means earlier tokens never attend to later ones. Recomputing them again and again at every single generation step is pure wasted work. The KV cache stores every previous token's $K$ and $V$ vectors in GPU memory, so each new generation step only needs to compute $Q$, $K$, $V$ for the one new token and attend against the cached history.
How much memory this actually costs
The KV cache size scales as $2 \times L \times H \times d_h \times N \times \text{batch} \times \text{bytes}$, where $L$ is the number of layers, $H$ the number of attention heads, $d_h$ the head dimension, $N$ the sequence length, and the factor of 2 accounts for storing both $K$ and $V$. For a 70B-parameter model with 80 layers and a 32K-token context at batch size 1 in fp16, the KV cache alone can consume tens of gigabytes — often rivaling or exceeding the memory used by the model's weights themselves. This is why long-context serving is primarily a memory engineering problem, not a compute one.
def kv_cache_bytes(num_layers, num_kv_heads, head_dim, seq_len, batch=1, bytes_per_val=2):
"""Returns KV cache size in bytes. bytes_per_val=2 for fp16/bf16."""
return 2 * num_layers * num_kv_heads * head_dim * seq_len * batch * bytes_per_val
# Example: Llama-2-70B-like config, 32K context, batch size 1
size = kv_cache_bytes(num_layers=80, num_kv_heads=8, head_dim=128, seq_len=32_000)
print(f"KV cache: {size / 1e9:.1f} GB")
# Note num_kv_heads=8 assumes grouped-query attention (fewer KV heads than query heads),
# which is exactly the optimization described below.Multi-query and grouped-query attention shrink the cache directly
Attention head sharing strategies
| Scheme | KV heads vs query heads | Cache size | Quality vs MHA |
|---|---|---|---|
| Multi-Head Attention (MHA) | 1:1 | Largest | Baseline (best) |
| Grouped-Query Attention (GQA) | Several query heads share 1 KV head | Moderate reduction | Close to MHA |
| Multi-Query Attention (MQA) | All query heads share 1 KV head | Largest reduction | Noticeable degradation |
PagedAttention and non-contiguous cache allocation
Beyond shrinking per-token cache size, serving systems like vLLM introduced PagedAttention, which borrows the idea of virtual memory paging from operating systems. Instead of reserving one large contiguous memory block per request (which wastes memory to fragmentation and over-provisioning for worst-case length), the KV cache is split into fixed-size blocks allocated on demand, with a lookup table mapping logical positions to physical blocks. This lets a server pack far more concurrent requests into the same GPU memory, and makes memory sharing across requests (e.g., identical prompt prefixes) straightforward.
The cache is why 'just increase context length' isn't free
What's Next?
Managing the KV cache efficiently is only half of serving — the other half is deciding, token by token, which candidate to actually emit from the model's probability distribution, which is a whole design space of its own.
I build these systems professionally.
Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.