Sliding Window Attention

Sliding Window Attention (SWA) restricts each token to nearby tokens instead of the entire sequence.

For bidirectional encoders, token \(i\) attends to keys \(j\) where \(|i-j| \le r\): up to \(r\) tokens on either side, plus itself. Causal attention additionally requires \(j \le i\).

Efficiency

A band mask on full-sequence attention does not automatically save compute. The implementation must skip out-of-window pairs.

• Dense attention: \(O(T^2)\) attention work.
• Local attention: \(O(Tw)\), where \(w=2r+1\) is the window width.
• Stacking layers expands the receptive field. Global tokens/layers can connect distant regions more directly.

Training on Apple Silicon

A practical PyTorch MPS approach is chunked local SDPA:

1. Split queries into blocks [s:e].
2. Slice keys/values to [max(0, s-r):min(T, e+r)].
3. Call torch.nn.functional.scaled_dot_product_attention with a small boolean mask enforcing abs(q_pos-k_pos) <= r (True means allowed).
4. Concatenate block outputs along the sequence dimension.

Autograd works through these ordinary operations. With query blocks no larger than the window, attention work and saved score storage scale as \(O(Tw)\) rather than \(O(T^2)\). Benchmark block sizes: small blocks increase launch overhead. Activation checkpointing can further reduce saved intermediates.

As of November 2026: MPS FlexAttention is forward-only, so it is not a training replacement. Tinygrad's core SWA path uses dense masked attention. Third-party MLX kernels need scrutiny: differentiable does not necessarily mean the backward pass is sparse.

Reference: PyTorch SDPA.

Leave a Comment