FlexAttention

Context. FlashAttention made exact attention fast by fusing the two matmuls and the softmax into one tiled kernel that never writes the $T\times T$ score matrix to memory, but each such kernel supports a short list of variants, and a new mask (packed documents, a sliding window, a prefix) meant a handwritten kernel or a dense boolean mask whose memory traffic grows as $T^2$. FlexAttention, from the PyTorch team, observes that almost every variant is a pointwise edit of the scores before the softmax, so the user writes the edit as a few lines of Python and the compiler inlines it into a handwritten fused template. A mask also tells the kernel which scores need not be computed. Summarized at the level of 128 by 128 tiles, it lets the kernel skip every fully masked tile, so attention costs what the visible tiles cost, whatever the sequence length. That is what let the speedrun train on 64K-token packed sequences in which each token sees only its own document.

docs = torch.cumsum(tokens == EOT, dim=0)    # document id of every position

def mask_mod(b, h, q_idx, kv_idx):
    causal = q_idx >= kv_idx
    same_doc = docs[q_idx] == docs[kv_idx]
    in_window = q_idx - kv_idx < window      # 1024 in record #12
    return causal & same_doc & in_window

block_mask = create_block_mask(mask_mod, None, None, T, T)
y = flex_attention(q, k, v, block_mask=block_mask)

A tile-level summary of the mask lets the kernel skip fully masked tiles. create_block_mask evaluates mask_mod over the index grid and stores, for each row of query tiles, how many tiles are not fully masked and their column indices. The kernel loops over that list only, so $$\text{work}\;\propto\;\#\{\text{tiles containing at least one visible score}\},$$ and the summary takes $O\big((T/128)^2\big)$ small integers. For the speedrun's shape the saving is large (my count): in one 65,536-token sequence causal attention touches 131,328 tiles, a causal 1024-token window touches 4,572 (3.5% of them), and causal attention inside independent 1024-token chunks touches 2,304, and document boundaries remove more. The same rule limits the gain: a mask whose visible scores are scattered over many tiles gains little, and a tile with one visible score costs a full tile. In the authors' benchmarks FlexAttention runs at 0.68 to 1.43 times FlashAttention-2 where that kernel supports the variant (the backward pass down to 0.86 times), and 5 to 8 times faster than PyTorch SDPA with a materialized mask, which must read a $B\times N\times N$ boolean mask; on packed sequences SDPA's throughput drops 25% from 2K to 8K tokens, while FlexAttention, needing only a document-id vector and the tile lists, keeps its throughput.

Record #12: document-masked packed sequences reduced the tokens needed by 37.5%. Record #11 trained on 1024-token chunks taken from the token stream at arbitrary points, so attention crossed document boundaries and a token near the start of a chunk had almost no context. Record #12 (7.2 to 5.03 minutes, 30% less time) keeps the same 524,288 tokens per step but gives each of the eight GPUs one 65,536-token sequence with the mask above. The step got slower, 145 to 162 ms, and the run needed 1875 steps instead of 3000, 37.5% fewer tokens to reach 3.28. Two things changed together, and the record does not separate them: no attention goes into unrelated documents, and every token sees up to 1024 tokens of its own document wherever it falls. The tile lists kept the cost low: about twice the attention work of 1024-token chunks, where dense causal attention would cost 57 times as much. The comparison also spans a PyTorch upgrade (2.5.1 to a 2.6 nightly), and validation moved to the same document-masked 64K format, which the rules allow but which gives validation tokens more context, so part of the gain may come from easier evaluation.

Record #13: growing the window over training saved 125 steps. Because mask_mod reads the window from a tensor, it can change without recompiling. Record #13 grows it linearly from 64 to 1792 tokens in steps of 64, which reduced the steps from 1875 to 1750 and the time from 5.03 to 4.66 minutes. The step time barely moved, so the gain is again loss per step, and the record does not separate a short window early from a longer one (1792) at the end. The same record also retuned the optimizer (Adam $\beta_1$, the Muon learning rate, the cooldown and the momentum warmup) and validated with the final 1792-token window, so the schedule's own share is not isolated. The only source is a tweet; the closest paper is the sequence-length warmup of Li et al. (2021).