Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Mistral Architecture › Implementation: Sliding Window Attention

Michael BrenndoerferUpdated August 6, 20257 min read

Language AI HandbookAll chapters

Part 29 of 67Modern Decoder Models

Chapter 211Mistral Architecture › Implementation: Sliding Window Attention

Reading section 7 of 14 · Implementation: Sliding Window Attention · 7 minSection 7 of 14 · 7 min

Section 7 of 14 · 7 min read

Mistral Architecture › Implementation: Sliding Window AttentionOne section · Sections in this chapter7 / 14 ↓

Sections in this chapter14

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

Having understood the mathematics behind sliding window attention, let's implement it from scratch. Building the mechanism ourselves will solidify the concepts and reveal implementation details that matter in practice.

Step 1: Building the Attention MaskLink Copied

The foundation of sliding window attention is the mask that enforces which positions can attend to which. We need to encode two constraints:

Causal constraint: Position $i$ cannot attend to any position $j > i$ (no looking into the future)
Window constraint: Position $i$ cannot attend to any position $j < i - w + 1$ (outside the sliding window)

Both constraints translate to blocking certain positions with $-\infty$ values:

In[22]:

Code

def create_sliding_window_mask(seq_len, window_size):
    """
    Create a causal mask with sliding window constraint.

    Returns a mask where:
    - 0 indicates positions that CAN be attended to
    - -inf indicates positions that CANNOT be attended to

    Args:
        seq_len: Length of the sequence
        window_size: Size of the sliding window

    Returns:
        Mask tensor of shape (seq_len, seq_len)
    """
    # Start with zeros (allowing all positions)
    mask = np.zeros((seq_len, seq_len))

    for i in range(seq_len):
        for j in range(seq_len):
            if j > i:
                # Future positions: block (causal constraint)
                mask[i, j] = float("-inf")
            elif i - j >= window_size:
                # Past positions beyond window: block
                mask[i, j] = float("-inf")

    return mask

The logic is straightforward: for each position pair $(i, j)$ , we check if $j$ is in the future ( $j > i$ ) or outside the window ( $i - j \geq w$ ). If either condition is true, we block with $-\infty$ .

Let's visualize the mask to verify it matches our expectations:

Out[23]:

Console

Sliding Window Mask (8 tokens, window=4):
0 = can attend, -inf = blocked

Token 0:   .  X  X  X  X  X  X  X
Token 1:   .  .  X  X  X  X  X  X
Token 2:   .  .  .  X  X  X  X  X
Token 3:   .  .  .  .  X  X  X  X
Token 4:   X  .  .  .  .  X  X  X
Token 5:   X  X  .  .  .  .  X  X
Token 6:   X  X  X  .  .  .  .  X
Token 7:   X  X  X  X  .  .  .  .

Reading this output row by row tells the story of the sliding window:

Token 0: Can only attend to itself (the sequence just started)
Token 1: Can attend to positions 0-1 (window not yet full)
Token 2: Can attend to positions 0-2 (window not yet full)
Token 3: Can attend to positions 0-3 (window now full with 4 positions)
Token 4: The window slides! Can attend to positions 1-4, blocking position 0
Tokens 5-7: The window continues sliding, always covering exactly 4 positions

This banded diagonal structure is what turns quadratic complexity into linear.

Step 2: The Complete Attention Forward PassLink Copied

With the mask in place, we can implement the full sliding window attention computation. The algorithm follows the standard attention formula, but with our sliding window mask applied before softmax:

In[24]:

Code

def sliding_window_attention(
    queries,  # (batch, seq_len, num_heads, head_dim)
    keys,  # (batch, seq_len, num_heads, head_dim)
    values,  # (batch, seq_len, num_heads, head_dim)
    window_size,
    scale=None,
):
    """
    Compute sliding window attention.

    Each query can only attend to keys within the sliding window.
    """
    batch_size, seq_len, num_heads, head_dim = queries.shape

    if scale is None:
        scale = 1.0 / np.sqrt(head_dim)

    # Reshape for batched matrix multiplication
    # (batch, num_heads, seq_len, head_dim)
    q = queries.transpose(0, 2, 1, 3)
    k = keys.transpose(0, 2, 1, 3)
    v = values.transpose(0, 2, 1, 3)

    # Compute attention scores: (batch, num_heads, seq_len, seq_len)
    scores = np.matmul(q, k.transpose(0, 1, 3, 2)) * scale

    # Apply sliding window mask
    mask = create_sliding_window_mask(seq_len, window_size)
    scores = scores + mask  # Broadcasting adds mask to all batches and heads

    # Softmax over keys dimension
    # Use stable softmax: subtract max before exp
    scores_max = np.max(scores, axis=-1, keepdims=True)
    scores_exp = np.exp(scores - scores_max)
    attention_weights = scores_exp / np.sum(scores_exp, axis=-1, keepdims=True)

    # Apply attention to values
    output = np.matmul(attention_weights, v)

    # Reshape back to (batch, seq_len, num_heads, head_dim)
    output = output.transpose(0, 2, 1, 3)

    return output, attention_weights

