25. Sliding-Window Causal Attention for Streaming Inference

Problem

Plain causal attention lets query position i attend to every position from 0 to i — fine for a short prompt, but for a real-time streaming system processing an unbounded, continuously growing input (audio frames, a long conversation), that means both the KV-cache and the attention cost at position i grow without bound as the stream continues. A system with a real-time latency budget can't tolerate per-token cost that keeps climbing the longer a session runs.

Sliding-window attention bounds this: each query only attends to the most recent window_size positions (including itself), never anything further back, regardless of how long the sequence has grown. This caps both per-token compute and (paired with a bounded-size KV-cache that evicts anything older than the window) memory, at the cost of the model losing access to context older than the window.

Source: src/25_sliding_window_attention.py

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

>>> q = k = v = torch.randn(1, 5, 3)
>>> out = sliding_window_causal_attention(q, k, v, window_size=2)
>>> out.shape
torch.Size([1, 5, 3])
>>> # with window_size >= seq_len this must equal ordinary causal attention

Step-by-Step Approach

1. Compute raw scaled scores exactly as in plain attention: q @ k.transpose(-1, -2) / sqrt(d), shape (batch, seq, seq).
2. Build query and key position index vectors with torch.arange(seq_len), then broadcast them into a (seq, seq) grid: query_pos = positions.unsqueeze(1) (row index i), key_pos = positions.unsqueeze(0) (column index j).
3. Combine two conditions into the allowed mask: (key_pos <= query_pos) for causality, AND (key_pos > query_pos - window_size) for the window bound — together, exactly the last window_size positions up to and including i.
4. Apply the mask the usual way: masked_fill(~allowed, float("-inf")) before softmax.
5. Softmax over the key axis and multiply by v, unchanged from plain attention — the entire difference from ordinary causal attention is this one extra boolean condition in the mask.
6. Verify two invariants directly: window_size >= seq_len must reduce exactly to ordinary full causal attention (the window condition becomes vacuously true), and corrupting a value strictly older than the window for a given query position must leave that query's output completely unchanged.

The key insight is that sliding-window attention is causal masking with a second, symmetric boundary added — instead of only blocking the future (j > i), it also blocks the distant past (j ≤ i - window_size), turning the triangular causal mask into a diagonal band of fixed width.

Reference solution

import math
import torch


def sliding_window_causal_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, window_size: int,
) -> torch.Tensor:
    seq_len, d = q.shape[-2], q.shape[-1]

    scores = q @ k.transpose(-1, -2) / math.sqrt(d)

    positions = torch.arange(seq_len, device=q.device)
    query_pos = positions.unsqueeze(1)  # (seq, 1)
    key_pos = positions.unsqueeze(0)    # (1, seq)

    # causal AND within the trailing window
    allowed = (key_pos <= query_pos) & (key_pos > query_pos - window_size)

    scores = scores.masked_fill(~allowed, float("-inf"))
    attn_weights = torch.softmax(scores, dim=-1)
    return attn_weights @ v

Key Functions & Tricks

• positions.unsqueeze(1) / unsqueeze(0) — the standard broadcast-to-grid idiom for building any position-pair-dependent mask (causal, windowed, or otherwise) without an explicit double loop.
• Combining two boolean conditions with & — the window mask is just causal masking intersected with a second, sliding lower bound; no new masking mechanism is needed, just a second predicate.
• query_pos - window_size — the lower bound of the allowed range; note it's a strict > (not >=), since position i - window_size itself is exactly one step too old.
• Reduction check against window_size >= seq_len — confirms the window condition degrades gracefully to a no-op rather than accidentally excluding valid positions at the boundary.
• Corrupt-and-compare independence check — the same technique as the causal-masking problem, adapted to confirm both boundaries of the window (not just the future one) actually block information flow.

How to Recognize This Pattern