KV Caching Deep-Dive: Memory-Efficient Transformer Inference

# Naive generation (what NOT to do)
tokens = [prompt_tokens]
for step in range(max_tokens):
    logits = model(tokens)  # Recomputes K,V for ALL tokens!
    next_token = sample(logits[-1])
    tokens.append(next_token)

# From nanochat/gpt.py CausalSelfAttention.forward() lines 79-91
def forward(self, x, cos_sin, kv_cache):
    B, T, C = x.size()
    
    # Project input to Q, K, V
    q = self.c_q(x).view(B, T, self.n_head, self.head_dim)      # Cost: O(T × d²)
    k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)  # Cost: O(T × d²)
    v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)  # Cost: O(T × d²)
    
    # Apply rotary embeddings + normalization
    q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
    q, k = norm(q), norm(k)
    
    # Transpose for attention computation
    q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
    
    # Attention computation
    # q: (B, H, T, D), k: (B, H, T, D), v: (B, H, T, D)
    scores = q @ k.transpose(-2, -1)  # (B, H, T, T) - Cost: O(T²)
    attn = softmax(scores / sqrt(D))   # Cost: O(T²)
    y = attn @ v                       # (B, H, T, D) - Cost: O(T²)

# From nanochat/engine.py lines 56-66
class KVCache:
    def __init__(self, batch_size, num_heads, seq_len, head_dim, num_layers):
        # Shape: (L, 2, B, H, T, D)
        #         ↑  ↑  ↑  ↑  ↑  ↑
        #         │  │  │  │  │  └─ head_dim (64-128 typically)
        #         │  │  │  │  └──── sequence length (max cache capacity)
        #         │  │  │  └─────── num_heads (or num_kv_heads for MQA)
        #         │  │  └────────── batch_size
        #         │  └───────────── 2 = [keys, values]
        #         └──────────────── num_layers
        self.kv_shape = (num_layers, 2, batch_size, num_heads, seq_len, head_dim)
        self.kv_cache = None  # Lazy initialization
        self.pos = 0          # Current position in sequence

# From nanochat/engine.py lines 101-104
def insert_kv(self, layer_idx, k, v):
    # Lazy initialize on first insert
    if self.kv_cache is None:
        self.kv_cache = torch.empty(self.kv_shape, dtype=k.dtype, device=k.device)

# From nanochat/engine.py lines 106-114
def insert_kv(self, layer_idx, k, v):
    B, H, T_add, D = k.size()
    t0, t1 = self.pos, self.pos + T_add
    
    # Dynamically grow the cache if needed
    if t1 > self.kv_cache.size(4):
        t_needed = t1 + 1024              # Current need + 1024 buffer
        t_needed = (t_needed + 1023) & ~1023  # Round up to nearest 1024
        current_shape = list(self.kv_cache.shape)
        current_shape[4] = t_needed
        self.kv_cache.resize_(current_shape)

# From nanochat/engine.py lines 115-124
def insert_kv(self, layer_idx, k, v):
    # ... (growth logic above)
    
    # Insert new k,v into cache at current position
    self.kv_cache[layer_idx, 0, :, :, t0:t1] = k
    self.kv_cache[layer_idx, 1, :, :, t0:t1] = v
    
    # Return views of ALL cached k,v up to current position
    key_view = self.kv_cache[layer_idx, 0, :, :, :t1]
    value_view = self.kv_cache[layer_idx, 1, :, :, :t1]
    
    # Advance position after last layer processes
    if layer_idx == self.kv_cache.size(0) - 1:
        self.pos = t1
    
    return key_view, value_view

# From nanochat/engine.py lines 180-192
def generate(self, tokens, num_samples=1, max_tokens=None, ...):
    # Step 1: Prefill with batch size 1
    m = self.model.config
    kv_cache_prefill = KVCache(
        batch_size=1,
        seq_len=len(tokens),
        num_heads=m.n_kv_head,
        head_dim=m.n_embd // m.n_head,
        num_layers=m.n_layer
    )
    
    ids = torch.tensor([tokens], dtype=torch.long, device=device)
    logits = self.model.forward(ids, kv_cache=kv_cache_prefill)
    logits = logits[:, -1, :]  # Take logits at last position only
    next_ids = sample_next_token(logits, rng, temperature, top_k)

# From nanochat/gpt.py lines 104-107
if kv_cache is None or Tq == Tk:
    # Tq == Tk means prefill (processing all tokens at once)
    # Use PyTorch's efficient causal attention implementation
    y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# From nanochat/engine.py lines 225-229
