Introduction to KV Cache Management for Long-Context Inference
The Key-Value (KV) cache is one of the most important architectural optimizations in modern transformer-based large language model (LLM) inference. As models are increasingly used with long contexts — sometimes exceeding 100,000 tokens — managing the KV cache efficiently becomes a critical engineering challenge. This tutorial provides a complete, practical guide to understanding, implementing, and optimizing KV cache management for long-context inference scenarios.
What Is the KV Cache?
In a transformer decoder, each token generation step requires computing attention over all previously generated tokens. Without caching, every new token would require recomputing the keys and values for the entire sequence from scratch, leading to quadratic computational overhead. The KV cache solves this by storing the projected key and value tensors for every previously processed token, so that during autoregressive generation, only the new token's key and value need to be computed and appended.
For each transformer layer, the cache stores two tensors: one for keys and one for values. Each has a shape of (batch_size, num_kv_heads, seq_len, head_dim). As the sequence grows, these tensors grow along the sequence dimension, which is the root cause of the memory challenges we will discuss.
The Math Behind the Memory Growth
Consider a model with the following parameters:
num_layers: number of transformer layersnum_kv_heads: number of key/value attention headshead_dim: dimension of each attention headseq_len: current sequence lengthbytes_per_element: typically 2 for FP16/BF16
The total KV cache memory in bytes is:
kv_cache_bytes = 2 * num_layers * num_kv_heads * head_dim * seq_len * batch_size * bytes_per_element
For example, a Llama-2-70B model processing a 32K token context with batch size 1 in FP16 requires roughly 5GB just for the KV cache — and that grows linearly with sequence length and batch size.
Why KV Cache Management Matters
As context windows expand to support use cases like document analysis, code repositories, and multi-turn conversations, the KV cache becomes the dominant memory consumer during inference. Poor KV cache management leads to several problems:
- Out-of-memory errors when sequences exceed available GPU memory
- Reduced throughput because fewer sequences can be batched simultaneously
- Increased latency due to memory bandwidth bottlenecks during attention computation
- Wasted memory from over-allocation or fragmentation
- Poor hardware utilization when cache transfers dominate compute time
Effective KV cache management is therefore essential for deploying long-context models in production environments where cost, latency, and throughput all matter.
How KV Cache Works in Practice
Basic KV Cache Implementation
Let's start with a simplified PyTorch implementation that demonstrates the core concept. This example shows how keys and values are cached and reused across generation steps:
import torch
import torch.nn as nn
import torch.nn.functional as F
class CachedAttentionLayer(nn.Module):
def __init__(self, hidden_dim, num_heads):
super().__init__()
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.head_dim = hidden_dim // num_heads
self.q_proj = nn.Linear(hidden_dim, hidden_dim)
self.k_proj = nn.Linear(hidden_dim, hidden_dim)
self.v_proj = nn.Linear(hidden_dim, hidden_dim)
self.o_proj = nn.Linear(hidden_dim, hidden_dim)
def forward(self, x, kv_cache=None, use_cache=True):
batch_size, seq_len, _ = x.shape
q = self.q_proj(x)
k = self.k_proj(x)
v = self.v_proj(x)
# Reshape to (batch, heads, seq, head_dim)
q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
if use_cache and kv_cache is not None:
# Append new keys and values to the cache
past_k, past_v = kv_cache
k = torch.cat([past_k, k], dim=2)
v = torch.cat([past_v, v], dim=2)
# Update cache for next step
new_cache = (k, v) if use_cache else None
# Scaled dot-product attention
attn_output = F.scaled_dot_product_attention(q, k, v)
# Reshape and project output
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(batch_size, seq_len, self.hidden_dim)
output = self.o_proj(attn_output)
return output, new_cache
class CachedTransformer(nn.Module):
def __init__(self, vocab_size, hidden_dim, num_heads, num_layers):
super().__init__()
self.embedding = nn.Embedding(vocab_size, hidden_dim)
self.layers = nn.ModuleList([
CachedAttentionLayer(hidden_dim, num_heads)
for _ in range(num_layers)
])
def forward(self, input_ids, caches=None, use_cache=True):
x = self.embedding(input_ids)
new_caches = []
for i, layer in enumerate(self.layers):
cache = caches[i] if caches is not None else None
x, new_cache = layer(x, kv_cache=cache, use_cache=use_cache)
new_caches.append(new_cache)
return x, new_caches
# Example usage: autoregressive generation with KV cache
def generate(model, input_ids, max_new_tokens=100):
model.eval()
caches = None
generated = input_ids
with torch.no_grad():
# Prefill phase: process the entire prompt
logits, caches = model(input_ids, caches=None, use_cache=True)
next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
generated = torch.cat([generated, next_token], dim=1)
# Decode phase: generate one token at a time using the cache
for _ in range(max_new_tokens - 1):
logits, caches = model(next_token, caches=caches, use_cache=True)
next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
generated = torch.cat([generated, next_token], dim=1)
return generated
Notice the two distinct phases: the prefill phase processes the entire prompt at once and populates the initial cache, while the decode phase processes one token at a time, appending to the cache at each step. This separation is fundamental to efficient inference.
Pre-Allocating the KV Cache
Dynamic tensor concatenation, as shown above, is simple but inefficient because it allocates new memory and copies data at every step. A better approach is to pre-allocate the maximum cache size and fill it incrementally:
class PreAllocatedKVCache:
def __init__(self, batch_size, num_layers, num_kv_heads,
max_seq_len, head_dim, dtype=torch.float16, device='cuda'):
self.max_seq_len = max_seq_len
self.num_layers = num_layers
self.current_len = 0
# Pre-allocate key and value tensors for all layers
# Shape: (num_layers, batch, num_kv_heads, max_seq_len, head_dim)
self.keys = torch.zeros(
num_layers, batch_size, num_kv_heads, max_seq_len, head_dim,
dtype=dtype, device=device
)
self.values = torch.zeros(
num_layers, batch_size, num_kv_heads, max_seq_len, head_dim,
dtype=dtype, device=device
)
def update(self, layer_idx, new_keys, new_values):
"""Append new keys/values for a specific layer."""
seq_len = new_keys.shape[2]
end = self.current_len + seq_len
self.keys[layer_idx, :, :, self.current_len:end, :] = new_keys
self.values[layer_idx, :, :, self.current_len:end, :] = new_values
def get(self, layer_idx):
"""Retrieve cached keys/values up to current length."""
return (
self.keys[layer_idx, :, :, :self.current_len, :],
self.values[layer_idx, :, :, :self.current_len, :]
)
def advance(self, tokens_added):
"""Update the current sequence length after appending tokens."""
self.current_len += tokens_added
def reset(self):
"""Clear the cache for a new sequence."""
self.current_len = 0
def get_memory_usage(self):
"""Return current memory usage in bytes."""
element_size = self.keys.element_size()
total_elements = 2 * self.keys.numel() # keys + values
# Only count used portion
used_elements = total_elements * (self.current_len / self.max_seq_len)
return int(used_elements * element_size)
This pre-allocation strategy avoids repeated memory allocation and copying, which is critical for achieving high throughput in production inference servers.
Memory Challenges with Long Contexts
The Linear Growth Problem
Unlike model weights, which are fixed, the KV cache grows linearly with sequence length. For a model like Llama-3-70B with 80 layers, 8 KV heads (with GQA), and a head dimension of 128, the per-token KV cache cost is:
per_token_bytes = 2 * 80 * 8 * 128 * 2 # 2 (K+V) * layers * kv_heads * head_dim * fp16
# = 327,680 bytes per token ≈ 320 KB per token
# For a 128K context:
total = 327680 * 131072 # ≈ 42.9 GB per sequence
This means a single 128K-token sequence can consume over 40GB of memory just for the KV cache, exceeding the capacity of most single-GPU configurations.
Batching and Throughput Trade-offs
The KV cache also directly impacts how many concurrent requests a server can handle. If each long-context request requires 40GB of cache, a GPU with 80GB of memory can only serve one or two requests at a time. This dramatically reduces throughput compared to short-context workloads where dozens of requests can be batched together.
Advanced KV Cache Management Techniques
1. PagedAttention
PagedAttention, introduced in the vLLM project, is inspired by operating system virtual memory and paging. Instead of allocating a contiguous block of memory for each sequence's KV cache, it divides the cache into fixed-size blocks (pages) that can be allocated and freed dynamically. This eliminates fragmentation and enables efficient memory sharing across sequences with common prefixes.
class PagedKVCache:
"""Simplified PagedAttention-style KV cache manager."""
def __init__(self, num_blocks, block_size, num_layers,
num_kv_heads, head_dim, dtype=torch.float16, device='cuda'):
self.block_size = block_size
self.num_layers = num_layers
self.num_kv_heads = num_kv_heads
self.head_dim = head_dim
# Block pool: shape (num_blocks, num_layers, 2, block_size, num_kv_heads, head_dim)
# The '2' dimension is for keys and values
self.block_pool = torch.zeros(
num_blocks, num_layers, 2, block_size,
num_kv_heads, head_dim,
dtype=dtype, device=device
)
self.free_blocks = list(range(num_blocks))
# Map sequence_id -> list of block indices
self.sequence_blocks = {}
def allocate_sequence(self, seq_id):
"""Allocate an initial block for a new sequence."""
if not self.free_blocks:
raise RuntimeError("No free blocks available")
block_idx = self.free_blocks.pop(0)
self.sequence_blocks[seq_id] = [block_idx]
return block_idx
def append_tokens(self, seq_id, new_keys, new_values):
"""Append new token KV data, allocating new blocks as needed."""
blocks = self.sequence_blocks[seq_id]
current_block = blocks[-1]
filled_in_block = self._get_filled_count(seq_id)
num_new_tokens = new_keys.shape[2]
for i in range(num_new_tokens):
if filled_in_block >= self.block_size:
# Allocate a new block
if not self.free_blocks:
raise RuntimeError("KV cache out of memory: no free blocks")
current_block = self.free_blocks.pop(0)
blocks.append(current_block)
filled_in_block = 0
# Write token data into the current block
for layer in range(self.num_layers):
self.block_pool[current_block, layer, 0, filled_in_block] = \
new_keys[layer, :, i, :]
self.block_pool[current_block, layer, 1, filled_in_block] = \
new_values[layer, :, i, :]
filled_in_block += 1
def free_sequence(self, seq_id):
"""Release all blocks held by a completed sequence."""
for block_idx in self.sequence_blocks[seq_id]:
self.free_blocks.append(block_idx)
del self.sequence_blocks[seq_id]
def _get_filled_count(self, seq_id):
"""Calculate how many tokens are in the last block."""
total_tokens = len(self.sequence_blocks[seq_id]) * self.block_size
# In practice, you'd track exact token counts per sequence
return total_tokens % self.block_size
PagedAttention dramatically improves memory utilization. In practice, vLLM reports throughput improvements of 2-4x compared to naive serving approaches, primarily because it eliminates the memory waste from over-provisioning cache for maximum sequence lengths.
2. Sliding Window Attention
Sliding window attention limits the cache to only the most recent N tokens, discarding older entries. This bounds the cache size to a fixed window regardless of total sequence length. Models like Mistral-7B use this technique natively:
class SlidingWindowKVCache:
def __init__(self, window_size, num_layers, num_kv_heads,
head_dim, dtype=torch.float16, device='cuda'):
self.window_size = window_size
self.num_layers = num_layers
# Fixed-size circular buffer
self.keys = torch.zeros(
num_layers, num_kv_heads, window_size, head_dim,
dtype=dtype, device=device
)
self.values = torch.zeros(
num_layers, num_kv_heads, window_size, head_dim,
dtype=dtype, device=device
)
self.write_ptr = 0
self.filled = 0
def append(self, new_keys, new_values):
"""Append new tokens, overwriting oldest if window is full."""
seq_len = new_keys.shape[2] # (layers, heads, seq, dim)
for i in range(seq_len):
for layer in range(self.num_layers):
self.keys[layer, :, self.write_ptr, :] = new_keys[layer, :, i, :]
self.values[layer, :, self.write_ptr, :] = new_values[layer, :, i, :]
self.write_ptr = (self.write_ptr + 1) % self.window_size
if self.filled < self.window_size:
self.filled += 1
def get_cache(self):
"""Return the cached keys and values in temporal order."""
if self.filled < self.window_size:
return self.keys[:, :, :self.filled, :], self.values[:, :, :self.filled, :]
# Rotate so the oldest entry is first
roll_amount = self.window_size - self.write_ptr
return (
torch.roll(self.keys, shifts=roll_amount, dims=2),
torch.roll(self.values, shifts=roll_amount, dims=2)
)
The trade-off is that sliding window attention loses information from tokens outside the window. This works well for tasks where local context is most important but may degrade performance on tasks requiring long-range dependencies.
3. KV Cache Quantization
Quantizing the KV cache from FP16 to INT8 or even INT4 can reduce memory usage by 2-4x with minimal quality loss. This is one of the most impactful optimizations for long-context inference:
import torch
class QuantizedKVCache:
"""INT8 quantized KV cache with per-channel scaling."""
def __init__(self, num_layers, num_kv_heads, max_seq_len,
head_dim, dtype=torch.float16, device='cuda'):
self.max_seq_len = max_seq_len
self.current_len = 0
# INT8 storage for keys and values
self.keys_int8 = torch.zeros(
num_layers, num_kv_heads, max_seq_len, head_dim,
dtype=torch.int8, device=device
)
self.values_int8 = torch.zeros(
num_layers, num_kv_heads, max_seq_len, head_dim,
dtype=torch.int8, device=device
)
# Per-head scale factors (FP16)
self.key_scales = torch.zeros(
num_layers, num_kv_heads, 1, 1, dtype=dtype, device=device
)
self.value_scales = torch.zeros(
num_layers, num_kv_heads, 1, 1, dtype=dtype, device=device
)
def quantize_tensor(self, tensor):
"""Quantize FP16 tensor to INT8 with per-head scaling."""
# tensor shape: (layers, heads, seq, dim)
abs_max = tensor.abs().amax(dim=(2, 3), keepdim=True)
scale = abs_max / 127.0
scale = torch.clamp(scale, min=1e-8)
quantized = torch.round(tensor / scale).to(torch.int8)
return quantized, scale.to(torch.float16)
def dequantize_tensor(self, quantized, scale):
"""Convert INT8 back to FP16."""
return quantized.to(torch.float16) * scale
def append(self, new_keys, new_values):
"""Quantize and store new KV data."""
# new_keys shape: (layers, heads, seq, dim)
q_keys, k_scales = self.quantize_tensor(new_keys)
q_values, v_scales = self.quantize_tensor(new_values)
seq_len = new_keys.shape[2]
end = self.current_len + seq_len
self.keys_int8[:, :, self.current_len:end, :] = q_keys
self.values_int8[:, :, self.current_len:end, :] = q_values
# Update scales (using the most recent batch's scale)
self.key_scales = k_scales
self.value_scales = v_scales
self.current_len = end
def get_cache(self):
"""Return dequantized keys and values."""
keys = self.dequantize_tensor(
self.keys_int8[:, :, :self.current_len, :], self.key_scales
)
values = self.dequantize_tensor(
self.values_int8[:, :, :self.current_len, :], self.value_scales
)
return keys, values
def memory_savings(self):
"""Calculate memory savings compared to FP16."""
fp16_size = 2 * self.keys_int8.numel() * 2 # 2 bytes per FP16 element
int8_size = 2 * self.keys_int8.numel() * 1 # 1 byte per INT8 element
return fp16_size - int8_size
In practice, INT8 KV cache quantization typically reduces quality by less than 1% on standard benchmarks while halving memory usage. INT4 quantization is more aggressive and may require more careful calibration, but can reduce cache memory by 4x.
4. KV Cache Eviction and Compression
For extremely long contexts, you can selectively evict less important KV entries based on attention scores. This approach, used in techniques like H2O (Heavy-Hitter Oracle) and StreamingLLM, identifies "heavy hitter" tokens that receive disproportionate attention and retains them while discarding others:
class EvictionKVCache:
"""KV cache with attention-score-based eviction."""
def __init__(self, max_cache_size, num_layers, num_kv_heads,
head_dim, dtype=torch.float16, device='cuda'):
self.max_cache_size = max_cache_size
self.current_len = 0
self.keys = torch.zeros(
num_layers, num_kv_heads, max_cache_size, head_dim,
dtype=dtype, device=device
)
self.values = torch.zeros(
num_layers, num_kv_heads, max_cache_size, head_dim,
dtype=dtype, device=device
)
# Track cumulative attention scores per position
self.attention_scores = torch.zeros(
max_cache_size, device=device
)
# Always keep the first few tokens (attention sinks)
self.sink_tokens = 4
def update_attention_scores(self, scores):
"""Accumulate attention scores for cached positions."""
# scores shape: (batch, heads, query_len, kv_len)
avg_scores = scores.mean(dim=(0, 1, 2)) # (kv_len,)
self.attention_scores[:self.current_len] += avg_scores
def maybe_evict(self, num_to_add):
"""Evict lowest-scoring tokens if cache is full."""
available = self.max_cache_size - self.current_len
if available >= num_to_add:
return # No eviction needed
num_to_evict = num_to_add - available
# Never evict sink tokens (first N tokens)
evictable_start = self.sink_tokens
evictable_scores = self.attention_scores[evictable_start:self.current_len]
# Find indices of lowest-scoring tokens
_, lowest_indices = torch.topk(evictable_scores, num_to_evict, largest=False)
lowest_indices = lowest_indices + evictable_start
# Compact the cache by removing evicted positions
keep_mask = torch.ones(self.max_cache_size, dtype=torch.bool, device=self.keys.device)
keep_mask[lowest_indices] = False
# Shift remaining entries to fill gaps
for layer in range(self.keys.shape[0]):
self.keys[layer] = self.keys[layer][:, keep_mask, :]
self.values[layer] = self.values[layer][:, keep_mask, :]
self.attention_scores = self.attention_scores[keep_mask]
self.current_len -= num_to_evict
def append(self, new_keys, new_values, attention_scores=None):
"""Add new tokens, evicting if necessary."""
seq_len = new_keys.shape[2]
if attention_scores is not None:
self.update_attention_scores(attention_scores)
self.maybe_evict(seq_len)
end = self.current_len + seq_len
self.keys[:, :, self.current_len:end, :] = new_keys
self.values[:, :, self.current_len:end, :] = new_values
self.current_len = end
5. Prefix Caching and Sharing
Many inference workloads share common prefixes — for example, a system prompt that is the same across all requests. Prefix caching stores the KV cache for common prefixes and reuses them across requests, avoiding redundant computation:
import hashlib
class PrefixCacheManager:
"""Manages reusable KV cache for shared prompt prefixes."""
def __init__(self, max_cached_prefixes=100, max_cache_memory_gb=10):
self.max_cached_prefixes = max_cached_prefixes
self.max_cache_memory_bytes = max_cache_memory_gb * 1024**3
self.current_memory = 0
self.cache = {} # prefix_hash -> (kv_cache, token_ids, size_bytes)
self.access_order = [] # For LRU eviction
def _hash_tokens(self, token_ids):
"""Create a hash key from token IDs."""
token_bytes = token_ids.cpu().numpy().tobytes()
return hashlib.sha256(token_bytes).hexdigest()
def get(self, token_ids):
"""Try to find a cached prefix. Returns (cache, matched_length) or (None, 0)."""
# Try progressively shorter prefixes
for length in range(len(token_ids), 0, -1):
prefix = token_ids[:length]
key = self._hash_tokens(prefix)
if key in self.cache:
kv_cache, cached_tokens, size = self.cache[key]
# Update LRU order
self.access_order.remove(key)
self.access_order.append(key)
return kv_cache, length
return None, 0
def put(self, token_ids, kv_cache, size_bytes):
"""Store a KV cache for the given token prefix."""
if size_bytes > self.max_cache_memory_bytes:
return # Don't cache if too large
key = self._hash_tokens(token_ids)
# Evict if at capacity
while (len(self.cache) >= self.max_cached_prefixes or
self.current_memory + size_bytes > self.max_cache_memory_bytes):
if not self.access_order:
break
evict_key = self.access_order.pop(0)
_, _, evict_size = self.cache[evict_key]
self.current_memory -= evict_size
del self.cache[evict_key]
self.cache[key] = (kv_cache, token_ids, size_bytes)
self.access_order.append(key)
self.current_memory += size_bytes
def stats(self):
return {
'cached_prefixes': len(self.cache),
'memory_used_gb': self.current_memory / 1024**3,
'memory_limit_gb': self.max_cache_memory_bytes / 1024**3
}
Prefix caching is particularly powerful in conversational AI and agentic workflows where the same system prompt or few-shot examples are used across many requests. Production systems like vLLM and SGLang implement this with significant throughput gains.
Putting It All Together: A Production-Ready Cache Manager
Here is a more complete cache manager that combines several of the techniques discussed above:
import torch
from dataclasses import dataclass, field
from typing import Optional, Dict, List, Tuple
@dataclass
class CacheConfig:
"""Configuration for KV cache management."""
num_layers: int = 32
num_kv_heads: int = 8
head_dim: int = 128
max_seq_len: int = 32768
block_size: int = 16
dtype: torch.dtype = torch.float16
device: str = 'cuda'
enable_quantization: bool = True
quantization_bits: int = 8
enable_eviction: bool = False
eviction_threshold: float = 0.9
sink_tokens: int = 4
max_batch_size: int = 32
class ProductionKVCacheManager:
"""Production-grade KV cache manager with multiple optimization strategies."""
def __init__(self, config: CacheConfig):
self.config = config
self.sequences: Dict[int, 'SequenceCache'] = {}
self.next_seq_id = 0
# Calculate per-sequence memory budget
bytes_per_element = 1 if config.enable_quantization else 2
self.element_dtype = torch.int8 if config.enable_quantization else config.dtype
per_token_bytes = (2 * config.num_layers * config.num_kv_heads *
config.head_dim * bytes_per_element)
self.max_tokens_per_batch = (
self._estimate_gpu_memory_budget() // per_token_bytes
)
def _estimate_gpu_memory_budget(self):
"""Estimate available GPU memory for KV cache."""
if torch.cuda.is_available():
free, _ = torch.cuda.mem_get_info()
# Reserve 30% for model weights and activations
return int(free * 0.7)
return 4 * 1024**3 # Default 4GB for CPU
def create_sequence(self) -> int:
"""Allocate a new sequence cache. Returns sequence ID."""
seq_id = self.next_seq_id
self.next_seq_id += 1
self.sequences[seq_id] = SequenceCache(
seq_id=seq_id,
config=self.config,
dtype=self.element_dtype
)
return seq_id
def append_tokens(self, seq_id: int, keys: List[torch.Tensor],
values: List[torch.Tensor]):
"""Append KV data for new tokens across all layers."""
if seq_id not in self.sequences:
raise ValueError(f"Unknown sequence ID: {seq_id}")
seq_cache = self.sequences[seq_id]
seq_cache.append(keys, values)
# Check if eviction is needed
if self.config.enable_eviction:
if seq_cache.current_len > self.config.max_seq_len * self.config.eviction_threshold:
seq_cache.evict(self.config.sink_tokens)
def get_cache(self, seq_id: int) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
"""Retrieve the full KV cache for a sequence."""
if seq_id not in self.sequences:
raise ValueError(f"Unknown sequence ID: {seq_id}")
return self.sequences[seq_id].get_cache()
def free_sequence(self, seq_id: int):
"""Release memory for a completed sequence."""
if seq_id in self.sequences:
del self.sequences[seq_id]
def get_memory_stats(self) -> dict:
"""Return current memory usage statistics."""
total_tokens = sum(s.current_len for s in self.sequences.values())
bytes_per_element = 1 if self.config.enable_quantization else 2
per_token_bytes = (2 * self.config.num_layers * self.config.num_kv_heads *
self.config.head_dim * bytes_per_element)
return {
'active_sequences': len(self.sequences),
'total_cached_tokens': total_tokens,
'estimated_memory_gb': (total_tokens * per_token_bytes) / 1024**3,
'max_batch_size': self.config.max_batch_size,
}
class SequenceCache:
"""Per-sequence KV cache with quantization and eviction support."""
def __init__(self, seq_id: int, config: CacheConfig, dtype: torch.dtype):
self.seq_id = seq_id
self.config = config
self.current_len = 0
self.dtype = dtype
# Pre-allocate cache tensors
self.keys = [
torch.zeros(
config.num_kv_heads, config.max_seq_len, config.head_dim,
dtype=dtype, device=config.device
)
for _ in range(config.num_layers)
]
self.values = [
torch.zeros(
config.num_kv_heads, config.max_seq_len, config.head_dim,
dtype=dtype, device=config.device
)
for _ in range(config.num_layers)
]
# Attention score tracking for eviction
self.attention_scores = torch.zeros(
config.max_seq_len, device=config.device
)
def append(self, keys: List[torch.Tensor], values: List[torch.Tensor]):
"""Append new KV data for all layers."""
seq_len = keys[0].shape[2] # (heads, seq, dim) or (batch, heads, seq, dim)
end = self.current_len + seq_len
if end > self.config.max_seq_len:
raise RuntimeError(
f"Sequence length {end} exceeds max {self.config.max_seq_len}"
)
for layer in range(self.config.num_layers):
# Handle quantization if enabled
if self.config.enable_quantization:
k_quantized, k_scale = self._quantize(keys[layer])
v_quantized, v_scale = self._quantize(values[layer])
self.keys[layer][:, self.current_len:end, :] = k_quantized
self.values[layer][:, self.current_len:end, :] = v_quantized
else:
self.keys[layer][:, self.current_len:end, :] = keys[layer]
self.values[layer][:, self.current_len:end, :] = values[layer]
self.current_len = end
def _quantize(self, tensor):
"""Quantize FP16 tensor to INT8."""
abs_max = tensor.abs().amax(dim=(2,), keepdim=True)
scale = abs_max / 127.0
scale = torch.clamp(scale, min=1e-8)
quantized = torch.round(tensor / scale).to(torch.int8)
return quantized, scale
def get_cache(self):
"""Return current KV cache (sliced to used portion)."""
keys = [self.keys[l][:, :self.current_len, :] for l in range(self.config.num_layers)]
values = [self.values[l][:, :self.current_len, :] for l in range(self.config.num_layers)]
if self.config.enable_quantization:
keys = [k.to(torch.float16) for k in keys]
values = [v.to(torch.float16) for v in values]
return keys, values
def evict(self, sink_tokens: int):
"""Evict lowest-attention-score tokens, keeping sink tokens."""
if self.current_len <= sink_tokens:
return
# Evict 10% of tokens
num_to_evict = max(1, self.current_len // 10)
evictable_scores = self.attention_scores[sink_tokens:self.current_len]
_, lowest_idx = torch.topk(evictable_scores, num_to_evict, largest=False)
lowest_idx = lowest_idx + sink_tokens
keep_mask = torch.ones(self.current_len, dtype=torch.bool,
device=self.config.device)
keep_mask[lowest_idx] = False
for layer in range(self.config.num_layers):
self.keys[layer] = self.keys[layer][:, keep_mask, :]
self.values[layer] = self.values[layer][:, keep_mask, :]
self.attention_scores = self.attention_scores[keep_mask]
self.current_len -= num_to_evict
# Example: Using the production cache manager
if __name__ == "__main__":
config = CacheConfig(
num_layers=32,
num_kv_heads=8,
head_dim=128,
max_seq_len=32768,
enable_quantization=True,
enable_eviction=True,
)
manager = ProductionKVCacheManager(config)
# Simulate serving multiple sequences
seq1 = manager.create_sequence()
seq2 = manager.create_sequence()
print(f"Created sequences: {seq1}, {seq2}")
print(f"Memory stats: {manager.get_memory_stats()}")
# Clean up
manager.free_sequence(seq1)
manager.free_sequence(seq2)
print("All sequences freed.")
Best Practices for KV Cache Management
Memory Planning and Sizing
- Profile before deploying: Measure actual KV cache memory usage with your target model, batch sizes, and sequence lengths before committing to a hardware configuration.
- Account for peak usage: Size your cache for the maximum expected concurrent sequences and sequence lengths, not just average usage.
- Use memory monitoring: Implement real-time monitoring of KV cache memory consumption and set alerts for when usage approaches limits.
- Consider GPU memory hierarchy: For extremely long contexts, consider offloading less frequently accessed cache entries to CPU memory or NVMe storage.
Choosing the Right Techniques
- Start with pre-allocation: Pre-allocated contiguous caches are the simplest optimization and should be your baseline.
- Add PagedAttention for multi-tenant serving: If you are serving multiple concurrent users with varying sequence lengths, PagedAttention provides the best memory utilization.
- Enable quantization for memory-constrained setups: INT8 KV cache quantization is nearly free in terms of quality and halves memory usage. Use it by default.
- Use prefix caching for repetitive prompts: If your workload involves