During generation, each new token only needs to compute its own Q, K, V — but it attends to all previous K and V vectors. Without caching, we'd recompute K,V for all prior tokens at every step. The KV-cache stores these, turning generation from O(n²) to O(n) per step.
KV-cache is why batch size during inference is heavily memory-constrained. GQA (fewer KV heads) directly reduces this cost. This is the main motivation behind MQA and GQA research.
Python — KV-cache in generation loop#
Generating 1,000 tokens without a cache recomputes keys and values for the whole prefix at every step: 500× as many projections as with one. The cache has a price: for Llama 3 8B it is 128 KB per token, so a 128k-token context holds 16 GB of cache — as much as the model’s own weights in fp16.
- The cache trades compute for memory; on long contexts memory, not arithmetic, limits how many users a GPU can serve.
- Reusing cache across requests (prefix caching) only works when the prefix is byte-identical.
- Compressing or evicting cache entries saves memory but can silently drop information the model needed.
T = 1_000 # tokens generated recompute = T * (T + 1) // 2 # K/V projections without a cache: the whole prefix every step cached = T # with a cache: one new token per step layers, kv_heads, d_k, bytes_per = 32, 8, 128, 2 # Llama-3-8B shapes, fp16 per_token_kb = 2 * layers * kv_heads * d_k * bytes_per / 1024 at_128k_gb = per_token_kb * 131_072 / 1024 ** 2
Assumes import numpy as np. Configurations named after real models use their published shapes; other numbers are made up or simulated, as the comments say.
The same idea with a library
def generate_with_cache(model, prompt_ids, max_new=50):
past_kv = None # will hold cached K,V tensors
input_ids = prompt_ids
for _ in range(max_new):
# Only feed new token(s) — cache handles the rest
logits, past_kv = model(input_ids, past_key_values=past_kv)
next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
input_ids = next_id # only the new token
if next_id.item() == eos_token_id:
break
return all_generated_ids- R. Pope et al., “Efficiently Scaling Transformer Inference”, Proceedings of Machine Learning and Systems 5, 2023
- N. Shazeer, “Fast Transformer Decoding: One Write-Head Is All You Need”, arXiv:1911.02150, 2019
- W. Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention”, SOSP, 2023
Last reviewed October 2026 · How this site checks its content · Report a mistake
1. What does the KV cache store?
Show answer
B. Without it, generating token n would mean re-running attention over all n − 1 earlier tokens from scratch.
2. KV cache memory grows with…
Show answer
B. Long contexts and large batches can make the cache larger than the weights themselves.
3. How does grouped-query attention (GQA) shrink the KV cache?
Show answer
A. With 8 KV heads serving 64 query heads, the cache is 8 times smaller than with full multi-head attention.