Let's walk through this implementation step by step:

Reshaping: We transpose from (batch, seq, heads, dim) to (batch, heads, seq, dim) to enable batched matrix multiplication across all heads simultaneously
Score computation: The matrix product q @ k.T computes dot products between all query-key pairs, giving raw attention scores. We scale by $1/\sqrt{d_k}$ to prevent the softmax from becoming too peaked
Mask application: Adding the mask ( $0$ or $-\infty$ ) to scores enforces our sliding window constraint. Positions with $-\infty$ will have zero attention weight after softmax
Stable softmax: We subtract the maximum score before exponentiating to prevent numerical overflow. This is mathematically equivalent to standard softmax but numerically stable
Value aggregation: The final matmul blends values according to attention weights, producing the output representation

Step 3: Testing and VisualizationLink Copied

Let's verify our implementation works correctly and visualize the resulting attention patterns:

In[25]:

Code

# Create sample data engineered so heads show contrasting within-window patterns:
# Head 0: queries aligned with the most-recent key  → strong recency bias
# Head 1: queries spread evenly across all window keys → uniform distribution
batch_size = 1
seq_len = 16
num_heads = 4
head_dim = 32
window_size = 4

# Base random tensors
queries = np.random.randn(batch_size, seq_len, num_heads, head_dim).astype(
    np.float32
)
keys = np.random.randn(batch_size, seq_len, num_heads, head_dim).astype(
    np.float32
)
values = np.random.randn(batch_size, seq_len, num_heads, head_dim).astype(
    np.float32
)

# Head 0: boost dot-product with the immediately preceding key so recency
# dominates.  For each position i, align q[i, head0] strongly toward k[i, head0].
scale = 1.0 / np.sqrt(head_dim)
for i in range(1, seq_len):
    queries[0, i, 0] = (
        keys[0, i, 0] * 6.0
    )  # q aligns with own key (most recent)

# Head 1: make queries nearly orthogonal to all keys so scores are near-zero
# before softmax → uniform attention within window
queries[0, :, 1] = np.random.randn(seq_len, head_dim).astype(np.float32) * 0.05

# Run sliding window attention
output, attention_weights = sliding_window_attention(
    queries, keys, values, window_size
)

Out[26]:

Console

Input shapes:
  Queries: (1, 16, 4, 32)
  Keys:    (1, 16, 4, 32)
  Values:  (1, 16, 4, 32)

Output shape: (1, 16, 4, 32)
Attention weights shape: (1, 4, 16, 16)

Window size: 4
Attention pattern is banded with width 4

The output maintains the same shape as the input queries, as expected. The attention weights tensor has shape (batch, heads, seq_len, seq_len), but only the diagonal band of width 4 contains non-zero values due to our sliding window mask.

Out[27]:

Visualization

Heatmap showing sliding window attention weights for head 0 with banded diagonal structure. — Attention pattern for head 0. Each row shows which keys a query attends to, with the sliding window creating a banded diagonal structure.

Heatmap showing sliding window attention weights for head 1 with different weight distribution. — Attention pattern for head 1. Different heads learn different attention patterns within the same sliding window constraint.

The attention heatmaps reveal several important properties of sliding window attention:

The banded structure: Each row shows attention concentrated in a diagonal band of width 4 (our window size). The strict boundaries are enforced by our mask, the positions outside the band have exactly zero attention weight.

Within-window variation: Inside the allowed window, different heads develop different patterns. Head 0 might emphasize recent positions while Head 1 focuses on slightly older ones. This diversity is valuable, different heads can specialize in different types of relationships.

Preserved expressiveness: Despite the hard constraint on which positions can be attended, the model retains full flexibility in how attention is distributed within the window. The softmax ensures weights sum to 1, but the specific allocation depends on the learned query-key interactions.

This implementation demonstrates that sliding window attention is a straightforward modification to standard attention. The only change is adding a carefully constructed mask before softmax. Everything else, the score computation, softmax normalization, and value aggregation, remains identical.