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

📄 arXiv: 2610.17901v1 📥 PDF

作者: Guoxia Wang, Qianyue He, Siming Wu, Haoyang Xie, Qiwen Bao, Jinle Zeng, Jiabin Yang, Tian Wu, Dianhai Yu

分类: cs.LG

发布日期: 2026-10-13

备注: Published as a main conference paper at NeurIPS 2026

💡 一句话要点

提出FlashMask-3以解决长上下文训练中的注意力掩码瓶颈问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文训练 注意力机制 多模态模型 稀疏掩码 分布式计算 GPU优化 模型训练

📋 核心要点

  1. 现有方法在长上下文训练中，注意力掩码的复杂度导致性能瓶颈，限制了模型的扩展性。
  2. 论文提出FlashMask-3，通过列式稀疏掩码和掩码感知的Hopper内核，优化了多节点上下文并行性，降低了通信开销。
  3. 实验结果显示，FlashMask-3在多个基准测试中显著提升了计算效率和吞吐量，尤其在大规模模型训练中表现优异。

📝 摘要（中文）

长上下文训练的大型语言和多模态模型将结构化注意力掩码与多节点上下文并行性结合，使得掩码成为主要瓶颈。FlashMask通过列式稀疏掩码将其复杂度从O(n^2)压缩至O(n)，并在Ampere GPU上实现。FlashMask-3引入了掩码感知的Hopper内核和分布式扩展，显著提高了计算效率和吞吐量。在128个H100 GPU上进行的118B-A14B MoE的端到端训练中，吞吐量提升了7.3%至11.0%。

🔬 方法详解

问题定义：论文要解决长上下文训练中注意力掩码的复杂度问题，现有方法在多节点并行时面临通信瓶颈，影响整体性能。

核心思路：FlashMask-3通过引入列式稀疏掩码，将掩码的复杂度从O(n^2)降低到O(n)，并设计掩码感知的Hopper内核，以优化计算和通信效率。

技术框架：整体架构包括列式稀疏掩码的生成、掩码感知的计算内核、分布式负载均衡器（HCC-IPO）以及掩码驱动的通信修剪，确保在多节点环境下高效运行。

关键创新：最重要的技术创新在于引入了掩码感知的Hopper内核和零通信掩码定位，显著提升了计算和通信的重叠效率，与现有方法相比，减少了通信开销。

关键设计：在参数设置上，FlashMask-3优化了掩码的存储结构，采用了基于关键位置的索引和查询轴间隔边界的存储方式，确保了高效的掩码处理。具体的损失函数和网络结构细节在论文中进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，FlashMask-3在单个H100上处理128K上下文时，相较于FlashMask、FlexAttention和MagiAttention，TFLOPs提升了45%至141%。在128个H100 GPU上进行的118B-A14B MoE训练中，吞吐量提升了7.3%至11.0%，注意力前向和反向计算速度提升了24.5%至44.9%。

🎯 应用场景

该研究的潜在应用领域包括大型语言模型的训练、图像处理和多模态学习等。通过优化长上下文的处理能力，FlashMask-3可以在自然语言处理、计算机视觉等领域实现更高效的模型训练，推动相关技术的发展与应用。

📄 摘要（原文）

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 \emph{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\times$--$3.71\times$, $1.71\times$--$2.92\times$, and $1.09\times$--$1.46\times$. 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.