FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention

FlashMask-3. Efficient and Expressive Mask-Aware Distributed Attention. By Guoxia Wang and eight others, from Baidu. Published on the thirteenth of October. Its code repository is worked out: about one hundred and twenty source files, with a readme and a licence.

Long-context training of large language and multimodal models combines two things: structured attention masks, such as document, prefix-LM, share-question, sliding-window and block-sparse masks, and context parallelism spread across many nodes. Together they make the mask a dominant bottleneck. The earlier FlashMask handled the first half well. Its column-wise sparse mask compresses these structured patterns from quadratic to linear size, using four boundary vectors along the query axis. But it was implemented for Ampere GPUs. It does not use the warp-specialized TMA and WGMMA pipeline that FlashAttention-3 offers on Hopper, and it does not support distribution at all.

The key observation is about how that mask is laid out. It is indexed by key positions, yet what it stores are interval boundaries along the query axis. Because of that, under any multi-chunk context-parallel assignment, working out each rank's local piece of the mask reduces to a clip, shift and max on the boundary vectors. The authors show this provably needs zero communication, and the result stays compatible with the single-GPU kernel.

FlashMask-3 is built on that. On a single GPU, it adds a mask-aware Hopper kernel with three-role warp specialization, sitting atop the FlashAttention-3 pipeline. For distribution, it adds three pieces: the zero-communication mask localization, a sparsity-aware load balancer called HCC-IPO, and a mask-aware overlap of compute and communication that prunes communication based on the mask. Its forward all-gather uses hierarchical routing that is aware of the network topology.

On a single H100, across twelve masks and contexts from four thousand to one hundred and twenty-eight thousand tokens, achieved TFLOPs improve by forty-five to one hundred and forty-one percent over FlashMask, by thirty-four to seventy-one percent over FlexAttention, and by two to forty-four percent over MagiAttention.

In the distributed setting, with context-parallel degrees from four to thirty-two and eight thousand tokens of sequence per device, it reaches thirty-one to seventy-five percent efficiency in overlapping compute with communication. Throughput improves over the Megatron-Core ring variant by one point eight zero to three point seven one times, and over MagiAttention by one point zero nine to one point four six times.

End to end, training a mixture-of-experts model of one hundred and eighteen billion parameters at a context of one hundred and twenty-eight thousand tokens on one hundred and twenty-eight H100 GPUs, throughput rises by seven point three to eleven percent over a Dual Chunk all-gather baseline, and the attention forward and backward passes run twenty-four point five to forty-four point nine percent faster.

The whole distributed design hangs on one property of a data structure. A mask that stores query-axis intervals can be cut up for each rank without any rank talking to another, and that turns what had been the bottleneck into a small local computation.