while True:
    # Forward pass with single token
    logits = self.model.forward(ids, kv_cache=kv_cache_decode)  # ids: (B, 1)
    logits = logits[:, -1, :]  # (B, vocab_size)
    next_ids = sample_next_token(logits, rng, temperature, top_k)
    sampled_tokens = next_ids[:, 0].tolist()

# From nanochat/gpt.py lines 108-111
elif Tq == 1:
    # Single query attending to all cached keys/values
    # No causal mask needed (query is at the end of sequence)
    y = F.scaled_dot_product_attention(q, k, v, is_causal=False)

# From nanochat/gpt.py lines 113-121
else:
    # Tq > 1 AND Tq < Tk: Processing a chunk during decode
    # Example: 50 cached tokens, adding 10 new tokens
    attn_mask = torch.zeros((Tq, Tk), dtype=torch.bool, device=q.device)
    prefix_len = Tk - Tq
    
    # All queries attend to all cached tokens (prefix)
    attn_mask[:, :prefix_len] = True
    
    # Causal attention within the new chunk
    attn_mask[:, prefix_len:] = torch.tril(torch.ones((Tq, Tq), dtype=torch.bool, device=q.device))
    
    y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)

# From nanochat/engine.py lines 183-202
# Step 1: Prefill once with batch size 1
kv_cache_prefill = KVCache(batch_size=1, seq_len=len(tokens), ...)
logits = self.model.forward(ids, kv_cache=kv_cache_prefill)
 
# Step 2: Replicate cache for batch generation
kv_length_hint = (len(tokens) + max_tokens) if max_tokens else model.config.sequence_len
kv_cache_decode = KVCache(
    batch_size=num_samples,  # ← Expand to N samples
    seq_len=kv_length_hint,
    **kv_model_kwargs
)
kv_cache_decode.prefill(kv_cache_prefill)  # Replicate cached data
del kv_cache_prefill  # Free memory
 
# Step 3: Decode in parallel for all N samples
for token_column, token_masks in self.generate(...):
    # Each sample generates independently
    ...

# From nanochat/engine.py lines 74-99
def prefill(self, other):
    """
    Prefill this cache with data from another cache.
    Optionally expand along batch dimension.
    """
    assert self.kv_cache is None, "Cannot prefill non-empty cache"
    assert other.kv_cache is not None, "Cannot prefill with None cache"
    
    # Validate dimensions (shapes must be compatible)
    for ix, (dim1, dim2) in enumerate(zip(self.kv_shape, other.kv_shape)):
        if ix in [0, 1, 3, 5]:  # num_layers, 2, num_heads, head_dim
            assert dim1 == dim2, f"Dimension {ix} mismatch: {dim1} != {dim2}"
        elif ix == 2:  # batch_size can expand
            assert dim1 >= dim2 or dim2 == 1, f"Batch dim mismatch"
        elif ix == 4:  # seq_len: target must be >= source
            assert dim1 >= dim2, f"Seq len mismatch"
    
    # Initialize and copy (with broadcasting)
    dtype, device = other.kv_cache.dtype, other.kv_cache.device
    self.kv_cache = torch.empty(self.kv_shape, dtype=dtype, device=device)
    self.kv_cache[:, :, :, :, :other.pos, :] = other.kv_cache
    self.pos = other.pos

# Source cache: (L, 2, 1, H, T, D)  ← batch_size = 1
# Target cache: (L, 2, N, H, T, D)  ← batch_size = N
# Assignment triggers PyTorch broadcasting: dimension 2 expands from 1 → N

# Standard MHA (not in nanochat, for comparison)
n_head = 10
head_dim = 128
 
# Separate projections for each head
c_q = nn.Linear(n_embd, n_head * head_dim)  # 10 query heads
c_k = nn.Linear(n_embd, n_head * head_dim)  # 10 key heads
c_v = nn.Linear(n_embd, n_head * head_dim)  # 10 value heads

# From nanochat/gpt.py lines 68-76
class CausalSelfAttention(nn.Module):
    def __init__(self, config, layer_idx):
        self.n_head = config.n_head          # e.g., 10 query heads
        self.n_kv_head = config.n_kv_head    # e.g., 1 shared KV head
        
        self.c_q = nn.Linear(n_embd, n_head * head_dim, bias=False)       # 10 heads
        self.c_k = nn.Linear(n_embd, n_kv_head * head_dim, bias=False)   # 1 head only!
        self.c_v = nn.Linear(n_embd, n_kv_head * head_dim, bias=False)   # 1 head only!

