Block-Sparse Attention

A sliding window alone needs $n/w$ layers for information to cross the sequence. Block-sparse patterns fix that by adding a second component: a small set of positions that every query may attend to and that may attend to everything.

Zaheer et al. make the graph-theoretic version of that argument: local plus random plus global edges give a sparse graph whose diameter is small, and the resulting attention is a universal approximator of sequence functions and Turing-complete, matching dense attention’s theoretical guarantees at a fraction of the edges. The theory is reassuring rather than decisive — it says the pattern can express what dense attention expresses, not that training finds it.

Blocks, not entries

The sparsity has to be structured at the granularity the hardware works in. A GPU’s matmul unit operates on tiles; skipping scattered individual entries within a tile saves nothing, because the tile is computed either way.

With $b = 64$ or $128$, each surviving block is a dense matmul of a size the hardware likes, and the pattern is a list of block coordinates rather than a mask tensor. This is why the implementations are kernels rather than masks: the saving is in blocks never launched.

Why decoders mostly stopped

Two things happened. FlashAttention removed the memory term from exact attention, so the reason to approximate dropped from “cannot fit” to “would be faster” — a much weaker reason. And the patterns interact badly with what decoders turned out to need: they are hand-designed per task, they complicate the KV cache, and a fixed stride is a poor prior for text, where the relevant earlier token is wherever it happens to be.

What survived is the part that was never really about sparsity: the global tokens. A handful of positions every query can see, whatever else is masked, is the same structure that attention sinks turn out to provide for free — and that is now more often obtained by pinning a few tokens in the cache than by designing a pattern.

import torch
from torch import Tensor


def block_pattern(
    n_blocks: int, window: int = 3, stride: int = 8, n_global: int = 2
) -> Tensor:
    """Which (query block, key block) pairs to compute. [Nb, Nb] bool, causal."""
    idx = torch.arange(n_blocks)
    q, k = idx[:, None], idx[None, :]

    local = (q - k).abs() <= window
    strided = (k % stride) == 0
    globals_ = (k < n_global) | (q < n_global)

    return (local | strided | globals_) & (k <= q)


def cost(pattern: Tensor) -> float:
    """The only number that matters: fraction of blocks actually launched."""
    causal = torch.ones_like(pattern).tril()
    return (pattern.sum() / causal.sum()).item()

Build the pattern once per sequence length and keep it. The kernel then iterates the nonzero coordinates, and correctness rests on one thing worth asserting directly: every query block must have at least one key block, or its softmax normalises over nothing and the row comes back NaN. The k <= q term makes that a real risk at the first block, which is one more reason the global set is usually the first few positions.