Medium Research

Sliding Window Attention

Sliding-window attention limits each query to nearby keys. This exercise uses a two-sided window: position i can attend to itself and up to window_size positions on either side. It can use later positions, so this is not a causal mask.

The window is a radius

Let w = window_size. Keep key j when abs(i - j) <= w, clipping the range to the sequence boundaries. There are at most 2 * w + 1 allowed keys, and fewer near either end.

For a three-position sequence with w = 1:

Query position	Allowed key positions
0	0, 1
1	0, 1, 2
2	1, 2

If queries and keys are zero and the values are [[10], [20], [30]], all allowed scores are equal. The outputs are the local means [[15], [20], [25]]. With w = 0, only the diagonal is allowed and the output equals V. With w >= seq_len - 1, all keys are allowed.

Compute attention within each window

Calculate scores = (Q @ K.T) / sqrt(d_k). Mask scores outside each window with negative infinity before applying a stable softmax along the last axis. Multiply the resulting weights by V.

Softmax must normalize over only the allowed keys. Every row includes its own position, so it has at least one allowed score. Assume the unmasked dot products are finite. A fixed finite penalty such as -1e9 is not an exact substitute for negative infinity when excluded scores can be very large.

A local attention pattern can support implementations that compute only nearby pairs. This exercise's dense score-and-mask approach still constructs a square matrix; the mask alone does not reduce its quadratic storage. It implements the local pattern, not a complete Longformer model with additional global attention or specialized kernels.

Inputs

Implement sliding_window_attention(Q, K, V, window_size).

Q, K and V: floating-point tensors of shape (seq_len, d_k). Both dimensions are positive.
window_size: a nonnegative integer radius, which may exceed the sequence length.

Inputs use compatible dtypes and, in PyTorch, the same device. Allocate any mask on that device. No batching, padding mask or dropout is required.

Output

Return a tensor of shape (seq_len, d_k) containing the weighted value vectors. Leave all input tensors unchanged.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Share keys and values across attention heads

Multi-query attention lets several query heads share one key head and one value head. Grouped-query attention shares keys and values within groups of query heads. For example, eight query heads and two key/value heads can form two groups of four query heads each. Queries remain distinct within each group.

Follow the specified mapping between query and key/value heads. Their counts must support that grouping. Sharing keys and values reduces their storage, including the cache used during generation.

Restrict which keys each query can read

Sliding-window attention limits a query to nearby keys. A symmetric window can include positions on either side; a causal window includes only permitted past positions and, where specified, the current position.

Masking a full score matrix changes which positions contribute, but still requires storage for that full matrix. To avoid quadratic score storage, the implementation must also avoid creating the full matrix.

Compute attention a block at a time

For each query, a blockwise calculation tracks three quantities: the largest score seen so far, the sum of exponentials relative to that maximum, and the sum of values weighted by those exponentials.

If a new block raises the maximum from m to m_new, multiply both old sums by exp(m - m_new) before adding terms from the new block. This puts old and new terms on the same scale. Divide the final weighted-value sum by the final exponential sum to get the attention output.

Do not simply average the attention outputs of independently normalized blocks. Their total probability weights generally differ. Handle fully masked blocks without subtracting negative infinity from itself, and follow the task's rule for a query with no allowed keys.

With valid rows, this method matches dense attention up to rounding. FlashAttention also organizes computation and memory transfers to use hardware efficiently; implementing equivalent arithmetic alone does not provide the same speedup. See the FlashAttention paper.

Detecting runtime…