KV Cache
KV Cache
The KV (Key-Value) cache is a critical optimization for autoregressive LLM inference. It stores the Key and Value tensors from previous attention computations, avoiding redundant recomputation during token-by-token generation.
The Problem
In autoregressive generation:
- Token 1 attends only to itself
- Token 2 attends to tokens 1 and 2
- Token 3 attends to tokens 1, 2, and 3
- ...
Without caching, computing attention for the Nth token would recompute all previous K and V vectors from scratch, leading to O(N²) computation per token.
The Solution
The KV cache stores the K and V tensors for all previous tokens in each layer. When generating a new token:
- Only compute Q, K, V for the new token
- Retrieve cached K and V from all previous tokens
- Compute attention with the new Q against all K and V
- Append the new K and V to the cache
Memory Impact
The KV cache is the primary memory bottleneck during inference:
KV Cache Size = 2 × (number of layers) × (hidden dim) × (sequence length) × (bytes per element) × (KV heads)
For a 70B model with 80 layers, 8192 hidden dim, and 128K context:
- ~5-10 GB with standard caching
- Requires GQA/MQA to reduce memory
Optimization Techniques
- Grouped Query Attention (GQA) — Fewer KV heads than query heads
- Multi-Query Attention (MQA) — Single KV head
- KV cache quantization — Store K/V in 4-bit or 8-bit
- PagedAttention (vLLM) — Manage cache as pages (like virtual memory)
- Prefix caching — Reuse KV cache across requests with shared prefixes
- Sliding window — Only cache recent tokens
- Shared KV cache — For batched inference
Impact on Context Window
The KV cache memory scales linearly with context length, making long contexts expensive. This is why models like Gemini 2 Pro with 1M context windows require sophisticated cache management.
Related
- Attention Mechanism — What the cache serves
- grouped-query-attention — Reduces KV cache size
- Context Window — Memory cost scales with window
- streaming — Real-time token delivery enabled by caching