FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention
FlashMask-3 adds Hopper and distributed mask-aware attention, raising long-context training throughput over prior kernels.
FlashMask-3 extends column-wise sparse attention masks to Hopper GPUs and multi-node context parallelism. It uses a three-role warp-specialized kernel on FlashAttention-3's TMA/WGMMA pipeline, zero-communication mask localization, a sparsity-aware load balancer, and mask-driven communication pruning. On one H100, across 12 masks at 4K-128K context, it improves achieved TFLOPs by 45-141% over FlashMask, 34-71% over FlexAttention, and 2-44% over MagiAttention. At context-parallel degrees 4-32 it reaches 31-75% compute-communication overlap, and training a 118B-A14B MoE at 128K on 128 H100s raises throughput 7.3-11.0% over a Dual Chunk all-gather baseline.
- FlashMask-3 adds a mask-aware Hopper kernel on FlashAttention-3's TMA/WGMMA pipeline.
- Mask localization across context-parallel ranks needs zero communication via clip-shift-max.
- Single-H100 TFLOPs rise 45-141% over FlashMask across 12 masks at 4K-128K context.
- On 128 H100s, a 118B-A14B MoE at 128K gains 7.3-11.0% end-to-end throughput.
Full article228 words · extracted from huggingface.co · click to collapse
Long-context training of large language and multimodal models couples structured attention masks (document, prefix-LM, share-question, sliding-window, block-sparse) with multi-node context parallelism, making the mask a dominant bottleneck. FlashMask's column-wise sparse mask compresses them from O(n^2) to O(n) via four query-axis boundary vectors, covering broad structured patterns. Implemented for Ampere GPUs, it neither exploits FlashAttention-3's TMA/WGMMA warp-specialized pipeline on Hopper nor supports distribution. The column-wise sparse mask is indexed by key positions but stores query-axis interval boundaries. Under any multi-chunk context-parallel assignment, it reduces per-rank mask localization to a provably zero-communication clip--shift--max on the boundary vectors, compatible with the single-GPU kernel. FlashMask-3 introduces a mask-aware Hopper kernel with three-role warp specialization atop FlashAttention-3's TMA/WGMMA pipeline, plus a mask-aware distributed extension: zero-communication mask localization, a sparsity-aware load balancer (HCC-IPO), and a mask-aware compute-communication overlap with mask-driven communication pruning, whose forward all-gather uses topology-aware hierarchical routing. On a single H100 across 12 masks at 4K--128K context, it improves achieved TFLOPs over FlashMask, FlexAttention, and MagiAttention by 45--141%, 34--71%, and 2--44%. At context-parallel degrees 4--32 with 8K sequence per device, it attains 31--75% compute-communication overlap efficiency and improves throughput over Megatron-Core CP ring, Megatron-Core CP all-gather, and MagiAttention by 1.80times--3.71times, 1.71times--2.92times, and 1.09times--1.46times. In end-to-end training of a 118B-A14B MoE at 128K on 128 H100 GPUs, it raises throughput by 7.3--11.0% over a Dual Chunk all-gather baseline, with attention forward+backward 24.5--44.9% faster.
Text extracted automatically; images, tables and formatting may be missing. Original: https://huggingface.co/papers/2610.17901