KV Cache

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)
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
# int8
# Total

2. Multi-Query Attention (MQA)

Use fewer key/value heads than query heads:

# Standard attention
# (e.g., 32 heads, so 32 K/V head sets)

# Multi-Query Attention
# Only 1 K/V head set for all 32 queries!

# Memory reduction for KV cache:
# Original
# MQA

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
system_prompt = "You are a helpful AI assistant..."

# Without caching:
# Each request

# 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

---

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
# - Llama 2

# Memory reduction
# Instead of

2. Sparse Attention Patterns

Only cache relevant tokens:

# Patterns for long sequences
# Strided attention
# Local attention
# Hierarchical

# Reduces KV cache from O(N) to O(N/k) for k-stride
```sql

### 3. Recency-Based Pruning

Drop old tokens that are less important:

```python
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
# Saves

-

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

1. Always use KV cache in inference (unless memory is unlimited)
2. Pre-allocate space based on max sequence length
3. Use quantization to reduce memory (int8, int4)
4. Monitor cache size relative to available GPU memory
5. Implement cache clearing for request batches
6. Test with realistic sequence length distributions
7. Profile bottlenecks (compute vs memory)
8. Use multi-query attention for large scale serving

Don'ts

1. Recompute K, V every token (kills performance)
2. Pre-allocate full sequence length if unnecessary
3. Use float32 for cache (float16 is sufficient)
4. Forget to clear cache between batches
5. Ignore memory fragmentation (use paged attention)
6. Cache indefinitely (implement eviction policies)
7. Mix different quantization levels in same batch
8. 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

1. KV Cache Compression: Lossy compression without quality loss
2. Adaptive Cache Pruning: Remove unimportant tokens dynamically
3. Speculative Decoding: Efficiently handle speculative tokens
4. CPU-GPU Offloading: Hybrid memory management
5. 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

-

Related Notes

• Pagedattention
• 01 Vllm
• Flash Attention
• 00 Attention Mechanisms
• Llm Inference Optimization