mlsanityData to 14 Oct 2026

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

arXiv 2610.17901 ↗First posted 13 Oct 2026cs.LGcited by 0

Not on the map. Why?

Abstract (arXiv)

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(n2) 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.80×--3.71×, 1.71×--2.92×, and 1.09×--1.46×. 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.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 21 of its 28 references that are in the map. First posted on arXiv 13 Oct 2026.

Builds on

21 of 28 references.

References · 21

  1. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  2. 2020Longformer: The Long-Document Transformercited by 3,413
  3. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
  4. 2019Generating Long Sequences with Sparse Transformerscited by 1,546
  5. 2025Emerging Properties in Unified Multimodal Pretrainingcited by 897
  6. 2021Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMcited by 632
  7. 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 519
  8. 2023Ring Attention with Blockwise Transformers for Near-Infinite Contextcited by 457
  9. 2025Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attentioncited by 407
  10. 2023DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Modelscited by 245
  11. 2025MoBA: Mixture of Block Attention for Long-Context LLMscited by 198
  12. 2018Online normalizer calculation for softmaxcited by 162
9 more
  1. 2021Self-attention Does Not Need O(n2) Memorycited by 154
  2. 2024USP: A Unified Sequence Parallelism Approach for Long Context Generative AIcited by 61
  3. 2024FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusioncited by 57
  4. 2023Striped Attention: Faster Ring Attention for Causal Transformerscited by 54
  5. 2023DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Trainingcited by 29
  6. 2024FlashMask: Efficient and Rich Mask Extension of FlashAttentioncited by 20
  7. 2026ERNIE 5.0 Technical Reportcited by 11
  8. 2025DCP: Addressing Input Dynamism In Long-Context Training via Dynamic Context Parallelismcited by 7
  9. 2025Mesh-Attention: A New Communication-Efficient Distributed Attention with Improved Data Localitycited by 2

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025DAM: Dynamic Attention Mask for Long-Context Large Language Model Inference Acceleration
  2. 2026Salca: A Sparsity-Aware Hardware Accelerator for Efficient Long-Context Attention Decoding
  3. 2025BlindSight: Harnessing Sparsity for Efficient Vision-Language Models
  4. 2026FlashMemory-DeepSeek-V4: Lightning Index Ultra-Long Context via Lookahead Sparse Attention
  5. 2026AB-Sparse: Sparse Attention with Adaptive Block Size for Accurate and Efficient Long-Context Inference
  6. 2026FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving
  7. 2026SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers
  8. 2025Polar Sparsity: High Throughput Batched LLM Inferencing with Scalable Contextual Sparsity