Full attention costs . Sliding-window attention, used by Mistral and Gemma, lets each token attend only to the last tokens, itself included, which makes long contexts affordable.
Implement:
sliding_window_mask(n, window): an (n, n) boolean array where mask[i, j] is True if query i may attend to key j, which is when j <= i (causal) and i - j < window.masked_attention(Q, K, V, mask): scaled dot-product attention where disallowed positions get before the softmax. Return (output, weights). Every row has at least one allowed position.attended_pairs(n, window): how many True entries the mask has, computed without building the mask (a formula or loop over rows).