KV Cache Compression › The H2O Algorithm
Language AI HandbookAll chapters KV Cache Compression › The H2O AlgorithmOne section · Sections in this chapter6 / 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
Heavy-Hitter Oracle (H2O) represents a more sophisticated eviction policy that extends beyond simple recency or sink preservation. The algorithm combines attention-based selection with sink preservation, addressing the limitations of both pure window-based and pure attention-based approaches. The key insight motivating H2O is that tokens exhibit different temporal importance patterns. Some tokens are "heavy hitters," consistently receiving high attention across many generation steps because they carry information the model repeatedly needs. Other tokens are transient, relevant only briefly before becoming obsolete as generation proceeds.
The name "oracle" is somewhat aspirational: a true oracle would know with certainty which future tokens would be needed and preserve exactly those. H2O approximates this oracle by using historical attention patterns as a proxy for future importance. This approximation works well in practice because semantic importance tends to be persistent. A sentence's subject noun remains important throughout the sentences that follow. A key fact mentioned in the introduction remains relevant when the conclusion references it. The oracle approximation is not perfect, but it is far better than ignoring historical patterns entirely.
Consider how these patterns manifest in practice. A subject noun at the beginning of a paragraph might be a heavy hitter, receiving attention throughout the paragraph as the model maintains coherence. A transition word, by contrast, might receive attention only from the immediately following tokens before becoming irrelevant. H2O captures this distinction by tracking cumulative attention over time and using it to identify tokens worthy of long-term preservation. The algorithm maintains two complementary categories of preserved tokens:
- Recent tokens: A sliding window of the most recent tokens. These provide immediate local context and ensure the model can maintain coherent generation based on what was just produced.
- Heavy hitters: Tokens with the highest cumulative attention scores, capped at tokens. These represent semantically important content that the model has repeatedly consulted, suggesting continued relevance.
The algorithm dynamically balances these sets throughout generation. This keeps both recency and historical importance are considered when making eviction decisions. Tokens can transition between categories: a recent token that accumulates high attention becomes a heavy hitter candidate, while a former heavy hitter that stops receiving attention may eventually be evicted as new heavy hitters emerge. This fluidity is a key advantage over static approaches that fix which tokens to preserve at the start of generation.
import torch
class H2OCache:
"""
Heavy-Hitter Oracle (H2O) KV cache compression.
Maintains recent tokens plus heavy-hitters (tokens with highest
cumulative attention scores).
Reference: "H2O: Heavy-Hitter Oracle for Efficient Generative Inference"
"""
def __init__(self, recent_size, heavy_hitter_size, num_heads, head_dim):
self.recent_size = recent_size
self.heavy_hitter_size = heavy_hitter_size
self.num_heads = num_heads
self.head_dim = head_dim
# Cache storage
self.keys = None # (batch, heads, cache_size, head_dim)
self.values = None
self.cumulative_attention = None # (batch, cache_size)
self.token_positions = None # Original position of each token
@property
def max_cache_size(self):
return self.recent_size + self.heavy_hitter_size
def update(self, new_key, new_value, attention_weights):
"""
Add new token and update heavy-hitter statistics.
Args:
new_key: (batch, heads, 1, head_dim)
new_value: (batch, heads, 1, head_dim)
attention_weights: (batch, heads, 1, current_cache_size)
"""
batch_size = new_key.size(0)
if self.keys is None:
# Initialize with first token
self.keys = new_key
self.values = new_value
self.cumulative_attention = torch.zeros(batch_size, 1)
self.token_positions = torch.tensor([[0]])
return
# Update cumulative attention for existing tokens
# Average across heads for simplicity
attn_update = attention_weights.mean(dim=1).squeeze(
-2
) # (batch, cache_size)
self.cumulative_attention = self.cumulative_attention + attn_update
# Add new token
new_position = self.token_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(batch_size, 1)], dim=1
)
self.token_positions = torch.cat(
[self.token_positions, torch.tensor([[new_position]])], dim=1
)
# Evict if necessary
current_size = self.keys.size(2)
if current_size > self.max_cache_size:
self._evict()
def _evict(self):
"""
Evict tokens to maintain cache size.
Keep: recent_size most recent + heavy_hitter_size highest attention
"""
batch_size = self.keys.size(0)
current_size = self.keys.size(2)
# Identify recent tokens (by position)
positions = self.token_positions[0] # (cache_size,)
_, recent_indices = positions.topk(self.recent_size)
recent_mask = torch.zeros(current_size, dtype=torch.bool)
recent_mask[recent_indices] = True
# Among non-recent tokens, find heavy hitters
non_recent_mask = ~recent_mask
non_recent_attention = self.cumulative_attention[0].clone()
non_recent_attention[recent_mask] = float("-inf") # Exclude recent
# Get top heavy hitters from non-recent tokens
_, heavy_hitter_indices = non_recent_attention.topk(
min(self.heavy_hitter_size, non_recent_mask.sum())
)
heavy_hitter_mask = torch.zeros(current_size, dtype=torch.bool)
heavy_hitter_mask[heavy_hitter_indices] = True
# Keep tokens that are either recent or heavy hitters
keep_mask = recent_mask | heavy_hitter_mask
# Apply mask
self.keys = self.keys[:, :, keep_mask, :]
self.values = self.values[:, :, keep_mask, :]
self.cumulative_attention = self.cumulative_attention[:, keep_mask]
self.token_positions = self.token_positions[:, keep_mask]import torch
import torch.nn.functional as F
def demonstrate_h2o(heavy_hitter_positions):
"""Show H2O cache behavior with varying attention patterns."""
batch_size = 1
num_heads = 8
head_dim = 64
cache = H2OCache(
recent_size=8,
heavy_hitter_size=8,
num_heads=num_heads,
head_dim=head_dim,
)
# Simulate 50 tokens with varying attention patterns
for t in range(50):
new_key = torch.randn(batch_size, num_heads, 1, head_dim)
new_value = torch.randn(batch_size, num_heads, 1, head_dim)
if cache.keys is not None:
current_size = cache.keys.size(2)
# Create attention weights
attn = torch.rand(batch_size, num_heads, 1, current_size)
# Give extra attention to heavy hitter positions if they're in cache
for pos_idx, pos in enumerate(cache.token_positions[0].tolist()):
if pos in heavy_hitter_positions:
attn[:, :, :, pos_idx] += 2.0 # Boost attention
# Normalize
attn = F.softmax(attn, dim=-1)
cache.update(new_key, new_value, attn)
else:
cache.update(new_key, new_value, None)
return cacheheavy_hitter_positions = {5, 10}
cache = demonstrate_h2o(heavy_hitter_positions)After 50 tokens: Cache size: 16 tokens Token positions in cache: [0, 1, 2, 3, 4, 5, 6, 10, 42, 43, 44, 45, 46, 47, 48, 49] Cumulative attention scores: Position 0: 4.059 Position 1: 3.039 Position 2: 2.462 Position 3: 2.214 Position 4: 1.966 Position 5: 12.586 (heavy hitter) Position 6: 1.736 Position 10: 10.055 (heavy hitter) Position 42: 0.254 Position 43: 0.213 Position 44: 0.161 Position 45: 0.146 Position 46: 0.120 Position 47: 0.070 Position 48: 0.034 Position 49: 0.000
Notice how positions 5 and 10 are retained in the cache despite not being among the most recent tokens. Their high cumulative attention scores qualify them as heavy hitters, preserving important context while still maintaining recency through the recent token budget. The algorithm has automatically identified these tokens as carrying important information that the model repeatedly needs, and it protects them from eviction even as dozens of newer tokens have been generated.
The cumulative attention scores reveal the distinction between heavy hitters and ordinary tokens. Positions 5 and 10 have accumulated substantially higher scores than other retained positions because the simulated attention pattern consistently attended to them. In a real model, these high-attention tokens would typically correspond to semantically important content: the subject of a sentence, a key entity being discussed, or a necessary fact that informs subsequent generation. The visualization makes it clear why H2O outperforms pure sliding windows: important early context is retained automatically, without any manual specification of what to keep.
Worked Example: H2O Eviction Step-by-Step
Let's trace through a concrete numerical example to see exactly how H2O makes eviction decisions. Suppose our cache has recent_size = 2 and heavy_hitter_size = 2, giving a maximum cache size of 4. We have just generated 5 tokens (positions 0-4), so the cache has grown to 5 entries and must evict one.
After 5 generation steps, the cumulative attention scores for the 5 cached tokens are:
- Position 0: cumulative attention = 3.8 (this is an anchor token, repeatedly consulted)
- Position 1: cumulative attention = 0.4 (barely consulted)
- Position 2: cumulative attention = 2.1 (moderately important)
- Position 3: cumulative attention = 0.3 (barely consulted, just added to window)
- Position 4: cumulative attention = 0.0 (brand new, no history yet)
The algorithm proceeds as follows. First, identify the recent_size = 2 most recent tokens by position. Positions 3 and 4 are the two most recent, so they enter the recent set unconditionally. Second, among the remaining tokens (positions 0, 1, 2), identify the heavy_hitter_size = 2 tokens with highest cumulative attention. Position 0 has score 3.8 and position 2 has score 2.1, making them the heavy hitters. Position 1 with score 0.4 is the eviction candidate. Third, the final cache contains positions {0, 2, 3, 4}, preserving the anchor token and the recent window while discarding the low-importance middle token at position 1.
This example illustrates the core logic: recency and importance are evaluated independently, and both contribute tokens to the final cache. A token can survive eviction for either reason: being recent or being an established heavy hitter. Only tokens that are neither recent nor historically important get evicted.