# From nanochat/gpt.py lines 99-101
# After retrieving k,v from cache (shape: B, n_kv_head, T, D)
nrep = self.n_head // self.n_kv_head  # e.g., 10 // 1 = 10
k, v = repeat_kv(k, nrep), repeat_kv(v, nrep)
# Result shape: (B, n_head, T, D)

# From nanochat/gpt.py lines 52-61
def repeat_kv(x, n_rep):
    """Replicate KV heads to match number of query heads."""
    if n_rep == 1:
        return x  # No replication needed
    
    bs, n_kv_heads, slen, head_dim = x.shape
    return (
        x[:, :, None, :, :]  # Add dimension: (B, KV, 1, T, D)
        .expand(bs, n_kv_heads, n_rep, slen, head_dim)  # Expand: (B, KV, nrep, T, D)
        .reshape(bs, n_kv_heads * n_rep, slen, head_dim)  # Merge: (B, KV*nrep, T, D)
    )

# From nanochat/gpt.py GPTConfig lines 31-32
n_head: int = 10      # Query heads
n_kv_head: int = 5    # Key/Value heads (GQA when 1 < n_kv_head < n_head)

cache_bytes = (num_layers × 2 × batch_size × num_kv_heads 
               × seq_len × head_dim × sizeof(bfloat16))
            = L × 2 × B × H_kv × T × D × 2

L = 20 layers
H_kv = 10 KV heads (1:1 with query heads, no MQA)
D = 128 head_dim
T = 2048 max sequence length
B = 1 batch size
 
cache_size = 20 × 2 × 1 × 10 × 2048 × 128 × 2 bytes
           = 210 MB per sample
 
With MQA (H_kv = 1):
cache_size = 20 × 2 × 1 × 1 × 2048 × 128 × 2 bytes
           = 21 MB per sample (10× reduction!)

import torch
from nanochat.engine import KVCache, Engine
from nanochat.checkpoint_manager import load_model
from nanochat.tokenizer import get_tokenizer
 
# Load model and tokenizer
device = torch.device("cuda")
model, tokenizer, meta = load_model("base", device, phase="eval")
engine = Engine(model, tokenizer)
 
# Prepare prompt
prompt = "The capital of France is"
tokens = tokenizer.encode(prompt, prepend="<|bos|>")
 
print(f"Prompt: {prompt}")
print(f"Prompt tokens: {len(tokens)}")
 
# Generate with KV caching (3 samples)
print("\nGenerating 3 samples with KV caching...")
torch.cuda.synchronize()
import time
t0 = time.time()
 
results, masks = engine.generate_batch(
    tokens,
    num_samples=3,
    max_tokens=50,
    temperature=0.8,
    top_k=50,
    seed=42
)
 
torch.cuda.synchronize()
t1 = time.time()
 
# Print results
for i, (result, mask) in enumerate(zip(results, masks)):
    text = tokenizer.decode(result)
    num_forced = sum(1 - m for m in mask)  # Count forced tokens (calculator tool)
    print(f"\nSample {i+1}:")
    print(f"  Text: {text}")
    print(f"  Tokens: {len(result)} ({num_forced} forced by tool)")
 
print(f"\nTotal time: {(t1-t0)*1000:.1f}ms")
print(f"Time per token: {(t1-t0)*1000/sum(len(r) for r in results):.1f}ms")

# Detailed cache statistics
m = model.config
kv_cache = KVCache(
    batch_size=4,
    num_heads=m.n_kv_head,
    seq_len=2048,
    head_dim=m.n_embd // m.n_head,
    num_layers=m.n_layer
)
 
print(f"Initial cache shape: {kv_cache.kv_shape}")
print(f"Initial cache pos: {kv_cache.pos}")
print(f"Cache initialized: {kv_cache.kv_cache is not None}")
 
# After some generation
# kv_cache will have been filled via insert_kv() calls
 
if kv_cache.kv_cache is not None:
    print(f"\nAfter generation:")
    print(f"  Cache position: {kv_cache.pos}")
    print(f"  Cache capacity: {kv_cache.kv_cache.size(4)}")
    print(f"  Utilization: {kv_cache.pos / kv_cache.kv_cache.size(4) * 100:.1f}%")
    
    # Memory usage
    cache_bytes = kv_cache.kv_cache.numel() * kv_cache.kv_cache.element_size()
    print(f"  Cache memory: {cache_bytes / 1024**2:.1f} MB")
 
# Overall GPU memory
allocated = torch.cuda.memory_allocated() / 1024**3
reserved = torch.cuda.memory_reserved() / 1024**3
print(f"\nGPU memory: {allocated:.2f} GB allocated, {reserved:.2f} GB reserved")