29. Sliding-Window Local Causal Attention

Problem

Full causal self-attention lets every token attend back to the entire history, which is O(T^2) in both compute and (materialized) memory. Many production long-context transformers instead restrict each token to a bounded local window of recent tokens — attention cost becomes O(T * W) for window size W, independent of total sequence length, which matters a lot for a system processing long-running audio/text streams where the oldest context stops being useful anyway.

Implement sliding-window causal self-attention: query position i may only attend to key positions j such that j <= i (causal) and i - j < window_size (local). Positions outside the window get zero attention weight, exactly as if they were masked out of a full causal attention computation.

Source: src/29_sliding_window_attention.py

def sliding_window_attention(
    q: torch.Tensor,  # (B, H, T, D)
    k: torch.Tensor,  # (B, H, T, D)
    v: torch.Tensor,  # (B, H, T, D)
    window_size: int,
) -> torch.Tensor:  # (B, H, T, D)
    ...

>>> q = k = v = torch.randn(1, 1, 6, 4)
>>> out = sliding_window_attention(q, k, v, window_size=2)
>>> out.shape
torch.Size([1, 1, 6, 4])
# token i=5 (0-indexed) can only see keys at positions 4 and 5

Step-by-Step Approach

1. Compute raw scaled dot-product scores over the whole sequence first, (Q @ K^T) / sqrt(D), giving a (T, T) matrix — correctness first, memory efficiency is a separate follow-up optimization.
2. Build two (T, T) position grids with torch.arange(T): one for query index i (as a column vector) and one for key index j (as a row vector), so they broadcast against each other.
3. Combine two boolean conditions into one allowed mask: j <= i (causal) AND i - j < window_size (inside the local window).
4. Use masked_fill to set every disallowed (i, j) score to -inf before the softmax — exactly the same masking mechanism as plain causal attention, just with a tighter allowed set.
5. Softmax over the last dimension and matmul with V as usual; the -inf entries automatically become exactly-zero attention weight after softmax.
6. Sanity-check the two boundary cases: window_size=1 should make every token attend only to itself, and window_size >= T should reduce exactly to standard full causal attention.

The key insight is that sliding-window attention is just causal attention with a stricter mask — the softmax/matmul machinery is unchanged, only the boolean condition defining which (i, j) pairs are "allowed" gets a second clause added.

Reference solution

import math

import torch


def sliding_window_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    window_size: int,
) -> torch.Tensor:
    B, H, T, D = q.shape
    scale = 1.0 / math.sqrt(D)

    scores = torch.matmul(q, k.transpose(-1, -2)) * scale  # (B, H, T, T)

    # allowed[i, j] = True iff j is causally valid (j<=i) AND inside the
    # local window (i-j < window_size); everything else gets -inf so softmax
    # zeroes it out just like a full causal mask would.
    qi = torch.arange(T, device=q.device).unsqueeze(1)  # (T, 1)
    kj = torch.arange(T, device=q.device).unsqueeze(0)  # (1, T)
    allowed = (kj <= qi) & (qi - kj < window_size)       # (T, T), broadcasts against scores
    scores = scores.masked_fill(~allowed, float("-inf"))

    attn = torch.softmax(scores, dim=-1)  # (B, H, T, T)
    return torch.matmul(attn, v)


TEST_CASES = [
    {"name": "window smaller than sequence length", "B": 1, "H": 1, "T": 8, "D": 4, "window_size": 3},
    {"name": "window of 1 (each token only attends to itself)", "B": 1, "H": 2, "T": 6, "D": 4, "window_size": 1},
    {"name": "window >= T (degenerates to full causal attention)", "B": 2, "H": 1, "T": 5, "D": 4, "window_size": 100},
]


def main():
    import torch.nn.functional as F

    torch.manual_seed(0)

    for i, case in enumerate(TEST_CASES):
        B, H, T, D, W = case["B"], case["H"], case["T"], case["D"], case["window_size"]
        print(f"Test {i}: {case['name']} (B={B}, H={H}, T={T}, D={D}, window_size={W})")
        q = torch.randn(B, H, T, D)
        k = torch.randn(B, H, T, D)
        v = torch.randn(B, H, T, D)

        out = sliding_window_attention(q, k, v, W)
        print(f"  out.shape={tuple(out.shape)}")
        assert out.shape == (B, H, T, D)

        qi = torch.arange(T).unsqueeze(1)
        kj = torch.arange(T).unsqueeze(0)
        allowed = (kj <= qi) & (qi - kj < W)
        additive_mask = torch.zeros(T, T).masked_fill(~allowed, float("-inf"))
        expected = F.scaled_dot_product_attention(q, k, v, attn_mask=additive_mask)

        torch.testing.assert_close(out, expected, atol=1e-4, rtol=1e-4)
        print("PASSED")


if __name__ == "__main__":
    main()

Key Functions & Tricks

• torch.arange(T).unsqueeze(1) / .unsqueeze(0) — builds the row-vector and column-vector index grids that broadcast into a full (T, T) position-pair matrix.
• boolean AND of two masks (&) — combines the causal condition and the window condition into a single allowed mask without materializing them separately.
• masked_fill(~allowed, float('-inf')) — the standard additive-masking pattern: invert the allowed set and fill disallowed positions with -inf pre-softmax.
• torch.softmax(scores, dim=-1) — converts masked logits into a proper probability distribution per query, with -inf entries collapsing to exactly 0.
• F.scaled_dot_product_attention(..., attn_mask=...) — used in the test as an independently-built oracle, taking an explicit additive mask so it doesn't rely on the student's masking logic.

How to Recognize This Pattern