KV Cache Compression › Attention Sink Preservation
Language AI HandbookAll chapters KV Cache Compression › Attention Sink PreservationOne section · Sections in this chapter5 / 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
As we discussed in Part XVIII, Chapter 5 on attention sinks, transformers exhibit a curious behavior that complicates attention-based eviction: initial tokens in a sequence accumulate disproportionate attention even when they carry no semantic importance. These "attention sinks" appear to act as computational anchors. This provides a stable target for attention heads to discharge probability mass when no relevant context exists. The phenomenon emerges from training dynamics: the model learns to use early tokens as a default attention target. This creates a dependency that persists even when those tokens are semantically meaningless.
To understand why this creates a problem for eviction, consider what happens during long generation. Attention-based eviction is supposed to preserve high-attention tokens because they carry important information. But initial tokens receive high attention for a completely different reason: they serve as a computational escape valve. If we sort all tokens by cumulative attention and discard the lowest-scoring ones, we will correctly identify many low-importance middle tokens for eviction. But we might also accidentally evict tokens that had middling cumulative attention scores while the true attention sinks at positions 0-3 consume a large fraction of the attention budget and dominate the preservation list.
When we evict tokens based purely on historical attention scores, we risk removing tokens that should stay while incorrectly preserving tokens that are sinks rather than semantically important. Without attention sinks available, models may produce incoherent outputs, fall into repetition loops, or exhibit generally degraded quality. The attention mechanism loses its computational anchor, disrupting the learned patterns that depend on its presence. This is not a minor quality degradation; for some models and tasks, removing attention sinks completely breaks generation.
This creates a paradox for attention-based eviction strategies. The tokens receiving highest attention (the sinks) may not be semantically important, while truly important content tokens may receive only moderate attention because they share the attention budget with the sinks. Pure attention-score-based eviction conflates two distinct sources of high attention: semantic importance and computational anchoring. These require different treatment in a principled eviction policy.
The solution is to protect initial tokens from eviction regardless of their attention scores. By unconditionally preserving the first few tokens, we ensure the attention sinks remain available while still allowing the eviction policy to remove low-importance tokens from the rest of the sequence. This separation lets the score-based eviction logic focus on semantic importance without interference from the sink phenomenon.
import torch
class StreamingLLMCache:
"""
Implements the StreamingLLM approach: preserve initial sink tokens
plus a sliding window of recent tokens.
Reference: "Efficient Streaming Language Models with Attention Sinks"
"""
def __init__(self, sink_size, window_size, num_heads, head_dim):
self.sink_size = sink_size # Number of initial tokens to always keep
self.window_size = window_size # Number of recent tokens to keep
self.num_heads = num_heads
self.head_dim = head_dim
self.keys = None
self.values = None
self.total_tokens_seen = 0
@property
def max_cache_size(self):
return self.sink_size + self.window_size
def update(self, new_key, new_value):
"""
Add new token, maintaining sink tokens and recent window.
Args:
new_key: (batch, heads, 1, head_dim)
new_value: (batch, heads, 1, head_dim)
"""
self.total_tokens_seen += 1
if self.keys is None:
self.keys = new_key
self.values = new_value
return
# Append new token
self.keys = torch.cat([self.keys, new_key], dim=2)
self.values = torch.cat([self.values, new_value], dim=2)
current_size = self.keys.size(2)
# If we exceed capacity, keep sinks + recent window
if current_size > self.max_cache_size:
# Keep first sink_size tokens and last window_size tokens
sink_keys = self.keys[:, :, : self.sink_size, :]
sink_values = self.values[:, :, : self.sink_size, :]
window_keys = self.keys[:, :, -self.window_size :, :]
window_values = self.values[:, :, -self.window_size :, :]
self.keys = torch.cat([sink_keys, window_keys], dim=2)
self.values = torch.cat([sink_values, window_values], dim=2)
def get_cache(self):
"""Return current key-value cache."""
return self.keys, self.valuesimport torch
# Demonstrate StreamingLLM cache behavior
def simulate_generation(cache, num_tokens):
"""Simulate generating num_tokens and track cache state."""
batch_size = 1
history = []
for i in range(num_tokens):
# Simulate a new key-value pair
new_key = torch.randn(batch_size, cache.num_heads, 1, cache.head_dim)
new_value = torch.randn(batch_size, cache.num_heads, 1, cache.head_dim)
cache.update(new_key, new_value)
if cache.keys is not None:
history.append(cache.keys.size(2))
return historycache = StreamingLLMCache(sink_size=4, window_size=32, num_heads=8, head_dim=64)
history = simulate_generation(cache, 100)Configuration: 4 sink tokens + 32 window Maximum cache size: 36 tokens Cache size over generation: After 10 tokens: 10 After 36 tokens: 36 After 50 tokens: 36 After 100 tokens: 36
The StreamingLLM approach maintains exactly sink_size + window_size tokens once the cache is full. This provides predictable memory usage while preserving the necessary attention sinks. The cache size grows linearly during the initial phase, then plateaus at the maximum size and remains constant regardless of how many additional tokens are generated. This bounded memory behavior enables to long sequences long streaming generation without memory growth, which is the defining property that gives this approach its name.
Research has shown that just 4 sink tokens are typically sufficient to maintain generation quality, making this a highly efficient solution. The small sink budget adds minimal memory overhead while giving the computational anchoring the model requires. Combined with a modest window size, the approach achieves substantial compression while maintaining coherent generation for streaming applications. The trade-off is that middle context (between the sinks and the recent window) is entirely lost, which can degrade quality on tasks requiring long-range reasoning.