论文页面 - FlashMask-3：高效且富有表达力的掩码感知分布式注意力

来源：https://huggingface.co/papers/2610.17901

摘要

FlashMask-3 通过一个针对 Hopper 架构优化的稀疏掩码内核以及分布式上下文并行扩展来加速长上下文注意力，消除了掩码通信开销并提升了训练吞吐量。

大语言模型和多模态模型的长上下文训练将结构化注意力掩码（文档、前缀 LM、共享问题、滑动窗口、块稀疏）与多节点上下文并行相结合，使得掩码成为主要瓶颈。FlashMask 的列式稀疏掩码通过四个查询轴边界向量将掩码从 O(n²) 压缩到 O(n)，涵盖广泛的结构化模式。其针对 Ampere GPU 的实现既未利用 FlashAttention-3 在 Hopper 上的 TMA/WGMMA（https://huggingface.co/papers?q=TMA%2FWGMMA）warp 特化流水线，也不支持分布式。该列式稀疏掩码以键位置索引，但存储查询轴的区间边界。在任意多分块上下文并行分配下，它将每个 rank 的掩码定位简化为可证明零通信的对边界向量的 clip-shift-max 操作，且与单 GPU 内核兼容。FlashMask-3 引入了一个掩码感知的 Hopper 内核，在 FlashAttention-3 的 TMA/WGMMA（https://huggingface.co/papers?q=TMA%2FWGMMA）流水线之上实现了三角色 warp 特化，并提供掩码感知的分布式扩展：零通信掩码定位、稀疏度感知的负载均衡器 HCC-IPO（https://huggingface.co/papers?q=HCC-IPO），以及掩码感知的计算-通信重叠（含掩码驱动的通信剪枝），其中前向 all-gather 采用拓扑感知的分层路由。在单张 H100 上覆盖 12 种掩码、4K–128K 上下文长度的测试中，其实测 TFLOPs 相比 FlashMask、FlexAttention 和 MagiAttention 分别提升 45–141%、34–71% 和 2–44%。在上下文并行度 4–32、每设备 8K 序列长度的条件下，达到 31–75% 的计算-通信重叠效率（https://huggingface.co/papers?q=compute-communication%20overlap），吞吐量相比 Megatron-Core CP ring、Megatron-Core CP all-gather 和 MagiAttention 分别提升 1.80–3.71 倍、1.71–2.92 倍和 1.09–1.46 倍。在 128 张 H100 GPU 上对 118B-A14B MoE 模型进行 128K 长度的端到端训练中，相比 Dual Chunk all-gather 基线吞吐量提升 7.3–11.0%，注意力前向+反向加速 24.5–44.9%。

查看 arXiv 页面 | 查看 PDF | 项目页面 | GitHub27 | 加入收藏

在你的 agent 中获取此论文：

hf papers read 2610.17901

没有最新版 CLI？curl -LsSf https://hf.co/cli/install.sh | bash

引用此论文的模型：0

暂无模型关联此论文。

在模型 README.md 中引用 arxiv.org/abs/2610.17901 以从此页面链接。

引用此论文的数据集：0

暂无数据集关联此论文。

在数据集 README.md 中引用 arxiv.org/abs/2610.17901 以从此页面链接。

引用此论文的 Spaces：0

暂无 Space 关联此论文。

在 Space README.md 中引用 arxiv.org/abs/2610.17901 以从此页面链接。

包含此论文的合集：0

暂无合集包含此论文。

将此论文加入合集以从此页面链接。