Sparse attention

Full causal attention compares every token with every earlier token. The work and the score memory both grow with $T^2$, and in a trained model most of those scores end up near zero anyway. Sparse attention decides ahead of time which scores to skip. The question for EXP-017 was what that costs on code, which refers much further back than prose does, like a variable defined 200 lines up is still in scope.

A window, plus a few global columns

A sliding window of size $w$ lets query $i$ read keys $j$ with $0 \le i - j < w$. Each row has at most $w$ allowed scores, so the cost grows with $T \cdot w$ instead of $T^2$. Longformer and Mistral use this pattern.

But a window loses everything older than $w$. A stride $s$ adds some of it back cheaply, every key whose position is a multiple of $s$ stays visible to all later queries. Those columns work like landmarks any token can reach.

So what happened ?

• With the window at 0 the mask is the full causal triangle, density 100 percent.
• A window of 4 turns it into a band. The last row only reads 4 keys, and density drops as the sequence grows, because the band's width is fixed while the triangle's area grows with $T^2$.
• A stride adds vertical stripes. Every $s$-th column stays visible all the way down.

One mask function for every pattern

The Day 1 causal mask and the Day 3 sparse masks come from the same function. With window = 0 it returns the plain triangle:

def attention_mask(
    length: int, window: int, stride: int, device: torch.device | None = None
) -> Tensor:
    """Causal mask, optionally narrowed to a sliding window with strided global columns."""
    rows = torch.arange(length, device=device).unsqueeze(1)
    columns = torch.arange(length, device=device).unsqueeze(0)
    distance = rows - columns
    allowed = distance >= 0
    if window:
        local = distance < window
        if stride:
            local = local | (columns % stride == 0)
        allowed = allowed & local
    return allowed.view(1, 1, length, length)

The model passes that mask as a boolean attn_mask into scaled_dot_product_attention. That's correct on every backend and every device, but it has a cost, and the note wrote it down as a measurement, not a defect. The moment an explicit mask replaces is_causal=True, SDPA can't use the flash kernel on CUDA anymore and falls back to a slower backend. PyTorch's flex_attention can compile a block-sparse mask into a fast kernel, but it needs torch.compile, and chasing that wasn't what EXP-017 was for.

A small probe, flash_accepts, records which backend actually ran. It has to run on the training device, because PyTorch's CPU flash path accepts an explicit mask and the CUDA kernel doesn't, so a probe on the laptop would report the wrong backend for a GPU run.

Worse on code, and no faster at this length

The note's hypothesis had two parts:

1. At a 512-token trained context, a 128-token window loses measurable bits per byte on code and less on prose.
2. The windowed run is no faster than full causal. Losing the flash kernel costs more than the skipped scores save at 512 tokens. The speed win belongs to lengths this lab never trained at.

Baseline, full causal attention at the same budget and context. Stop condition, windows of 128 and 256, strides of 0 and 64, three seeds, reporting code and prose bits per byte, seconds per step, the backend that ran and the computed density.

Built, tested, not scheduled

The plan changed before any Day 3 run. None of the sparse patterns help the tool-use harness the new plan is building toward. The code stays behind flags that default to off, so every other model is unaffected, and the mask tests still run.

What we didn't build, and why

• flex_attention kernels. Needs torch.compile, and the Day 3 GPU is a T4, where the gain was untested.
• Dilated or random patterns. The window and stride answered the question the plan asked.
• The read on how Longformer and Mistral pick a window size. Listed in the note as still to read, and it wasn't read before the plan changed.