Sliding Window Attention › Implementation
Language AI HandbookAll chapters Sliding Window Attention › ImplementationOne section · Sections in this chapter7 / 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
With the theory established, let's translate these ideas into working code. We'll build sliding window attention step by step, starting with the most straightforward approach and progressively optimizing. Along the way, you'll see how the mathematical formulas map directly to implementation choices.
Mask-Based Implementation
The most direct implementation follows the formula exactly: compute all attention scores, apply the mask, run softmax, and aggregate values. This approach prioritizes clarity over efficiency, making it ideal for understanding the mechanism.
First, we need helper functions. The softmax function normalizes scores into probability distributions. We implement a numerically stable version that subtracts the maximum before exponentiating to prevent overflow:
import numpy as np
def softmax(x, axis=-1):
"""Numerically stable softmax."""
x_max = np.max(x, axis=axis, keepdims=True)
exp_x = np.exp(x - x_max)
return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
def create_sliding_window_mask(seq_len, window_size, causal=False):
"""
Create a sliding window attention mask.
Args:
seq_len: Length of the sequence
window_size: Size of the attention window
causal: If True, mask future positions (for autoregressive models)
Returns:
Mask of shape (seq_len, seq_len) with 0 for allowed positions,
-inf for masked positions
"""
mask = np.full((seq_len, seq_len), -1e9)
half_window = window_size // 2
for i in range(seq_len):
if causal:
# Only attend to current and previous positions within window
start = max(0, i - window_size + 1)
end = i + 1
else:
# Symmetric window
start = max(0, i - half_window)
end = min(seq_len, i + half_window + 1)
mask[i, start:end] = 0
return maskThe mask creation function iterates through each position and sets the mask to 0 for positions within the window, for positions outside. Notice how the code directly implements our mathematical definition: for symmetric windows, we check if the distance is at most half_window; for causal windows, we also require that the key position is not in the future.
This implementation explicitly constructs the full mask, which works well for moderate sequence lengths. For very long sequences, we'll need a more efficient approach, but this version is perfect for understanding the pattern.
With the mask function ready, we can implement the complete attention computation:
def sliding_window_attention(Q, K, V, window_size, causal=False):
"""
Compute sliding window attention.
Args:
Q: Query matrix of shape (seq_len, d_k)
K: Key matrix of shape (seq_len, d_k)
V: Value matrix of shape (seq_len, d_v)
window_size: Size of the attention window
causal: If True, use causal masking
Returns:
Output of shape (seq_len, d_v)
"""
seq_len, d_k = Q.shape
# Compute attention scores
scores = Q @ K.T / np.sqrt(d_k)
# Apply sliding window mask
mask = create_sliding_window_mask(seq_len, window_size, causal)
masked_scores = scores + mask
# Apply softmax
weights = softmax(masked_scores, axis=-1)
# Compute weighted sum of values
output = weights @ V
return output, weightsThis function follows our four-step process exactly: compute scores (), scale by , add the mask, apply softmax, and multiply by values. The function returns both the output and the attention weights, which we'll visualize to verify the windowed pattern.
Let's test this implementation with a concrete example. We'll create random QKV representations, then examine the attention pattern:
# Create sample input
seq_len = 16
d_k = 32
Q = np.random.randn(seq_len, d_k) * 0.1
K = np.random.randn(seq_len, d_k) * 0.1
V = np.random.randn(seq_len, d_k) * 0.1
# Compute attention with window size 5
output, weights = sliding_window_attention(Q, K, V, window_size=5, causal=False)Input shapes: Q=(16, 32), K=(16, 32), V=(16, 32) Output shape: (16, 32) Attention weights shape: (16, 16) Attention pattern for token 8: Positions attended: [ 6 7 8 9 10] Weights: [0.19854472 0.20024106 0.19944171 0.20083186 0.20094065]
The attention weights confirm that token 8 attends only to positions 6-10 (a window of 5 centered on position 8), with all other weights effectively zero.
Efficient Block-Based Implementation
The mask-based implementation has a problem: it still computes and stores the full attention matrix, even though most entries are zero. For a 10,000-token sequence, we'd allocate 100 million floats just for the attention weights. This defeats the purpose of sliding window attention.
The solution is to never materialize the full matrix. Instead, we process each position independently, computing attention only over its local window. This insight transforms memory complexity from to , where is the window size:
def efficient_sliding_window_attention(Q, K, V, window_size):
"""
Memory-efficient sliding window attention.
Instead of creating the full n x n attention matrix, we process
each position individually, computing only the relevant window.
Args:
Q, K, V: Query, Key, Value matrices of shape (seq_len, d)
window_size: Size of the attention window
Returns:
Output of shape (seq_len, d)
"""
seq_len, d_k = Q.shape
d_v = V.shape[1]
half_window = window_size // 2
output = np.zeros((seq_len, d_v))
for i in range(seq_len):
# Determine window boundaries
start = max(0, i - half_window)
end = min(seq_len, i + half_window + 1)
# Extract local keys and values
K_local = K[start:end] # shape: (window, d_k)
V_local = V[start:end] # shape: (window, d_v)
# Compute attention scores for this position
scores = Q[i : i + 1] @ K_local.T / np.sqrt(d_k) # shape: (1, window)
weights = softmax(scores, axis=-1)
# Compute output for this position
output[i] = (weights @ V_local).squeeze()
return outputThe key insight is in the loop body. For each position , we compute the window boundaries, extract only the relevant keys and values, compute attention scores over just positions, and aggregate. Each iteration uses memory instead of , and we never store the full attention matrix.
This approach trades sequential processing for memory efficiency. In practice, production implementations use batched operations and GPU-optimized kernels to process multiple positions in parallel while still avoiding the full matrix.
Let's verify that both implementations produce identical results:
# Verify both implementations produce the same result
output_mask_based, _ = sliding_window_attention(Q, K, V, window_size=5)
output_efficient = efficient_sliding_window_attention(Q, K, V, window_size=5)
difference = np.max(np.abs(output_mask_based - output_efficient))Maximum difference between implementations: 1.39e-17 Implementations match: True
The negligible difference (on the order of ) confirms that both implementations compute identical results. The efficient version achieves the same output while using significantly less memory for long sequences.
Causal Sliding Window for Language Models
For autoregressive language models like GPT, we need an additional constraint: tokens must not attend to future positions. If token 10 could see token 11, it would be "cheating" during text generation, looking at words it's supposed to predict.
Combining causality with sliding windows creates a pattern where each token sees at most previous tokens, including itself. The window is asymmetric: all context comes from the past. Token 10 with attends to tokens 7, 8, 9, and 10, never tokens 11 and beyond:
def causal_sliding_window_attention(Q, K, V, window_size):
"""
Causal sliding window attention for autoregressive models.
Each position attends only to itself and the previous (window_size - 1) positions.
"""
seq_len, d_k = Q.shape
d_v = V.shape[1]
output = np.zeros((seq_len, d_v))
for i in range(seq_len):
# Only attend to current and previous positions within window
start = max(0, i - window_size + 1)
end = i + 1
K_local = K[start:end]
V_local = V[start:end]
scores = Q[i : i + 1] @ K_local.T / np.sqrt(d_k)
weights = softmax(scores, axis=-1)
output[i] = (weights @ V_local).squeeze()
return output