From the source
A selection-based hierarchical attention that runs the same forward+backward pass ~17× faster than standard attention at 512K context on a single B200, and delivers a 1.4–1.7× end-to-end pretraining speedup at 98K context.
Q, K, V are pooled symmetrically across an L-level pyramid; per-head \ell_2 norms pick a small dense sub-sequence; FlashAttention runs on the gather — no custom sparse attention kernel, no straight-through estimator, no auxiliary loss.
After the sparse stage, a brief standard attention resume converts the checkpoint back into a dense attention model: every recovered run matches or beats dense-from-scratch at the same token budget.
Validated at 530M Llama-3, 16k optimiser steps, 50B tokens, with 1M-token training across 32 B200s under context parallelism.
Long-context pretraining is bottlenecked by attention's quadratic compute cost.
FlashAttention shaves the constants, but the wall is still there: you train at the contexts you can afford.
…