KV Cache: Comprehensive Technical Guide¶
Overview¶
KV Cache (Key-Value Cache) is a fundamental optimization technique in transformer-based LLMs that stores precomputed Key and Value matrices during inference to avoid redundant calculations and dramatically speed up token generation.
- Critical for: LLM inference efficiency
- Trade-off: Memory for Speed
- Impact: 5-10x speedup in token generation
- Challenge: Memory-bound operation in modern GPUs
- Evolution: From basic caching to quantization and multi-query variants
The Fundamental Problem: Redundant Computation¶
Understanding Transformer Attention¶
The self-attention mechanism in transformers computes:
Attention(Q, K, V) = softmax(QK^T / √d_k) V
Where: - Q (Query): What we're looking for - K (Key): What information is available - V (Value): The information itself - d_k: Dimension of keys
Token Generation Process (Without Caching)¶
Generating token 1:
- Input: "Hello"
- Compute Q₁, K₁, V₁ for "Hello"
- Attention: softmax(Q₁K₁ᵀ / √d_k) V₁
- Output: token₁
Generating token 2:
- Input: "Hello world" (full sequence!)
- Compute Q₂, K₂ for "Hello world" ← REDUNDANT!
- Compute V₂ for "Hello world" ← REDUNDANT!
- (K₁, V₁ for "Hello" already computed!)
- Attention: softmax(Q₂K₂ᵀ / √d_k) V₂
- Output: token₂
Generating token 3:
- Input: "Hello world foo" (full sequence!)
- Compute Q₃, K₃ for "Hello world foo" ← REDUNDANT!
- Compute V₃ for "Hello world foo" ← REDUNDANT!
- (K₁, V₁, K₂, V₂ already computed!)
- Output: token₃
WASTE: Recompute K and V for all previous tokens every step!
Computational Complexity¶
For a sequence of length N:
Without KV Cache:
Token 1: Compute K, V for 1 token
Token 2: Compute K, V for 2 tokens (1 redundant)
Token 3: Compute K, V for 3 tokens (2 redundant)
...
Token N: Compute K, V for N tokens (N-1 redundant)
Total redundant computations: 1 + 2 + 3 + ... + (N-1) = N(N-1)/2
Complexity: O(N²) operations!
For N=1000: 1000 × 999 / 2 = 499,500 redundant K, V computations!
The Solution: KV Cache¶
Basic Concept¶
Instead of recomputing K and V for all previous tokens, store them:
Generating token 1:
- Compute Q₁, K₁, V₁
- Attention: softmax(Q₁K₁ᵀ / √d_k) V₁
- Cache: Save K₁, V₁ to memory
- Output: token₁
Generating token 2:
- Compute Q₂, K₂_new, V₂_new (only for current token)
- K₂ = concat(K₁_cached, K₂_new)
- V₂ = concat(V₁_cached, V₂_new)
- Attention: softmax(Q₂K₂ᵀ / √d_k) V₂
- Cache: Update cache with K₂_new, V₂_new
- Output: token₂
Generating token 3:
- Compute Q₃, K₃_new, V₃_new (only for current token)
- K₃ = concat([K₁, K₂, K₃_new]) (from cache!)
- V₃ = concat([V₁, V₂, V₃_new]) (from cache!)
- Attention: softmax(Q₃K₃ᵀ / √d_k) V₃
- Cache: Update cache
- Output: token₃
NEW: Only compute K, V for 1 new token per step!
Complexity: O(N) operations!
Speed Improvement¶
Sequence length: 1000 tokens
Without KV Cache:
Total operations: 1 + 2 + 3 + ... + 1000 ≈ 500,000
Time: ~50 seconds (very slow!)
With KV Cache:
Total operations: 1000 (1 per token, only new)
Time: ~5 seconds (10x faster!)
How KV Cache Works: Detailed Walkthrough¶
Attention Computation Steps¶
Step 1: Prefill (Prompt Processing)¶
Input Prompt: "The quick brown fox"
(4 tokens in prompt)
Processing:
- ┌─────────────────────────────────────────┐
- For each token in prompt: │
- ┤
- Token 1 "The": │
- Compute Q, K, V (dim: batch, seq, d) │
- Cache: K[1], V[1] │
- Attention: 1×1 matrix │
│ │
- Token 2 "quick": │
- Compute Q, K, V │
- Cache: K[1:2], V[1:2] │
- Attention: 1×2 matrix │
│ │
- Token 3 "brown": │
- Compute Q, K, V │
- Cache: K[1:3], V[1:3] │
- Attention: 1×3 matrix │
│ │
- Token 4 "fox": │
- Compute Q, K, V │
- Cache: K[1:4], V[1:4] │
- Attention: 1×4 matrix │
- ┘
Result: Cache contains all K, V for prompt
Step 2: Decoding (Token Generation)¶
Generate next tokens one-by-one:
Token 5 (generate "jumps"):
- Input: Only new token "?" (1 token)
- Compute: Q[5], K[5], V[5]
- Retrieve: K[1:4], V[1:4] from cache
- Combine: K_full = [K[1:4], K[5]] (length 5)
- V_full = [V[1:4], V[5]] (length 5)
- Attention: Q[5] @ K_full^T / √d_k → (1×5) scores
- Update Cache: K[1:5], V[1:5]
- Output: token for "jumps"
Token 6 (generate next):
- Input: Only new token "?" (1 token)
- Compute: Q[6], K[6], V[6]
- Retrieve: K[1:5], V[1:5] from cache
- Combine: K_full = [K[1:5], K[6]] (length 6)
- V_full = [V[1:5], V[6]] (length 6)
- Attention: Q[6] @ K_full^T / √d_k → (1×6) scores
- Update Cache: K[1:6], V[1:6]
- Output: next token
Memory Consumption Analysis¶
KV Cache Size Calculation¶
For a transformer model:
KV cache per token =
2 × (num_layers × num_heads × head_dim)
× (batch_size × seq_length)
× (bytes_per_value)
Formula:
KV_cache_size = 2 × L × H × D × B × S × dtype_size
Where:
- L = number of layers
- H = number of attention heads
- D = head dimension
- B = batch size
- S = sequence length
- dtype_size = 2 (float16) or 4 (float32) bytes
Example: Llama 2 7B Model¶
Model: Llama 2 7B
Hidden dim: 4096
Num heads: 32
Head dim: 128
Num layers: 32
Batch size: 1 request
Sequence length: 2048 tokens
Data type: float16 (2 bytes)
Calculation:
KV_cache_per_layer = 2 × 32 × 128 × 1 × 2048 × 2 bytes
= 2 × 32 × 128 × 1 × 2048 × 2 bytes
= 33,554,432 bytes ≈ 32MB per layer
Total KV_cache = 32 layers × 32MB = 1024MB = 1GB
For 4 concurrent requests (batch_size=4):
Total KV_cache = 4 × 1GB = 4GB
For longer sequence (4096 tokens):
KV_cache per request = 2GB
For 4 requests = 8GB
Memory Breakdown: Full Model vs KV Cache¶
Llama 2 7B (float16):
- Model weights: 7B × 2 bytes = 14GB (loaded once)
- Activations during forward pass: ~2GB (temporary)
- KV Cache per request: 1GB (grows with sequence length)
│
Total for inference:
- Single request (2048 tokens): 14GB + 1GB = 15GB
- 4 concurrent requests: 14GB + 4GB = 18GB
- 8 concurrent requests: 14GB + 8GB = 22GB
KV Cache dominates once batch size is large!
Comparison: With vs Without KV Cache¶
Model: Llama 2 7B
Sequence length: 2048 tokens
WITHOUT KV Cache (recompute every step):
- Step 1: Compute K, V for 1 token
- Step 2: Compute K, V for 2 tokens (recompute all!)
- Step 3: Compute K, V for 3 tokens (recompute all!)
- ...
- Step 2048: Compute K, V for 2048 tokens
Total computations: 1+2+3+...+2048 = ~2M operations
Memory: Temporary (stream out immediately)
WITH KV Cache:
- Step 1: Compute K, V for 1 token, cache it
- Step 2: Compute K, V for 1 token, cache it
- Step 3: Compute K, V for 1 token, cache it
- ...
- Step 2048: Compute K, V for 1 token, cache it
Total computations: 2048 operations (2000x fewer!)
Memory: 1GB persistent (growing cache)
Trade: Extra 1GB memory → 2000x compute reduction!
Worth it? YES! Compute is the bottleneck in inference.
Implementation Details¶
Data Structure¶
class KVCache:
def __init__(self,
num_layers: int,
num_heads: int,
head_dim: int,
max_seq_len: int,
batch_size: int = 1):
# Pre-allocate buffers
self.key_cache = torch.zeros(
(batch_size, num_layers, max_seq_len, num_heads, head_dim),
dtype=torch.float16
)
self.value_cache = torch.zeros(
(batch_size, num_layers, max_seq_len, num_heads, head_dim),
dtype=torch.float16
)
# Track sequence lengths per request in batch
self.seq_lengths = torch.zeros(batch_size, dtype=torch.long)
def update(self,
layer_idx: int,
k_new: torch.Tensor, # (batch, 1, num_heads, head_dim)
v_new: torch.Tensor, # (batch, 1, num_heads, head_dim)
positions: torch.Tensor): # Where to place in cache
"""
Add new K, V to cache
"""
batch_size = k_new.shape[0]
seq_len = positions.max().item() + 1
for b in range(batch_size):
pos = positions[b].item()
self.key_cache[b, layer_idx, pos] = k_new[b, 0]
self.value_cache[b, layer_idx, pos] = v_new[b, 0]
self.seq_lengths[b] = seq_len
def get(self,
layer_idx: int,
batch_idx: int = None) -> tuple:
"""
Retrieve cached K, V up to current position
"""
if batch_idx is not None:
seq_len = self.seq_lengths[batch_idx].item()
return (
self.key_cache[batch_idx, layer_idx, :seq_len],
self.value_cache[batch_idx, layer_idx, :seq_len]
)
else:
# Return for all batches
return self.key_cache[:, layer_idx], self.value_cache[:, layer_idx]
Attention with KV Cache¶
def attention_with_kv_cache(
query: torch.Tensor, # (batch, 1, num_heads, head_dim)
key_new: torch.Tensor, # (batch, 1, num_heads, head_dim)
value_new: torch.Tensor, # (batch, 1, num_heads, head_dim)
kv_cache: KVCache,
layer_idx: int,
positions: torch.Tensor # Where in sequence we are
) -> torch.Tensor:
"""
Compute attention using KV cache
"""
batch_size, _, num_heads, head_dim = query.shape
# Get full K, V from cache (without new tokens)
k_cached, v_cached = kv_cache.get(layer_idx)
# Get sequence length from cache
seq_len = kv_cache.seq_lengths.max().item()
# Concatenate: [cached K, new K]
k_full = torch.cat(
[k_cached[:, :seq_len], key_new],
dim=1
) # (batch, seq_len+1, num_heads, head_dim)
v_full = torch.cat(
[v_cached[:, :seq_len], value_new],
dim=1
) # (batch, seq_len+1, num_heads, head_dim)
# Standard attention computation
scores = torch.matmul(query, k_full.transpose(-2, -1)) / math.sqrt(head_dim)
# scores: (batch, num_heads, 1, seq_len+1)
weights = torch.softmax(scores, dim=-1)
output = torch.matmul(weights, v_full)
# output: (batch, num_heads, 1, head_dim)
# Update cache with new K, V
kv_cache.update(layer_idx, key_new, value_new, positions)
return output
Two-Phase Generation with KV Cache¶
def generate_with_kv_cache(
model: nn.Module,
prompt_ids: torch.Tensor,
max_new_tokens: int,
kv_cache: KVCache
) -> torch.Tensor:
"""
Two-phase generation: prefill + decode
"""
# Phase 1: Prefill (process entire prompt)
print("Phase 1: Prefill (process prompt)")
input_ids = prompt_ids
position = 0
with torch.no_grad():
for token_idx in range(prompt_ids.shape[1]):
# Process one token at a time (or batch if possible)
output, _ = model(
input_ids=input_ids[:, token_idx:token_idx+1],
kv_cache=kv_cache,
position=torch.tensor([position])
)
position += 1
# Phase 2: Decode (generate tokens one-by-one)
print("Phase 2: Decode (generate new tokens)")
generated_ids = []
current_token = prompt_ids[:, -1:] # Start with last prompt token
for _ in range(max_new_tokens):
with torch.no_grad():
output, _ = model(
input_ids=current_token,
kv_cache=kv_cache,
position=torch.tensor([position])
)
position += 1
# Get token with highest probability
next_token = output.argmax(dim=-1)
generated_ids.append(next_token)
current_token = next_token
return torch.cat(generated_ids, dim=1)
Optimization Techniques¶
1. KV Cache Quantization¶
Reduce precision to save memory:
class QuantizedKVCache:
"""Store KV cache in int8 instead of float16"""
def __init__(self, ...):
# Store as int8 instead of float16
self.key_cache = torch.zeros(..., dtype=torch.int8)
self.value_cache = torch.zeros(..., dtype=torch.int8)
# Quantization parameters
self.key_scale = torch.ones(...) # Scaling factors
self.value_scale = torch.ones(...)
def update(self, layer_idx, k_new, v_new, positions):
"""Quantize before storing"""
# Quantize K to int8
k_scale = k_new.abs().max() / 127.0
k_quantized = (k_new / k_scale).to(torch.int8)
# Store quantized value and scale
self.key_cache[..., positions] = k_quantized
self.key_scale[..., positions] = k_scale
def get(self, layer_idx):
"""Dequantize when retrieving"""
k = self.key_cache[layer_idx].float()
k = k * self.key_scale[layer_idx]
return k
# Memory savings
# float16: 2 bytes per value
# int8: 1 byte per value + scale (minimal)
# Total: ~50% memory reduction!
2. Multi-Query Attention (MQA)¶
Use fewer key/value heads than query heads:
# Standard attention: K, V have same num_heads as Q
# (e.g., 32 heads, so 32 K/V head sets)
# Multi-Query Attention: K, V have 1 head (shared)
# Only 1 K/V head set for all 32 queries!
# Memory reduction for KV cache:
# Original: 32 × (K size)
# MQA: 1 × (K size) = 32x smaller KV cache!
class MultiQueryAttention(nn.Module):
def __init__(self, num_q_heads, num_kv_heads, head_dim):
self.num_q_heads = num_q_heads # 32
self.num_kv_heads = num_kv_heads # 1 or 8
self.head_dim = head_dim
def forward(self, x):
q = self.q_proj(x) # (batch, seq, num_q_heads * head_dim)
k = self.k_proj(x) # (batch, seq, num_kv_heads * head_dim)
v = self.v_proj(x) # (batch, seq, num_kv_heads * head_dim)
# Reshape and repeat K, V heads to match Q heads
q = q.view(batch, seq, self.num_q_heads, self.head_dim)
k = k.view(batch, seq, self.num_kv_heads, self.head_dim)
k = k.repeat(1, 1, self.num_q_heads // self.num_kv_heads, 1)
# Now k has same shape as q
# Standard attention
scores = q @ k.transpose(-2, -1) / sqrt(self.head_dim)
...
MQA models: Llama 2 70B, Mistral 7B
3. Grouped Query Attention (GQA)¶
Middle ground between standard and MQA:
Standard Attention: 32 Q heads, 32 K heads, 32 V heads
Multi-Query: 32 Q heads, 1 K head, 1 V head (extreme)
Grouped Query: 32 Q heads, 8 K heads, 8 V heads (balanced)
Memory comparison:
- Standard: Full KV cache (baseline)
- Grouped (8): 4x smaller KV cache (32/8)
- Multi-Query: 32x smaller KV cache (32/1)
Trade: Quality vs KV cache size
GQA provides good balance!
4. Prefix Caching (Paged Attention)¶
Share KV cache for common prompts:
# Scenario: 100 requests with same system prompt
system_prompt = "You are a helpful AI assistant..."
# Without caching:
# Each request: Full KV cache = 100 requests × cache_size
# With prefix caching:
# Store system prompt KV once (shared)
# Each request adds its own continuation KV
class PrefixKVCache:
def __init__(self):
self.prefix_cache = {} # Hash of prefix → KV cache
def get_or_create(self, prefix_tokens, prefix_hash):
"""Reuse cache for identical prefixes"""
if prefix_hash not in self.prefix_cache:
# Compute and cache prefix KV
self.prefix_cache[prefix_hash] = self._compute_prefix_kv(
prefix_tokens
)
return self.prefix_cache[prefix_hash]
# Memory savings:
# With 100 requests sharing 50-token prefix:
# Saved: 99 × prefix_cache_size
Memory Growth Analysis¶
Sequence Length Dependency¶
KV Cache Memory = 2 × num_layers × num_heads × head_dim × seq_len
Grows LINEARLY with sequence length!
For Llama 2 7B:
- ┌─────────────────┬──────────────┐
- Sequence Length │ KV Cache │
- ┼──────────────┤
- 512 │ 256 MB │
- 1024 │ 512 MB │
- 2048 │ 1 GB │
- 4096 │ 2 GB │
- 8192 │ 4 GB │
- ┴──────────────┘
Doubling sequence length = Doubling KV cache
Multi-Request Memory Management¶
Scenario: A100 GPU with 40GB memory
Model weights: 14GB
Activations: ~2GB
Available for KV cache: ~24GB
Each 7B request at 2048 tokens: 1GB KV cache
→ Can serve: 24 concurrent requests
But batch size depends on:
- Sequence length variations
- Model size
- Hardware memory
Challenges and Solutions¶
Challenge 1: Memory Bottleneck¶
Problem: As batch size increases, KV cache dominates memory
- Small batch (1-2): Model weights are bottleneck
- Large batch (8-16): KV cache is bottleneck
- Result: Limited throughput despite fast compute
Solution:
✓ KV cache quantization (50% reduction)
✓ Multi-query attention (32x reduction)
✓ Paged attention (non-contiguous allocation)
✓ Offload to CPU (trade latency for memory)
Challenge 2: Variable Sequence Length¶
Problem: Requests have different lengths
- Request 1: 100 tokens
- Request 2: 2000 tokens
- Request 3: 500 tokens
Traditional: Pre-allocate max (2000) for all → waste
PagedAttention: Allocate only what's needed → efficient
Challenge 3: Attention Computation Cost¶
With KV cache, bottleneck shifts:
- Compute: O(seq_len) for KV projection
- Memory: O(seq_len²) for attention scores/weights
- Result: Long sequences still slow (but faster than before)
For seq_len=2048:
- Attention matrix: 2048 × 2048 = 4M values
- Memory: 16 MB per head
- With 32 heads: 512 MB temporary
KV Cache Variants & Modern Techniques¶
1. Sliding Window Attention¶
Only cache recent tokens:
class SlidingWindowKVCache:
def __init__(self, window_size: int):
self.window_size = window_size # e.g., 512 tokens
def get(self, seq_len: int):
"""Only keep last window_size tokens"""
start = max(0, seq_len - self.window_size)
return self.key_cache[:, start:seq_len], \
self.value_cache[:, start:seq_len]
# Models using sliding window:
# - Mistral: 4096 token window
# - Llama 2: Can use sliding window during inference
# Memory reduction: Cache_size = 2 × model_params × window_size
# Instead of: Cache_size = 2 × model_params × full_seq_length
2. Sparse Attention Patterns¶
Only cache relevant tokens:
# Patterns for long sequences
# Strided attention: Attend to every Nth token
# Local attention: Attend to nearby tokens only
# Hierarchical: Different patterns per layer
# Reduces KV cache from O(N) to O(N/k) for k-stride
3. Recency-Based Pruning¶
Drop old tokens that are less important:
class PruningKVCache:
def __init__(self, pruning_ratio: float = 0.5):
self.pruning_ratio = pruning_ratio
def prune(self, importance_scores):
"""Remove least important tokens"""
threshold = importance_scores.quantile(self.pruning_ratio)
mask = importance_scores > threshold
# Keep only important tokens
self.key_cache = self.key_cache[:, mask]
self.value_cache = self.value_cache[:, mask]
# Aggressive: Remove 50% of cache with <5% quality loss
# Saves: 50% memory at cost of slight quality degradation
Real-World Performance Impact¶
Case Study: Chat Application¶
Scenario: Serve 100 concurrent chat requests
Model: Llama 2 13B
Average conversation length: 1000 tokens
Hardware: 8x A100 GPUs (320GB total)
WITHOUT KV Cache:
- Per request: ~26GB (14GB model + 12GB recomputation)
- Concurrent: Max 2-3 requests per GPU
- Total throughput: ~20 requests/sec
- Cost: $80,000/month
WITH KV Cache:
- Per request: ~14GB model + 2GB KV cache = 16GB
- Concurrent: 8-10 requests per GPU
- Total throughput: ~100 requests/sec (5x better!)
- Cost: $20,000/month (75% savings!)
WITH KV Cache + Quantization:
- KV cache: 1GB (quantized int8)
- Concurrent: 12-14 requests per GPU
- Total throughput: ~150 requests/sec (7.5x better!)
- Cost: $12,000/month (85% savings!)
Best Practices¶
✅ Do's¶
- Always use KV cache in inference (unless memory is unlimited)
- Pre-allocate space based on max sequence length
- Use quantization to reduce memory (int8, int4)
- Monitor cache size relative to available GPU memory
- Implement cache clearing for request batches
- Test with realistic sequence length distributions
- Profile bottlenecks (compute vs memory)
- Use multi-query attention for large scale serving
❌ Don'ts¶
- ❌ Recompute K, V every token (kills performance)
- ❌ Pre-allocate full sequence length if unnecessary
- ❌ Use float32 for cache (float16 is sufficient)
- ❌ Forget to clear cache between batches
- ❌ Ignore memory fragmentation (use paged attention)
- ❌ Cache indefinitely (implement eviction policies)
- ❌ Mix different quantization levels in same batch
- ❌ Disable cache without measuring impact
Comparison: With vs Without KV Cache¶
| Metric | Without KV | With KV | Improvement |
|---|---|---|---|
| Compute per token | O(seq_len) | O(1) | 1000x faster |
| Memory per token | O(seq_len) | O(1) | 1000x less |
| Inference time (2048 tokens) | 50s | 5s | 10x faster |
| GPU utilization | Low (compute-bound) | Medium | Better |
| Batch throughput | 0.1 req/s | 10 req/s | 100x better |
| Memory bottleneck | Compute | KV cache | Manageable |
Implementation in Popular Frameworks¶
PyTorch (Manual)¶
# See implementation details section above
Hugging Face Transformers¶
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
# KV cache is automatically used
inputs = tokenizer("Hello", return_tensors="pt")
# Generation with KV cache (default)
outputs = model.generate(
inputs["input_ids"],
use_cache=True, # Enable KV cache (default)
max_length=256
)
# Without cache (slow, for comparison)
outputs_no_cache = model.generate(
inputs["input_ids"],
use_cache=False, # Disable KV cache
max_length=256
)
vLLM¶
from vllm import LLM, SamplingParams
# vLLM handles KV cache automatically and efficiently
llm = LLM(model="meta-llama/Llama-2-7b-hf")
outputs = llm.generate(
["The quick brown fox"],
SamplingParams(max_tokens=256)
)
# KV cache is managed by vLLM's PagedAttention
Future Directions¶
Research Areas¶
- KV Cache Compression: Lossy compression without quality loss
- Adaptive Cache Pruning: Remove unimportant tokens dynamically
- Speculative Decoding: Efficiently handle speculative tokens
- CPU-GPU Offloading: Hybrid memory management
- Distributed KV Cache: Multi-GPU cache management
Emerging Techniques¶
- Token Merging: Merge similar tokens to reduce cache
- Attention Pattern Prediction: Pre-compute likely attention patterns
- Hierarchical Caching: Multi-level cache (hot/cold)
Key Takeaways¶
🔑 KV Cache reduces inference compute from O(N²) to O(N)
💾 Trade: Extra memory for dramatically reduced computation
⚡ 10x faster token generation (5 seconds vs 50 seconds)
📊 Memory scales with sequence length (linear growth)
🎯 Critical optimization for production LLM serving
🔬 Still an active research area with new techniques emerging
Further Reading¶
- Transformer Attention Is All You Need: Original attention mechanism
- KV Cache Quantization Papers: Recent work on memory reduction
- vLLM PagedAttention: Efficient cache management
- Multi-Query Attention: Ainslie et al., 2023
- Grouped Query Attention: Ainslie et al., 2023