FlashMask-3 kernel claims 3.71× attention speedup at 128K

Scale a context-parallel training run from 4 to 32 ranks and FlashMask-3's distributed attention clocks 1.80× to 3.71× the throughput of Megatron-Core CP ring at an 8K local sequence per device, according to the authors' arXiv paper. Against Megatron-Core CP all-gather the margin is 1.71× to 2.92×; against MagiAttention it narrows to 1.09× to 1.46×.

The trick is in the mask. FlashMask's column-wise sparse mask "compresses such masks from O(n²) to O(n) via four query-axis boundary vectors and covers a broad family of structured patterns," the abstract says. FlashMask-3 lifts that into a Hopper kernel with a three-role warp specialization on the FlashAttention-3 TMA/WGMMA pipeline, and bolts on "zero-communication mask localization, a sparsity-aware load balancer (HCC-IPO), and a mask-aware compute-communication overlap with mask-driven communication pruning."

On a single H100 across 12 structured masks at context lengths from 4K to 128K, achieved TFLOPs improve over FlashMask, FlexAttention, and MagiAttention by 45% to 141%, 34% to 71%, and 2% to 44% respectively. The end-to-end picture is more measured: training a 118B-A14B MoE at a 128K context on 128 H100 GPUs raises throughput 7.3% to 11.0% over a Dual Chunk all-gather baseline, with attention forward+backward 24.5% to 44.9% faster.

The pitch is explicit about what it replaces. FlashMask, the authors write, "does not exploit FlashAttention-3's TMA/WGMMA warp-specialized pipeline on Hopper GPUs, and offers no distributed support." Code lives in the PaddlePaddle organisation on GitHub.