27. Sliding-Window Attention

Problem

Mistral 7B's headline architectural trick is that no query ever attends to the full history: each query position i is restricted to a local window of the W most recent key positions (including itself), rather than every position from 0 to i. Stacking L layers of window W still gives a theoretical receptive field of roughly L * W tokens, but each layer's attention matmul only ever needs to touch an O(seq_len * W) band instead of the full O(seq_len2) matrix.

This problem asks for the windowed attention output itself. A full O(seq_len2) scores matrix with the band masked out is an acceptable, interview-reasonable solution here — a truly O(seq_len * W) rolling-buffer version is the natural follow-up once this is working. Given q, k, v already split into heads, compute attention where position i may only attend to positions j with j <= i (causal) and i - j < window_size (local).

Source: src/27_sliding_window_attention.py

def sliding_window_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, window_size: int,
) -> torch.Tensor: ...

>>> q = k = v = torch.randn(1, 2, 6, 4)
>>> sliding_window_attention(q, k, v, window_size=3).shape
torch.Size([1, 2, 6, 4])

Step-by-Step Approach

1. Compute the ordinary scaled dot-product scores matrix q @ k.transpose(-1, -2) / sqrt(head_dim), shape (batch, heads, seq, seq) — the windowing happens entirely in the mask, not the matmul itself.
2. Build a (seq, seq) boolean band mask with torch.arange: i = arange(seq).view(seq, 1), j = arange(seq).view(1, seq), mask = (j <= i) & ((i - j) < window_size).
3. Apply the mask with scores.masked_fill(~mask, float("-inf")) so masked-out positions get exactly zero probability after softmax, without special-casing them.
4. Note the mask always keeps at least the diagonal j == i for every row, so no row is ever all -inf — softmax is safe from NaNs by construction.
5. torch.softmax(scores, dim=-1) then weights @ v to get the windowed output, exactly as in ordinary attention.
6. Sanity-check with window_size >= seq_len: the band mask degenerates to plain causal masking, so the output must match ordinary causal self-attention exactly.

The key insight is that sliding-window attention is not a different attention mechanism, it's the same softmax(QK^T)V with a stricter mask: standard causal attention only removes future keys, this removes future keys and keys further than window_size in the past.

Reference solution

import math
import torch


def sliding_window_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, window_size: int,
) -> torch.Tensor:
    batch, heads, seq, hd = q.shape

    # full scores matrix, then mask out everything outside the causal + local
    # window band -- masked_fill with -inf before softmax is what keeps the
    # excluded positions at exactly 0 probability without special-casing them
    scores = q @ k.transpose(-1, -2) / math.sqrt(hd)  # (batch, heads, seq, seq)

    idx = torch.arange(seq, device=q.device)
    i, j = idx.view(seq, 1), idx.view(1, seq)
    # every row i always keeps at least the diagonal (j == i), so no row of
    # the mask is ever all-False -- softmax never sees an all -inf row (no NaNs)
    band_mask = (j <= i) & ((i - j) < window_size)  # (seq, seq)

    scores = scores.masked_fill(~band_mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)
    return weights @ v  # (batch, heads, seq, head_dim)

Key Functions & Tricks

• torch.arange(seq).view(seq, 1) / .view(1, seq) — broadcasts to an (seq, seq) pairwise index-difference grid without any Python loop.
• tensor.masked_fill(~mask, float("-inf")) — sets disallowed logits to -inf so softmax drives their probability to exactly 0.
• torch.softmax(scores, dim=-1) — normalizes each query's kept-key scores into a probability distribution over just its window.
• Boolean & between two broadcastable masks — composes the causal condition and the window condition into a single band mask in one expression.
• q @ k.transpose(-1, -2) — batched matmul that broadcasts over both the batch and heads leading dimensions automatically.

How to Recognize This Pattern