Sliding Window Attention › Mistral-Style Windowed Attention
Language AI HandbookAll chapters Sliding Window Attention › Mistral-Style Windowed AttentionOne section · Sections in this chapter8 / 12 ↓
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
Mistral 7B popularized sliding window attention in modern large language models. Their approach uses a 4,096-token window size, which covers most practical context lengths while enabling efficient inference.
Key Design Choices
Mistral's implementation incorporates several practical considerations:
- Large window size (4096): Covers entire documents in many use cases, reducing the need for global tokens
- Rolling buffer KV cache: During generation, only the most recent key-value pairs are stored
- Pre-fill and chunking: Long prompts are processed in chunks to manage memory
The rolling buffer approach is particularly elegant. Traditional transformers store all past key-value pairs, requiring memory that grows with sequence length:
where:
- : memory required for the full KV cache
- : current sequence length (grows during generation)
- : number of transformer layers
- : model dimension
With sliding window attention, we only store the most recent key-value pairs, giving constant memory usage:
where is the window size. Since , , and are all constants for a given model, memory usage becomes independent of sequence length. This is the key insight that enables generating arbitrarily long sequences without running out of memory.
class RollingKVCache:
"""
Rolling key-value cache for efficient autoregressive generation.
Maintains only the most recent 'window_size' key-value pairs,
enabling constant memory usage regardless of sequence length.
"""
def __init__(self, window_size, num_heads, d_k, d_v):
self.window_size = window_size
self.num_heads = num_heads
# Pre-allocate buffers
self.k_cache = np.zeros((window_size, num_heads, d_k))
self.v_cache = np.zeros((window_size, num_heads, d_v))
# Current position in the circular buffer
self.position = 0
self.length = 0
def update(self, k, v):
"""
Add new key-value pair to the cache.
Args:
k: New key of shape (num_heads, d_k)
v: New value of shape (num_heads, d_v)
"""
# Write to current position (wraps around)
idx = self.position % self.window_size
self.k_cache[idx] = k
self.v_cache[idx] = v
self.position += 1
self.length = min(self.length + 1, self.window_size)
def get_kv(self):
"""Get current keys and values in correct order."""
if self.length < self.window_size:
# Haven't filled the buffer yet
return self.k_cache[: self.length], self.v_cache[: self.length]
else:
# Buffer is full, need to reorder
start = self.position % self.window_size
indices = np.concatenate(
[np.arange(start, self.window_size), np.arange(0, start)]
)
return self.k_cache[indices], self.v_cache[indices]# Demonstrate rolling cache behavior
cache = RollingKVCache(window_size=4, num_heads=2, d_k=8, d_v=8)
# Simulate adding tokens
for i in range(6):
k = np.full((2, 8), i) # Simple keys to track order
v = np.full((2, 8), i)
cache.update(k, v)
K, V = cache.get_kv()
first_values = K[:, 0, 0] # First element of each cached keyRolling KV Cache Demo (window=4): After 6 updates, cache contains: 4 entries Cached positions (by value): [2. 3. 4. 5.] Only most recent 4 tokens are stored
After adding 6 tokens to a cache with window size 4, tokens 0 and 1 have been evicted. The cache now contains only tokens 2, 3, 4, and 5. This shows how the circular buffer automatically discards the oldest entries when it reaches capacity. This behavior mirrors what happens during autoregressive generation in models like Mistral.
Memory Comparison
The rolling buffer approach provides substantial memory savings for long sequences:
def kv_cache_memory_mb(
seq_len, num_layers, num_heads, d_head, window_size=None, dtype_bytes=2
):
"""
Calculate KV cache memory usage.
Args:
seq_len: Current sequence length
num_layers: Number of transformer layers
num_heads: Number of attention heads
d_head: Dimension per head
window_size: If provided, use rolling buffer; otherwise full cache
dtype_bytes: 2 for fp16, 4 for fp32
"""
# Full cache: 2 (K+V) * layers * heads * seq_len * d_head * dtype
if window_size is None:
elements = 2 * num_layers * num_heads * seq_len * d_head
else:
# Rolling buffer: only window_size positions
effective_len = min(seq_len, window_size)
elements = 2 * num_layers * num_heads * effective_len * d_head
return elements * dtype_bytes / (1024**2) # Convert to MB
# Mistral 7B config (approximate)
config = {"num_layers": 32, "num_heads": 32, "d_head": 128, "window_size": 4096}
seq_lengths = [1024, 4096, 16384, 65536]KV Cache Memory (fp16, Mistral 7B config):
Seq Len Full Cache Rolling (w=4096) Savings
--------------------------------------------------------
1,024 512.0 MB 512.0 MB 0.0%
4,096 2048.0 MB 2048.0 MB 0.0%
16,384 8192.0 MB 2048.0 MB 75.0%
65,536 32768.0 MB 2048.0 MB 93.8%The memory comparison reveals the scaling advantage of the rolling buffer. At 1,024 tokens (below the window size), both approaches use the same memory. At 4,096 tokens (exactly the window size), they remain equivalent. Beyond that point, the savings become dramatic: at 16K tokens, the rolling buffer uses 75% less memory, and at 64K tokens, savings reach over 93%.
This enables processing long documents on hardware that could not otherwise accommodate them.