Attention Sinks › Implementation
Michael BrenndoerferUpdated July 4, 20257 min read
Language AI HandbookAll chapters Part 18 of 67Long Context
Attention Sinks › ImplementationOne section · Sections in this chapter7 / 16 ↓
Sections in this chapter16
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
With the mathematical foundation in place, let's translate these concepts into working code. The implementation centers on a key insight from our analysis: the cache must always contain the sink tokens, so we need a data structure that preserves specific positions while sliding others.
Our StreamingLLMCache class manages this by treating the cache as two distinct regions: a fixed sink region at the beginning that never changes, and a sliding window region that shifts as new tokens arrive. When the cache fills up, we evict the oldest window tokens while the sink tokens remain untouched.
In[8]:
Codeimport torch
class StreamingLLMCache:
"""
KV cache for streaming inference that preserves attention sinks.
The cache maintains:
- First `num_sink_tokens` positions (attention sinks)
- Most recent `window_size` positions (sliding window)
"""
def __init__(
self,
num_sink_tokens=4,
window_size=1020,
num_layers=12,
num_heads=12,
head_dim=64,
):
self.num_sink_tokens = num_sink_tokens
self.window_size = window_size
self.max_cache_size = num_sink_tokens + window_size
self.num_layers = num_layers
self.num_heads = num_heads
self.head_dim = head_dim
# Initialize empty caches for keys and values
# Shape: (num_layers, batch_size, num_heads, cache_len, head_dim)
self.k_cache = None
self.v_cache = None
self.cache_len = 0
def update(self, layer_idx, new_k, new_v):
"""
Add new key-value pairs to the cache, evicting old non-sink tokens if needed.
Args:
layer_idx: Which transformer layer this update is for
new_k: New key tensor of shape (batch, heads, new_len, head_dim)
new_v: New value tensor of shape (batch, heads, new_len, head_dim)
Returns:
Combined key and value tensors for attention computation
"""
batch_size = new_k.shape[0]
new_len = new_k.shape[2]
# Initialize cache on first update
if self.k_cache is None:
self.k_cache = torch.zeros(
self.num_layers,
batch_size,
self.num_heads,
self.max_cache_size,
self.head_dim,
device=new_k.device,
dtype=new_k.dtype,
)
self.v_cache = torch.zeros_like(self.k_cache)
# If cache has room, simply append
if self.cache_len + new_len <= self.max_cache_size:
self.k_cache[
layer_idx, :, :, self.cache_len : self.cache_len + new_len
] = new_k
self.v_cache[
layer_idx, :, :, self.cache_len : self.cache_len + new_len
] = new_v
if layer_idx == self.num_layers - 1:
self.cache_len += new_len
else:
# Cache is full: preserve sinks, slide window
# Keep first num_sink_tokens, evict oldest non-sink tokens
sink_k = self.k_cache[layer_idx, :, :, : self.num_sink_tokens]
sink_v = self.v_cache[layer_idx, :, :, : self.num_sink_tokens]
# Calculate how many window tokens we can keep
tokens_to_keep = self.window_size - new_len
if tokens_to_keep > 0:
# Keep recent window tokens
window_k = self.k_cache[layer_idx, :, :, -tokens_to_keep:]
window_v = self.v_cache[layer_idx, :, :, -tokens_to_keep:]
# Reconstruct cache: sinks + kept window + new tokens
self.k_cache[layer_idx, :, :, : self.num_sink_tokens] = sink_k
self.k_cache[
layer_idx,
:,
:,
self.num_sink_tokens : self.num_sink_tokens
+ tokens_to_keep,
] = window_k
self.k_cache[
layer_idx,
:,
:,
self.num_sink_tokens + tokens_to_keep : self.num_sink_tokens
+ tokens_to_keep
+ new_len,
] = new_k
self.v_cache[layer_idx, :, :, : self.num_sink_tokens] = sink_v
self.v_cache[
layer_idx,
:,
:,
self.num_sink_tokens : self.num_sink_tokens
+ tokens_to_keep,
] = window_v
self.v_cache[
layer_idx,
:,
:,
self.num_sink_tokens + tokens_to_keep : self.num_sink_tokens
+ tokens_to_keep
+ new_len,
] = new_v
else:
# Window is smaller than new tokens, just keep sinks and new
self.k_cache[layer_idx, :, :, : self.num_sink_tokens] = sink_k
self.k_cache[
layer_idx,
:,
:,
self.num_sink_tokens : self.num_sink_tokens + new_len,
] = new_k
self.v_cache[layer_idx, :, :, : self.num_sink_tokens] = sink_v
self.v_cache[
layer_idx,
:,
:,
self.num_sink_tokens : self.num_sink_tokens + new_len,
] = new_v
if layer_idx == self.num_layers - 1:
self.cache_len = (
self.num_sink_tokens + max(0, tokens_to_keep) + new_len
)
# Return the active portion of the cache for this layer
return (
self.k_cache[layer_idx, :, :, : self.cache_len],
self.v_cache[layer_idx, :, :, : self.cache_len],
)
def get_cache_length(self):
return self.cache_lenThe cache management is straightforward: when adding new tokens would exceed the maximum cache size, we preserve the sink tokens, keep as many recent tokens as possible, and append the new tokens. Let's test the cache behavior by adding tokens in batches:
In[9]:
Code# Test the cache behavior
cache = StreamingLLMCache(
num_sink_tokens=4, window_size=12, num_layers=2, num_heads=4, head_dim=32
)
# Simulate adding tokens in batches
batch_size = 1
cache_lengths = []
for step in range(5):
new_k = torch.randn(batch_size, 4, 4, 32) # 4 new tokens per step
new_v = torch.randn(batch_size, 4, 4, 32)
for layer in range(2):
k, v = cache.update(layer, new_k, new_v)
cache_lengths.append(cache.get_cache_length())Out[10]:
ConsoleStreamingLLM Cache Demonstration ================================================== Sink tokens: 4 Window size: 12 Max cache size: 16 Step 1: Added 4 tokens, cache length = 4 Step 2: Added 4 tokens, cache length = 8 Step 3: Added 4 tokens, cache length = 12 Step 4: Added 4 tokens, cache length = 16 Step 5: Added 4 tokens, cache length = 16
The cache grows from 4 to 16 tokens over the first four steps, then stabilizes at 16 (the maximum cache size of 4 sink tokens + 12 window tokens). Once the cache is full, the window slides forward while always preserving the sink tokens.
Now let's implement the attention mechanism that uses this cache:
In[11]:
Codeimport torch.nn.functional as F
def streaming_attention(query, cache, layer_idx, new_k, new_v, scale=None):
"""
Compute attention with StreamingLLM cache.
Args:
query: Query tensor of shape (batch, heads, 1, head_dim) for single token
cache: StreamingLLMCache instance
layer_idx: Current layer index
new_k: Key for new token(s)
new_v: Value for new token(s)
scale: Attention scale factor (default: 1/sqrt(head_dim))
Returns:
Attention output and updated cache
"""
head_dim = query.shape[-1]
if scale is None:
scale = 1.0 / (head_dim**0.5)
# Update cache and get all keys/values
keys, values = cache.update(layer_idx, new_k, new_v)
# Compute attention scores: (batch, heads, query_len, cache_len)
scores = torch.matmul(query, keys.transpose(-2, -1)) * scale
# Apply softmax to get attention weights
attn_weights = F.softmax(scores, dim=-1)
# Compute output
output = torch.matmul(attn_weights, values)
return output, attn_weightsIn[12]:
Code# Demonstrate attention weight distribution
torch.manual_seed(42)
cache = StreamingLLMCache(
num_sink_tokens=4, window_size=20, num_layers=1, num_heads=1, head_dim=64
)
# Fill cache with some initial tokens
initial_k = torch.randn(1, 1, 20, 64)
initial_v = torch.randn(1, 1, 20, 64)
# Make all 4 sink tokens have distinctive key patterns (simulating trained sinks)
initial_k[:, :, 0, :] = initial_k[:, :, 0, :] * 3.0 + 1.0 # Sink token 0
initial_k[:, :, 1, :] = initial_k[:, :, 1, :] * 2.5 + 0.8 # Sink token 1
initial_k[:, :, 2, :] = initial_k[:, :, 2, :] * 2.0 + 0.6 # Sink token 2
initial_k[:, :, 3, :] = initial_k[:, :, 3, :] * 1.5 + 0.4 # Sink token 3
cache.update(0, initial_k, initial_v)
# Now add a new query token
query = torch.randn(1, 1, 1, 64)
new_k = torch.randn(1, 1, 1, 64)
new_v = torch.randn(1, 1, 1, 64)
output, attn_weights = streaming_attention(query, cache, 0, new_k, new_v)Out[13]:
Visualization
The attention weights show the expected pattern: higher attention on the first few positions (the sink tokens) with the remaining attention distributed across the window tokens.