33. Top-k Sparse Attention

Problem

Sliding-window attention makes attention sparse by position (only nearby tokens are visible); an alternative is to make it sparse by content: let each query look at the whole causal history, score every candidate key as usual, but keep only its k highest-scoring keys and discard the rest, regardless of how far back they are. This is a content-adaptive form of sparse attention — a query can reach far into the past for the few tokens that actually matter to it, while still cutting the number of keys that contribute to its output from O(T) down to a fixed k.

Implement top-k sparse causal self-attention: for each query position i, compute scores against every causally valid key (j <= i) as usual, then zero out the attention weight for every key except the top_k highest-scoring ones (if fewer than top_k keys are causally valid yet, keep all of them). Softmax is applied only over the surviving top-k scores.

Source: src/33_topk_sparse_attention.py

def topk_sparse_attention(
    q: torch.Tensor,  # (B, H, T, D)
    k: torch.Tensor,  # (B, H, T, D)
    v: torch.Tensor,  # (B, H, T, D)
    top_k: int,
) -> torch.Tensor:  # (B, H, T, D), causal
    ...

>>> q = k = v = torch.randn(1, 1, 6, 4)
>>> out = topk_sparse_attention(q, k, v, top_k=2)
>>> out.shape
torch.Size([1, 1, 6, 4])
# each query attends to only its 2 highest-scoring causally-valid keys

Step-by-Step Approach

1. Compute ordinary scaled dot-product scores over the whole sequence and apply the standard causal mask first (j <= i, else -inf) — top-k selection must only ever choose among causally valid keys.
2. Clamp top_k against the sequence length: eff_k = min(top_k, T), since torch.topk cannot ask for more entries than exist along a dimension.
3. Use scores.topk(eff_k, dim=-1) to get, per query row, the eff_k highest scores and their key indices — this works row-by-row automatically since topk operates along the last dimension.
4. Reconstruct a full (T, T) sparse score tensor that starts as all -inf, then use scatter_ to write the top-k values back in at their original column indices, leaving everything else at -inf.
5. Recognize the early-row edge case handles itself: for a query with fewer than eff_k causally-valid keys, some of its "top" picks will themselves be -inf (there weren't enough real scores) — scattering an -inf back into an already--inf tensor is a no-op, so nothing needs special-casing.
6. Softmax over the sparse scores and matmul with V as usual; the surviving top-k entries get real (renormalized) probability mass, everything else gets exactly zero.

The key insight is that topk + scatter_ is the general pattern for "keep only the best k entries per row, mask the rest" — and that letting the causal mask's -inf values flow naturally through the top-k selection (rather than special-casing rows with too little history) keeps the early-sequence edge case correct for free.

Reference solution

import math

import torch


def topk_sparse_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    top_k: int,
) -> torch.Tensor:
    B, H, T, D = q.shape
    scale = 1.0 / math.sqrt(D)

    scores = torch.matmul(q, k.transpose(-1, -2)) * scale  # (B, H, T, T)

    # first apply the ordinary causal mask -- top-k selection must only ever
    # choose among causally valid keys, never a future one
    qi = torch.arange(T, device=q.device).unsqueeze(1)
    kj = torch.arange(T, device=q.device).unsqueeze(0)
    causal = kj <= qi
    scores = scores.masked_fill(~causal.view(1, 1, T, T), float("-inf"))

    # clamp k so topk never asks for more entries than exist along the last dim
    eff_k = min(top_k, T)
    topk_vals, topk_idx = scores.topk(eff_k, dim=-1)  # (B, H, T, eff_k) each

    # rebuild a full (B, H, T, T) score tensor that's -inf everywhere except at
    # the top-k positions, where it holds the original score -- scatter_ writes
    # topk_vals into the positions named by topk_idx along the last dimension.
    # For early rows with fewer than eff_k causally-valid (finite) keys, some
    # of the "top" entries selected will themselves be -inf; scattering an
    # -inf back into an already -inf tensor is a no-op, so this naturally
    # reduces to "keep everything valid" without any special-casing.
    sparse_scores = torch.full_like(scores, float("-inf"))
    sparse_scores.scatter_(-1, topk_idx, topk_vals)

    attn = torch.softmax(sparse_scores, dim=-1)
    return torch.matmul(attn, v)


TEST_CASES = [
    {"name": "top_k smaller than sequence length", "B": 1, "H": 1, "T": 8, "D": 4, "top_k": 3},
    {"name": "top_k=1 (each query attends to a single best key)", "B": 1, "H": 2, "T": 6, "D": 4, "top_k": 1},
    {"name": "top_k >= T (degenerates to full causal attention)", "B": 2, "H": 1, "T": 5, "D": 4, "top_k": 100},
]


def main():
    torch.manual_seed(0)

    for i, case in enumerate(TEST_CASES):
        B, H, T, D, top_k = case["B"], case["H"], case["T"], case["D"], case["top_k"]
        print(f"Test {i}: {case['name']} (B={B}, H={H}, T={T}, D={D}, top_k={top_k})")
        q = torch.randn(B, H, T, D)
        k = torch.randn(B, H, T, D)
        v = torch.randn(B, H, T, D)

        out = topk_sparse_attention(q, k, v, top_k)
        print(f"  out.shape={tuple(out.shape)}, each query keeps <= {min(top_k, T)} keys instead of up to {T}")
        assert out.shape == (B, H, T, D)

        scale = 1.0 / math.sqrt(D)
        scores = torch.matmul(q, k.transpose(-1, -2)) * scale
        qi = torch.arange(T).unsqueeze(1)
        kj = torch.arange(T).unsqueeze(0)
        causal = kj <= qi
        scores = scores.masked_fill(~causal, float("-inf"))
        eff_k = min(top_k, T)
        topk_vals, topk_idx = scores.topk(eff_k, dim=-1)
        sparse_scores = torch.full_like(scores, float("-inf"))
        sparse_scores.scatter_(-1, topk_idx, topk_vals)
        attn = torch.softmax(sparse_scores, dim=-1)
        expected = torch.matmul(attn, v)

        torch.testing.assert_close(out, expected, atol=1e-4, rtol=1e-4)
        print("PASSED")


if __name__ == "__main__":
    main()

Key Functions & Tricks

• scores.topk(eff_k, dim=-1) — returns both the k largest values and their indices along the last dimension, computed independently per row (per query).
• torch.full_like(scores, float('-inf')).scatter_(-1, idx, vals) — rebuilds a full-size tensor that's -inf everywhere except at the selected indices, which is the standard way to turn a top-k selection back into a same-shaped masked tensor.
• min(top_k, T) — guards against requesting more entries than exist along the dimension being reduced, which would otherwise raise a runtime error.
• masked_fill for the causal mask, applied before topk — ordering matters: causal masking has to happen before top-k selection, not after, or a query could end up attending to a future key.
• torch.softmax over a mostly-(-inf) tensor — relies on the same max-subtraction stability guarantee as ordinary masked softmax; no separate numerical handling is needed for the sparse case.

How to Recognize This Pattern