KV Cache Compression › Attention-Based Eviction
Language AI HandbookAll chapters KV Cache Compression › Attention-Based EvictionOne section · Sections in this chapter4 / 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
A more sophisticated approach selects tokens based on their historical attention scores rather than their recency. The intuition behind this method is straightforward: tokens that have received high attention in the past are likely to be important for future generation. These tokens have already proven their value, so preserving them makes sense. By tracking cumulative attention and evicting low-attention tokens, we can maintain a compact cache of the most relevant context while adapting to the specific content being processed.
Think of attention-based eviction as a hotel concierge who tracks which guests request the most services. Guests who never ring the service bell can safely be checked out early to free up rooms. Guests who constantly request assistance clearly need to stay. The analogy breaks down slightly because we cannot know in advance which guests will become high-demand, but historical patterns give us a reliable signal.
The key insight here is that attention scores provide a signal of token importance. When the model attends heavily to a particular token, it extracts information from that token to inform its output. Tokens that consistently receive high attention across multiple generation steps are carrying information that the model repeatedly needs. These "frequently consulted" tokens are prime candidates for preservation, while tokens that receive negligible attention can likely be removed without consequence. The cumulative nature of this tracking is important: a token that receives a small amount of attention at each of 100 steps has accumulated substantial importance, even if no single step stands out as particularly high.
import torch
class AttentionBasedCache:
"""
Maintains a cache of fixed size by evicting tokens with lowest
cumulative attention scores.
"""
def __init__(self, max_cache_size, num_heads, head_dim):
self.max_cache_size = max_cache_size
self.num_heads = num_heads
self.head_dim = head_dim
# Initialize empty cache
self.keys = None # (batch, heads, cache_size, head_dim)
self.values = None
self.attention_scores = None # Cumulative attention per token
def update(self, new_key, new_value, attention_weights):
"""
Add new token to cache and evict if necessary.
Args:
new_key: (batch, heads, 1, head_dim)
new_value: (batch, heads, 1, head_dim)
attention_weights: (batch, heads, 1, seq_len) - attention to all cached tokens
"""
batch_size = new_key.size(0)
if self.keys is None:
# First token - initialize cache
self.keys = new_key
self.values = new_value
# Initialize with the attention this token receives (from itself)
self.attention_scores = attention_weights.mean(dim=1).squeeze(
-2
) # (batch, 1)
return
# Append new token
self.keys = torch.cat([self.keys, new_key], dim=2)
self.values = torch.cat([self.values, new_value], dim=2)
# Update cumulative attention scores for existing tokens
# Average attention across heads, then add to cumulative scores
current_attention = attention_weights.mean(dim=1).squeeze(
-2
) # (batch, seq_len)
# Existing tokens get their attention updated (all but last position)
old_scores = self.attention_scores + current_attention[:, :-1]
# New token starts with its self-attention score
new_score = current_attention[:, -1:]
self.attention_scores = torch.cat([old_scores, new_score], dim=1)
# Evict if over capacity
if self.keys.size(2) > self.max_cache_size:
self._evict_lowest_attention()
def _evict_lowest_attention(self):
"""Remove the token with lowest cumulative attention."""
# Find index of minimum attention score
_, min_idx = self.attention_scores.min(dim=1)
batch_size = self.keys.size(0)
# Create mask for tokens to keep (all except min_idx)
keep_mask = torch.ones(batch_size, self.keys.size(2), dtype=torch.bool)
for b in range(batch_size):
keep_mask[b, min_idx[b]] = False
# Apply mask
# For simplicity, we'll handle batch_size=1 case
self.keys = self.keys[:, :, keep_mask[0], :]
self.values = self.values[:, :, keep_mask[0], :]
self.attention_scores = self.attention_scores[:, keep_mask[0]]This approach adapts dynamically to the content being processed. Important tokens accumulate high attention scores over time and remain in the cache, while less relevant tokens get evicted as their scores remain low. The method naturally identifies semantically meaningful tokens like subjects, key entities, and important verbs that the model needs to reference repeatedly. Unlike the sliding window approach, attention-based eviction can retain a token from position 5 while evicting tokens from positions 50 through 200 if those middle tokens have not been useful.
However, the approach introduces computational overhead for tracking scores and, more critically, has a basic flaw that we must address. The flaw relates to a peculiar attention pattern we covered earlier: attention sinks. We will explore this issue and its solution in the next section.