Brayden Zhang

Sliding Window Attention

2 min read

Restrict each query to the most recent keys, making attention and bounding the KV Cache to tokens per layer.

Equation

Stacking windowed layers gives an effective receptive field of tokens, since information hops one window per layer.

From scratch

import torch
 
def sliding_window_mask(n_q, n_k, w, device='cpu'):
    i = torch.arange(n_k - n_q, n_k, device=device)[:, None]   # (n_q, 1)
    j = torch.arange(n_k, device=device)[None, :]              # (1, n_k)
    keep = (j <= i) & (j > i - w)                              # causal AND in-window
    return ~keep                                               # True = masked
 
def swa(q, k, v, w):                                           # (b, h, n, dh)
    s = torch.einsum('bhid,bhjd->bhij', q, k) / q.shape[-1]**0.5
    m = sliding_window_mask(q.shape[2], k.shape[2], w, q.device)
    return torch.einsum('bhij,bhjd->bhid',
                        s.masked_fill(m, float('-inf')).softmax(-1), v)
 
class RollingCache:                     # ring buffer, never exceeds w
    def __init__(self, w): self.w, self.k, self.v = w, None, None
    def update(self, k, v):
        self.k = k if self.k is None else torch.cat([self.k, k], 2)[:, :, -self.w:]
        self.v = v if self.v is None else torch.cat([self.v, v], 2)[:, :, -self.w:]
        return self.k, self.v

Notes

  • Common pattern: interleave windowed layers with a few full-attention layers (e.g. 3:1) so global information still has a direct path.
  • Attention sinks: keeping the first few tokens permanently in the cache alongside the rolling window prevents the perplexity blowup seen when the initial tokens are evicted (StreamingLLM).
  • The dense mask above is pedagogical; real kernels skip out-of-window blocks entirely, which is where the speedup comes from.
  • Pairs naturally with ALiBi and Relative Position, since both encode a recency prior.

Related: Scaled Dot-Product Attention, KV Cache Implementation, Linear Attention

Graph View