kv_cache.py
06_llm_inference/kv_cache.py · 240 lines · view on GitHub
"""
KV Cache from Scratch
Interview question: "How does KV caching work?"
KEY IMPROVEMENT OVER STANDARD INFERENCE:
STANDARD INFERENCE (WITHOUT KV CACHE):
- At each generation step, processes ENTIRE sequence from scratch
- Recomputes K and V for ALL previous tokens every time
- Example: Generating token 3 requires recomputing K, V for tokens 0, 1, 2
- Complexity: O(n³d) total, O(i²d) per step (where i is current sequence length)
KV CACHE INFERENCE (WITH KV CACHE):
- At each generation step, processes ONLY the new token
- Reuses cached K and V for all previous tokens
- Example: Generating token 3 only computes K, V for token 3, reuses cached K, V for tokens 0, 1, 2
- Complexity: O(n²d) total, O(id) per step
- Speedup: ~n× for sequences of length n
THE KEY CODE DIFFERENCE:
Standard (without cache):
input_ids = [token_0, token_1, ..., token_i] # Entire sequence
K = compute_K(input_ids) # Recomputes K for ALL tokens
V = compute_V(input_ids) # Recomputes V for ALL tokens
KV Cache (with cache):
input_ids = [token_i] # Only new token
K_new = compute_K(input_ids) # Only computes K for new token
V_new = compute_V(input_ids) # Only computes V for new token
K = concatenate([K_cache, K_new]) # Reuses cached K
V = concatenate([V_cache, V_new]) # Reuses cached V
See kv_cache_detailed.md for complete explanation with step-by-step comparisons!
"""
import numpy as np
from typing import Dict, List, Optional
class KVCache:
"""
KV Cache stores Key and Value matrices to avoid recomputation
Critical optimization for autoregressive generation
How it works:
1. First token: Compute full Q, K, V, cache K and V
2. Next tokens: Only compute Q for new token, reuse cached K, V
3. Result: Much faster generation
"""
def __init__(self, num_layers: int, num_heads: int, head_dim: int):
self.num_layers = num_layers
self.num_heads = num_heads
self.head_dim = head_dim
# Cache structure: {layer_idx: {'keys': [...], 'values': [...]}}
self.cache: Dict[int, Dict[str, List[np.ndarray]]] = {}
def initialize_layer(self, layer_idx: int):
"""Initialize cache for a layer"""
self.cache[layer_idx] = {
'keys': [],
'values': []
}
def update(self, layer_idx: int, keys: np.ndarray, values: np.ndarray):
"""
Update cache with new keys and values
THIS IS WHAT MAKES KV CACHE WORK:
- Stores computed K, V so they can be reused in future steps
- Without this, we'd have to recompute K, V every time
- With this, we compute K, V once, store them, reuse them
STANDARD (without cache):
- Doesn't store K, V
- Recomputes them every step
- Wasteful: same computation repeated many times
KV CACHE (with cache):
- Stores K, V after computing them
- Reuses stored values in future steps
- Efficient: each K, V computed only once
Args:
layer_idx: Which transformer layer
keys: New key vectors (num_heads, head_dim) or (1, num_heads, head_dim)
values: New value vectors (num_heads, head_dim) or (1, num_heads, head_dim)
"""
if layer_idx not in self.cache:
self.initialize_layer(layer_idx)
# Handle different input shapes
if keys.ndim == 2:
keys = keys[np.newaxis, :, :] # Add batch dimension
if values.ndim == 2:
values = values[np.newaxis, :, :]
self.cache[layer_idx]['keys'].append(keys)
self.cache[layer_idx]['values'].append(values)
def get(self, layer_idx: int) -> Optional[Dict[str, np.ndarray]]:
"""Get cached keys and values for a layer"""
if layer_idx not in self.cache or len(self.cache[layer_idx]['keys']) == 0:
return None
# Concatenate all cached keys/values along sequence dimension
keys = np.concatenate(self.cache[layer_idx]['keys'], axis=0)
values = np.concatenate(self.cache[layer_idx]['values'], axis=0)
return {'keys': keys, 'values': values}
def clear(self):
"""Clear cache (start new sequence)"""
self.cache = {}
def get_cache_size(self) -> int:
"""Get total cache size in elements"""
total = 0
for layer_idx in self.cache:
for key in ['keys', 'values']:
for item in self.cache[layer_idx][key]:
total += item.size
return total
def attention_with_kv_cache(Q: np.ndarray, K_cache: Optional[np.ndarray],
V_cache: Optional[np.ndarray], K_new: np.ndarray,
V_new: np.ndarray, d_k: int) -> np.ndarray:
"""
Attention with KV cache
Only compute attention for new token, reuse cached K/V
THIS IS THE KEY OPTIMIZATION:
STANDARD (without cache):
- Would recompute K, V for ALL tokens: K_all = compute_K([token_0, ..., token_i])
- Processes entire sequence every time
- Redundant computation: recomputes K_0, V_0, K_1, V_1, ... every step
KV CACHE (with cache):
- Only computes K, V for NEW token: K_new = compute_K([token_i])
- Reuses cached K, V for previous tokens: K = [K_cache, K_new]
- No redundant computation: each K, V computed only once
THE CONCATENATION IS THE KEY:
K = concatenate([K_cache, K_new]) # This line reuses cached values!
V = concatenate([V_cache, V_new]) # This line reuses cached values!
This concatenation allows us to:
1. Use cached K, V for all previous tokens (no recomputation)
2. Only compute K, V for the new token
3. Combine them for attention computation
Args:
Q: Query for new token (num_heads, head_dim)
K_cache: Cached keys (seq_len-1, num_heads, head_dim) or None
V_cache: Cached values (seq_len-1, num_heads, head_dim) or None
K_new: New key (num_heads, head_dim) - ONLY computed for new token
V_new: New value (num_heads, head_dim) - ONLY computed for new token
d_k: Key dimension
"""
# If no cache, this is first token
if K_cache is None:
# First token: no cache yet, just use new K, V
K = K_new[np.newaxis, :, :] # (1, num_heads, head_dim)
V = V_new[np.newaxis, :, :]
else:
# THIS IS THE KEY OPTIMIZATION:
# Instead of recomputing K, V for all previous tokens,
# we concatenate cached (already computed) + new (just computed)
#
# Standard would do: K = compute_K([token_0, ..., token_i]) # Recomputes all!
# KV Cache does: K = [K_cache, K_new] # Reuses cache, only computes new!
K = np.concatenate([K_cache, K_new[np.newaxis, :, :]], axis=0)
V = np.concatenate([V_cache, V_new[np.newaxis, :, :]], axis=0)
# Compute attention scores
# Q: (num_heads, head_dim), K: (seq_len, num_heads, head_dim)
# We need to compute Q @ K^T for each head
seq_len = K.shape[0]
num_heads = Q.shape[0]
scores = np.zeros((num_heads, seq_len))
for head in range(num_heads):
scores[head] = Q[head] @ K[:, head, :].T / np.sqrt(d_k)
# Softmax
exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
attention_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
# Apply to values
output = np.zeros_like(Q)
for head in range(num_heads):
output[head] = attention_weights[head] @ V[:, head, :]
return output
# Usage Example
if __name__ == "__main__":
print("KV Cache Example")
print("=" * 60)
# Initialize cache
num_layers = 12
num_heads = 8
head_dim = 64
cache = KVCache(num_layers, num_heads, head_dim)
# Simulate generation: 3 tokens
for token_idx in range(3):
layer_idx = 0 # First layer
# Generate random K, V for this token
K_new = np.random.randn(num_heads, head_dim)
V_new = np.random.randn(num_heads, head_dim)
# Get cached K, V
cached = cache.get(layer_idx)
if cached is None:
print(f"Token {token_idx}: First token, no cache")
else:
print(f"Token {token_idx}: Using cache with {len(cache.cache[layer_idx]['keys'])} previous tokens")
# Compute attention (simplified - would use Q in real implementation)
Q = np.random.randn(num_heads, head_dim)
output = attention_with_kv_cache(
Q,
cached['keys'] if cached else None,
cached['values'] if cached else None,
K_new, V_new, head_dim
)
# Update cache
cache.update(layer_idx, K_new, V_new)
print(f"\nCache size: {cache.get_cache_size()} elements")
print(f"Memory saved: Instead of recomputing, we reuse cached K/V")