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

Overview

The attention mask has become the bottleneck

Packed long-context batches carry structured masks (document, prefix-LM, share-question, sliding-window, block-sparse). Under multi-node context parallelism (CP), the mask itself dominates step time.

Three costs appear once masks meet CP:

Materialization: dense n×n masks do not fit long contexts.
Mask-aware scheduling: kernels must skip masked work without stalling the pipeline.
Mask-induced imbalance: equal-length shards carry very unequal work.

FlashMask already compresses masks to O(n) with four column-wise boundary vectors (LTS, LTE, UTS, UTE), but it targets Ampere and has no distributed support.

Key observation. The column-wise mask is indexed by key but stores query-row boundaries. That single asymmetry makes per-rank localization a zero-communication elementwise map, and the same full / partial / unmasked tile classification can drive every layer of the stack.

one tile classification

§3.1

Mask-aware Hopper kernel

Three-role warp specialization skips fully-masked tiles and masks only partial ones, off the WGMMA critical path.

§3.2

Zero-communication localization

clip–shift–max rewrites global boundaries into rank-local ones: lossless and without communication.

§3.3

Sparsity-aware balancing

Tile costs {0, α_p, α_u} give per-chunk loads; HCC-IPO balances them and proves the result optimal whenever the peak load reaches the lower bound.

§3.4

Mask-aware overlap

Hierarchical all-gather and segmented reduce-scatter hide behind attention; fully-masked chunks are never sent.

tile classes: unmaskedpartialfully-masked

Method · §3.1

A third warp role that reads the mask

FA3's producer warp group has 128 threads, yet one thread is enough to issue each TMA load, so 96 threads sit idle every iteration. FlashMask-3 turns them into a KV Index Generator that classifies tiles from the boundary vectors, so masking never enters the WGMMA critical path.

Mask

Q-block (click a tile row)

64 × 64 mask, 8 × 8 tiles

row = Q-block · column = KV block

unmasked tile partial tile (pale cells masked) fully-masked (skipped)

Tile rows [q₀, q₁) × a KV block, using block-wise max/min of the vectors:
skip if q₀ ≥ max LTS or q₁ ≤ min UTE · unmasked if q₁ ≤ min LTS and q₀ ≥ max UTE · otherwise partial.

Q-block 5: KV blocks visited right → left

producer WG · 96 idle threads

KV Index Generator

Scans the boundary vectors and emits a queue of valid KV blocks, each tagged partial or unmasked.

producer WG · TMA

KV Loader

Issues TMA loads only for queued blocks. In CP it also polls the all-gather write pointer (§3.4).

2× consumer WG · WGMMA

MMA Consumer

Runs score, softmax and output, applying the row-interval mask only on partial blocks.

Method · §3.2

Zero-communication mask localization

Keys are all-gathered, so every rank already holds every key column of the O(n) mask. It only has to renumber each row boundary into its own rows. That is a purely local count, so zero bytes cross ranks.

x^(r) = max_i [ clip(x − O_i, 0, C) + i·C·𝟙(clip>0) ]= number of rank r's rows before global row x

Mask

Chunk → rank

Global mask FM_gheld by every rank · O(n)

Local masks FM_l^(r)what each rank's kernel consumes

attended · row owned by Rank 0 attended · row owned by Rank 1 masked selected key column (hover to change)

Theorem 1 (lossless). For any chunk assignment, the local intervals equal the global intervals at the corresponding rows. Cost: O(n_r·N_K) elementwise work, zero communication.

Method · §3.3

Balance the work, not the length

CP cuts the query axis into N = P·κ chunks and gives each rank κ of them. Under structured masks, chunks of equal length carry very different work, and step time is set by the slowest rank. HCC-IPO minimizes that makespan Φ under the cardinality constraint.

Mask

Assignment

Chunks per rank κ

Ranks P

Mask, chunks & owners

unmasked tile (α_u=1) partial tile (α_p=1.3, illustrative) right strip = chunk load w_i

Per-rank load L_j = Σ w_i

each segment is one chunk

0

average load lower bound LB = max(⌈Σw/P⌉, max w)

The full HCC-IPO. Stage 1 (shown above) balances rank compute. Stage 2 swaps chunks one-for-one across nodes to cut the worst inter-node KV traffic while keeping every rank within an ε band of the average. Stage 3 reruns IPO inside each node. Stage 2 only moves chunks across nodes and Stage 3 only within them, so they do not interfere.

Method · §3.4

Hide communication behind sparse attention

CP must all-gather K,V in the forward pass and reduce-scatter dK,dV in the backward pass. FlashMask-3 runs both on 4–8 SMs of high-priority streams, overlaps them with a single unsplit attention kernel, and never sends a chunk that the mask fully hides (≈50% of AG/RS traffic on causal and prefix-LM document masks). Setup here: 2 nodes × 4 GPUs, viewed from rank R0.

Mask-aware pruning Hierarchical routing

schematic timing model · time in arbitrary units

Experiments · §4

Results

H100 80GB SXM5, 8 GPUs per node, NVLink inside a node and RoCE across nodes. CUDA 12.4, NVSHMEM 3.2.5, BF16, head dim 128. CP runs use weak scaling with 8K tokens per device, B = 32/P, H_q = 64, H_kv = 1.

TFLOPs/s gain of FlashMask-3

geo-mean · whisker = min–max over 12 masks × 4K–128K

Gain over FlashMask grows with length

geo-mean over masks

FlashMask-3 beats all three baselines on every mask and length. The MagiAttention mean excludes Random Eviction, where its block-sparse format is over 50× slower. Gains grow with length (+66.2% → +97.3% over FlashMask) as attention becomes more memory-bound.

