An inference optimization technique that stores the Key and Value tensors from previous tokens in transformer attention layers, avoiding redundant recomputation during autoregressive generation — the critical mechanism that makes long-context LLM inference practical.
Imagine you’re writing a long essay. After writing each paragraph, you need to review the entire essay so far to ensure consistency.
Without KV cache: Every time you write a new sentence, you re-read the entire essay from the beginning. For a 100-page essay, this becomes incredibly slow.
With KV cache: You keep notes (summaries) of each paragraph as you write it. When writing a new sentence, you consult your notes instead of re-reading everything. Much faster!
KV cache does the same for AI. When generating text token-by-token, instead of recomputing attention for all previous tokens at each step, the model stores the “keys” and “values” (intermediate computations) from previous tokens. When generating the next token, it only computes attention for the new token, using the cached keys and values for all previous tokens.
This optimization is what makes generating long responses (thousands of tokens) feasible in reasonable time.
In transformer models, the self-attention mechanism computes relationships between all tokens in a sequence. For each new token generated, the model needs to attend to all previous tokens.
Without KV Cache:
With KV Cache:
What’s Stored in KV Cache:
Memory Requirements: For a model like Llama-2-70B:
KV Cache Optimizations:
1. Quantized KV Cache:
2. Sliding Window Attention:
3. PagedAttention (vLLM):
4. Multi-Query Attention (MQA):
5. Grouped-Query Attention (GQA):
KV cache directly impacts inference costs and capabilities:
Why It Matters:
Enterprise Implications:
Cost Example:
Popular Implementations:
A librarian helping patrons find books. Without KV cache, the librarian re-catalogs the entire library every time someone asks a question. With KV cache, the librarian maintains an index card catalog — when a new book arrives, they add one card instead of re-cataloging everything. The catalog (KV cache) grows with the library, but lookups stay fast.
# KV Cache in transformer inference (PyTorch)
import torch
import torch.nn as nn
class TransformerBlockWithKVCache(nn.Module):
def __init__(self, dim, num_heads):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.q_proj = nn.Linear(dim, dim)
self.k_proj = nn.Linear(dim, dim)
self.v_proj = nn.Linear(dim, dim)
self.o_proj = nn.Linear(dim, dim)
def forward(self, x, past_kv=None):
batch_size, seq_len, dim = x.shape
# Compute Q, K, V for current tokens
q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
# If past_kv exists, concatenate with cached K, V
if past_kv is not None:
past_k, past_v = past_kv
k = torch.cat([past_k, k], dim=1) # Cache all previous K
v = torch.cat([past_v, v], dim=1) # Cache all previous V
# Compute attention (simplified)
# q: [batch, seq_len, heads, head_dim]
# k, v: [batch, total_seq_len, heads, head_dim]
# Transpose for attention computation
q = q.transpose(1, 2) # [batch, heads, seq_len, head_dim]
k = k.transpose(1, 2) # [batch, heads, total_seq_len, head_dim]
v = v.transpose(1, 2) # [batch, heads, total_seq_len, head_dim]
# Attention scores
scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
attn_weights = torch.softmax(scores, dim=-1)
# Apply attention to values
output = torch.matmul(attn_weights, v)
output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, dim)
output = self.o_proj(output)
# Return output and current KV (to be cached for next step)
current_kv = (k, v)
return output, current_kv
# Usage example
block = TransformerBlockWithKVCache(dim=512, num_heads=8)
# First token
x1 = torch.randn(1, 1, 512) # batch=1, seq=1, dim=512
out1, kv1 = block(x1)
# Second token (uses cached KV from first token)
x2 = torch.randn(1, 1, 512)
out2, kv2 = block(x2, past_kv=kv1) # Pass cached KV
# Third token (uses cached KV from first two tokens)
x3 = torch.randn(1, 1, 512)
out3, kv3 = block(x3, past_kv=kv2)
print(f"KV cache size after 3 tokens: {kv3[0].shape}")
# Output: torch.Size([1, 8, 3, 64]) - cached K for all 3 tokens
Reality: KV cache eliminates redundant K, V computation, but attention scores still need to be recomputed for each new token. It’s a significant optimization, not a complete elimination.
Reality: KV cache consumes significant GPU memory, often the primary bottleneck for long-context inference. Managing KV cache memory is a critical engineering challenge.
Reality: KV cache size varies by model architecture. Models with MQA or GQA have much smaller KV caches than standard multi-head attention models.