Fugu-MT 論文翻訳(概要): FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention

論文の概要: FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention

arxiv url: http://arxiv.org/abs/2610.17901v1
Date: Tue, 13 Oct 2026 07:55:51 GMT
ステータス: 翻訳完了
システム内更新日: 2026-10-17 08:15:50.212575
Title: FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention
Title（参考訳）: FlashMask-3: 効率的かつ表現力のあるマスク対応分散注意
Abstract要約: FlashMask-3は、FlashAttention-3のTMA/WGMMAパイプライン上に3ロールのワープを特化したマスク対応のHopperカーネルを導入した。 4K--128Kコンテキストで12のマスクにまたがる1つのH100では、FlashMask、FlexAttention、MagiAttentionよりも45-141%、34-71%、および2--44%のTFLOPを実現している。
参考スコア（独自算出の注目度）: 4.549299243411132
License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
Abstract: 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.
Abstract（参考訳）: 大規模言語とマルチモーダルモデルの長いコンテキストトレーニングでは、マルチノードのコンテキスト並列性を備えたアテンションマスク(ドキュメント、プレフィックス-LM、シェア-クエスト、スライディング-ウインドウ、ブロック-スパース)が組み合わされ、マスクが支配的なボトルネックとなっている。 FlashMaskのカラムワイドスパースマスクは、4つのクエリ軸境界ベクトルを介して$O(n^2)$から$O(n)$に圧縮し、広い構造化パターンをカバーする。 Ampere GPU向けに実装されているが、FlashAttention-3のTMA/WGMMAワープ専用パイプラインをホッパーで利用せず、配布もサポートしていない。カラムワイズスパースマスクはキー位置でインデックスされるが、クエリ軸間隔境界を格納する。マルチチャンクコンテキスト並列割り当てでは、シングルGPUカーネルと互換性のある境界ベクトル上で、ランクごとのマスクのローカライゼーションを証明可能なゼロコミュニケーション \emph{clip--shift--max} に還元する。 FlashMask-3では、FlashAttention-3のTMA/WGMMAパイプライン上に3ロールのワープ専門化を備えたマスク対応ホッパーカーネルと、マスク対応分散拡張(ゼロコミュニケーションマスクローカライゼーション、スペーサリティ対応ロードバランサ(HCC-IPO)、マスク駆動通信プルーニングによるマスク対応計算通信オーバランサのオーバーラップなど)が導入されている。 4K--128Kコンテキストで12のマスクにまたがる1つのH100では、FlashMask、FlexAttention、MagiAttentionよりも45-141%、34-71%、および2--44%のTFLOPを実現している。コンテクスト並列度 4--32 では、デバイス毎の 8K シーケンスで計算通信のオーバーラップ効率が 31-75% に達し、Megatron-Core CP リング、Megatron-Core CP の全ガターでのスループットが向上し、MagiAttention が $1.80\times$--$3.71\times$, $1.71\times$--$2.92\times$, $1.09\times$--$1.46\times$ に向上した。 128 H100 GPU上で118Kの118B-A14B MoEのエンドツーエンドトレーニングでは、Dual Chunkの全ガターベースラインでスループットを7.3--11.0%向上し、前向きに24.5-44.9%高速化した。

論文の概要: FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention

関連論文リスト