Mask

Per-GPU TFLOPs/s vs CP degree

hover to compare

Throughput speedup of FlashMask-3

min–max over masks × CP 4–32

CP 4–8 stays on NVLink; at CP 16 inter-node RoCE appears, and the Megatron variants dip while FlashMask-3 keeps rising (430.5 → 501.6 TFLOPs/s on Causal Doc). Megatron supports no Prefix-LM Document mask.

Mask

Overlap efficiency E = 1 − T_exposed / T_comm

forward vs backward · 32K–256K total

Forward (AG)Backward (AG + RS)

Across masks and CP degrees, 0.31–0.75 of the communication time is hidden behind computation.

Peak / average per-device compute

CP · lower is better · 1.0 = ideal

Tile scheduler on one H100

normalized TFLOPs/s (geo-mean over 12 masks)

IPO phase ablation vs an exact P-way branch-and-bound optimum

200 random instances each (100 for N = 128)

Configuration (N, P, κ, w_max)	LPT subopt.	LPT gap	+P2 subopt.	+P2 gap	+P2+P3 subopt.
Uniform (64, 16, 4, 10³)	99.0%	6.54%	86.0%	6.29%	0%
Uniform wide (64, 16, 4, 10⁴)	100%	8.67%	93.5%	3.42%	0%
Causal (64, 8, 8, 200)	93.5%	2.51%	53.5%	1.39%	0%
Uniform† (128, 16, 8, 200)	85.0%	1.08%	48.0%	0.79%	0%

LPT alone is sub-optimal in 85–100% of instances. Exact pairwise repartition (P2) narrows the gap; adding three-rank repartition (P3) reaches a 100% LB-certified optimal rate on every tested configuration.

118B-A14B MoE, 128K context, 128 H100s on 16 nodes, CP degree 16, 8K tokens per device, H_q = 64, H_kv = 8, document mask. The baseline runs the same FlashMask-3 kernel with Dual Chunk balancing and plain all-gather. Each run is 70 steps; the 10 fastest and 10 slowest are dropped and the remaining 50 averaged. The three datasets have an average mask sparsity of 70%, 80% and 86%.

End-to-end vs attention speedup

full FlashMask-3 over the baseline

End-to-end gain by component

over the baseline · one change at a time

Training throughput, 118B-A14B MoE on 128 H100s

mean ± std over 50 steps · Δ_attn = attention fwd+bwd speedup

	70% sparsity		80% sparsity		86% sparsity
Configuration	tokens/s/GPU	Δ_attn	tokens/s/GPU	Δ_attn	tokens/s/GPU	Δ_attn
Baseline (Dual Chunk + all-gather)	712.6 ± 8.7	–	753.3 ± 34.7	–	812.4 ± 57.1	–
+ HCC-IPO balancing	732.2 ± 24.5	+10.4%	786.2 ± 65.2	+19.1%	858.5 ± 62.1	+10.5%
+ mask-aware overlap	750.8 ± 6.4	+15.9%	785.0 ± 33.6	+13.9%	860.0 ± 55.0	+20.5%
+ both, hierarchical routing off	700.6 ± 28.8	+8.4%	742.0 ± 72.6	+21.5%	828.0 ± 60.6	+23.2%
+ both, w/o chunk-skip pruning	750.0 ± 20.6	+20.5%	788.4 ± 55.0	+25.4%	888.6 ± 69.6	+27.9%
FlashMask-3 (balance + overlap)	764.7 ± 28.8	+24.5%	811.4 ± 75.9	+42.1%	902.1 ± 72.6	+44.9%

End-to-end throughput rises by 7.3–11.0% while attention forward+backward gets 24.5–44.9% faster. The gap is expected: in the full configuration attention is only 43.5% / 25.8% / 24.4% of forward+backward compute, and the end-to-end step also includes gradient synchronization and the optimizer update, which FlashMask-3 does not touch. Hierarchical routing is what lets balancing and overlap compose across nodes: without it the balancer spreads chunks over nodes, inter-node traffic grows, and the combination (−1.7% / −1.5% / +1.9%) falls below overlap alone.

Cite

BibTeX

@inproceedings{wang2026flashmask3,
  title     = {FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention},
  author    = {Wang, Guoxia and He, Qianyue and Wu, Siming and Xie, Haoyang and
               Bao, Qiwen and Zeng, Jinle and Yang, Jiabin and Wu, Tian and Yu, Dianhai},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  url       = {https://arxiv.org/pdf/2610.17901},
  year      = {2026}
}

The attention mask has become the bottleneck

Mask-aware Hopper kernel

Zero-communication localization

Sparsity-aware balancing

Mask-aware overlap

A third warp role that reads the mask

64 × 64 mask, 8 × 8 tiles

Q-block 5: KV blocks visited right → left

KV Index Generator

KV Loader

MMA Consumer

Zero-communication mask localization

Balance the work, not the length

Mask, chunks & owners

Per-rank load Lj = Σ wi

Hide communication behind sparse attention

Results

TFLOPs/s gain of FlashMask-3

Gain over FlashMask grows with length

Per-GPU TFLOPs/s vs CP degree

Throughput speedup of FlashMask-3

Overlap efficiency E = 1 − Texposed / Tcomm

Peak / average per-device compute

Tile scheduler on one H100

IPO phase ablation vs an exact P-way branch-and-bound optimum

End-to-end vs attention speedup

End-to-end gain by component

Training throughput, 118B-A14B MoE on 128 H100s

BibTeX

Per-rank load L_j = Σ w_i

Overlap efficiency E = 1 − T_exposed / T_comm