KV Cache Compression › Combining Techniques
Language AI HandbookAll chapters KV Cache Compression › Combining TechniquesOne section · Sections in this chapter8 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
The most effective KV cache compression strategies combine multiple approaches to achieve multiplicative benefits. Rather than choosing between eviction and quantization, we can apply both: first reduce the number of cached tokens through intelligent eviction, then store the remaining tokens at reduced precision through quantization. This layered approach addresses memory from two complementary angles, and the benefits multiply rather than merely add.
The combination works so well because the two techniques address fundamentally different aspects of the memory problem. Eviction reduces the number of tokens we need to store, which scales with sequence length. Quantization reduces the bytes needed per token, which scales with the vector dimensions and precision. These are orthogonal dimensions of the memory tensor, and reducing both independently provides multiplicative compression. A 16x reduction in token count combined with a 2x reduction in bytes per token gives 32x total compression.
A practical configuration might use the following combination:
- Attention sink preservation: Keep the first 4 tokens unconditionally to maintain computational stability.
- H2O-style eviction: Maintain heavy hitters plus recent tokens to preserve both semantic importance and local context.
- INT8 quantization: Store all cached values at 8-bit precision to halve the per-token memory footprint.
import torch
class CompressedKVCache:
"""
Combined compression: sink preservation + H2O eviction + quantization.
"""
def __init__(
self,
sink_size=4,
recent_size=64,
heavy_hitter_size=64,
num_heads=32,
head_dim=128,
quant_bits=8,
):
self.sink_size = sink_size
self.recent_size = recent_size
self.heavy_hitter_size = heavy_hitter_size
self.num_heads = num_heads
self.head_dim = head_dim
# Quantizer for storage
self.quantizer = PerChannelQuantizedCache(bits=quant_bits)
# Runtime state (full precision for computation)
self.keys = None
self.values = None
self.cumulative_attention = None
self.positions = None
@property
def max_cache_size(self):
return self.sink_size + self.recent_size + self.heavy_hitter_size
def update(self, new_key, new_value, attention_weights=None):
"""Add new token with combined compression strategy."""
if self.keys is None:
self.keys = new_key
self.values = new_value
self.cumulative_attention = torch.zeros(new_key.size(0), 1)
self.positions = torch.tensor([[0]])
return
# Update attention scores
if attention_weights is not None:
attn_update = attention_weights.mean(dim=1).squeeze(-2)
self.cumulative_attention = self.cumulative_attention + attn_update
# Add new token
new_pos = self.positions.max() + 1
self.keys = torch.cat([self.keys, new_key], dim=2)
self.values = torch.cat([self.values, new_value], dim=2)
self.cumulative_attention = torch.cat(
[self.cumulative_attention, torch.zeros(new_key.size(0), 1)], dim=1
)
self.positions = torch.cat(
[self.positions, torch.tensor([[new_pos]])], dim=1
)
# Apply eviction if needed
if self.keys.size(2) > self.max_cache_size:
self._evict()
def _evict(self):
"""Evict using sink + H2O strategy."""
current_size = self.keys.size(2)
positions = self.positions[0]
# 1. Always keep sink tokens (first sink_size positions)
sink_mask = positions < self.sink_size
# 2. Keep recent tokens
non_sink_positions = positions.clone()
non_sink_positions[sink_mask] = -1
_, recent_indices = non_sink_positions.topk(self.recent_size)
recent_mask = torch.zeros(current_size, dtype=torch.bool)
recent_mask[recent_indices] = True
# 3. From remaining, keep heavy hitters
remaining_mask = ~(sink_mask | recent_mask)
remaining_attention = self.cumulative_attention[0].clone()
remaining_attention[~remaining_mask] = float("-inf")
n_heavy = min(self.heavy_hitter_size, remaining_mask.sum().item())
if n_heavy > 0:
_, heavy_indices = remaining_attention.topk(n_heavy)
heavy_mask = torch.zeros(current_size, dtype=torch.bool)
heavy_mask[heavy_indices] = True
else:
heavy_mask = torch.zeros(current_size, dtype=torch.bool)
# Combine masks
keep_mask = sink_mask | recent_mask | heavy_mask
# Apply eviction
self.keys = self.keys[:, :, keep_mask, :]
self.values = self.values[:, :, keep_mask, :]
self.cumulative_attention = self.cumulative_attention[:, keep_mask]
self.positions = self.positions[:, keep_mask]
def get_quantized_state(self):
"""Return quantized cache state for storage."""
self.quantizer.store(self.keys, self.values)
return (
self.quantizer.keys_quantized,
self.quantizer.values_quantized,
self.quantizer.key_params,
self.quantizer.value_params,
)
def memory_stats(self):
"""Report memory usage statistics."""
if self.keys is None:
return {}
batch, heads, seq_len, head_dim = self.keys.shape
# Full precision (FP16): 2 bytes per element
full_precision_bytes = (
batch * heads * seq_len * head_dim * 2 * 2
) # keys + values
# Quantized: bits/8 bytes per element
quantized_bytes = (
batch * heads * seq_len * head_dim * 2 * self.quantizer.bits / 8
)
return {
"tokens_cached": seq_len,
"max_cache_size": self.max_cache_size,
"full_precision_mb": full_precision_bytes / (1024 * 1024),
"quantized_mb": quantized_bytes / (1024 * 1024),
"compression_ratio": full_precision_bytes / quantized_bytes,
}import torch
def benchmark_compressed_cache(num_tokens):
"""Simulate long sequence generation with compressed cache."""
cache = CompressedKVCache(
sink_size=4,
recent_size=32,
heavy_hitter_size=28, # Total: 4 + 32 + 28 = 64 tokens
num_heads=32,
head_dim=128,
quant_bits=8,
)
# Simulate generation
batch_size = 1
for t in range(num_tokens):
new_key = torch.randn(batch_size, 32, 1, 128)
new_value = torch.randn(batch_size, 32, 1, 128)
# Simulate attention weights
if cache.keys is not None:
current_size = cache.keys.size(2)
attn = torch.softmax(
torch.randn(batch_size, 32, 1, current_size), dim=-1
)
else:
attn = None
cache.update(new_key, new_value, attn)
return cachenum_tokens = 1000
cache = benchmark_compressed_cache(num_tokens)
stats = cache.memory_stats()
# Calculate what uncompressed would be
uncompressed_bytes = 1 * 32 * num_tokens * 128 * 2 * 2 # KV, FP16
uncompressed_mb = uncompressed_bytes / (1024 * 1024)
overall_compression = uncompressed_mb / stats["quantized_mb"]Compressed KV Cache Results (1000 tokens generated): Without compression: 15.6 MB (1000 tokens) With eviction only: 1.0 MB (64 tokens) With eviction + INT8: 0.5 MB Overall compression: 31.2x
The combined approach achieves dramatic memory reduction through the multiplication of two independent compression factors. Eviction reduces the token count from 1000 to 64. This provides approximately 15.6x compression. INT8 quantization then provides an additional 2x reduction by halving the bytes per token. The total compression ratio exceeds 30x, transforming what would require large memory into a compact representation that fits easily in limited GPU memory.
This level of compression enables generation at context lengths that would otherwise exceed available memory. A configuration that might require 16GB of KV cache memory without compression can operate within 500MB using these combined techniques. For deployment scenarios where memory is constrained or costs must be minimized, such aggressive compression makes previously infeasible applications practical. Production systems serving thousands of concurrent users on shared GPU hardware routinely rely on these techniques to make the economics viable.