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

Guoxia Wang; Qianyue He; Siming Wu; Haoyang Xie; Qiwen Bao; Jinle Zeng; Jiabin Yang; Tian Wu

cs.LGOct 13, 2026

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

Authors: Guoxia Wang, Qianyue He, Siming Wu, Haoyang Xie, Qiwen Bao, Jinle Zeng, Jiabin Yang, Tian Wu,

Organizations: Baidu Inc.

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 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.

Figures & tables

Figure 1: Overview of FlashMask-3 across four layers of the execution stack.

Figure 2: Three-role warp specialization in the FlashMask-3 Hopper kernel: KV Index Generator (mask scan), KV Loader (TMA load), and MMA Consumer (WGMMA score), connected by three asynchronous pipelines.

Figure 3: Zero-communication mapping of global mask boundaries (LTS, UTE) to rank-local coordinates using the clip–shift–max transform (Eq. 1 ).

Figure 4: Mask-aware hierarchical compute-communication overlap. (a) AG skips fully-masked chunks; compute polls a write_ptr updated by the comm kernel. (b) Hierarchical AG ( $N{=}2,G{=}2$ ): inter-node congruence-group fetch (Phase 1), then intra-node NVLink fetch (Phase 2). (c) Multi-stage backward overlap ( $N{=}2,G{=}8$ ): RS of segment $t$ overlaps AG and grad compute of segment $t{+}1$ ; the tail RS runs non-overlapped.

Figure 5: Kernel FWD & BWD throughput (TFLOPs/s) on a single H100 GPU. Error bars are omitted for legibility; the underlying per-mask measurements with standard deviations are reported in Figure 9 (Appendix D.1.3 ).

Figure 6: Distributed scaling. Per-GPU TFLOPs/s versus CP degree

P

for FlashMask-3, Megatron all-gather, Megatron Ring Attention, and Magi-Attention on three practically relevant mask patterns.

	70% sparsity		80% sparsity		86% sparsity
Configuration	tokens/s/GPU	$\Delta_{\text{attn}}$	tokens/s/GPU	$\Delta_{\text{attn}}$	tokens/s/GPU	$\Delta_{\text{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%

Table 1: End-to-end training throughput of a 118B-A14B MoE model (128K context, CP degree 16,

S_{\text{local}}{=}8

K tokens per device,

H_{q}{=}64

H_{kv}{=}8

, document mask) on 128 H100 GPUs, for three datasets of differing average mask sparsity. Each run covers 70 steps; we discard the 10 highest and 10 lowest throughput steps and report the mean ± one standard deviation over the remaining 50.

\Delta_{\text{attn}}

is the corresponding attention forward+backward speedup over the baseline.

Panel of Figure 7 — Figure 7: (a) Speedup of the Preemptive Persistent Tile (PPT) scheduler over the static scheduler across different mask patterns and sequence lengths on a single H100 GPU. (b) Peak-to-average ratio of per-device computation time under different workload partitioning strategies with Context Parallelism.

Figure 8: Computation-communication overlap efficiency for various mask patterns and CP degrees.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Operation	Time	Memory	Kernel interface
Single-chunk clip (Def. 1 )	$O(N_{K})$	$O(N_{K})$	shared
Multi-chunk shift+merge (Def. 2 )	$O(n_{r}N_{K})$	$O(N_{K})$	shared
Total mask preprocessing	$O(n_{r}N_{K})$	$O(N_{K})$	shared

Appendix

Table 2: Resource profile of the FlashMask-3 distributed mask preprocessing on rank

r

owning

n_{r}

chunks.

N_{K}

denotes the global key length. All operations are elementwise and parallelized over the key dimension.

Configuration	$(N,P,\kappa,w_{\max})$	Subopt	Gap	Subopt	Gap	Subopt	P3 wins
		LPT		LPT+P2		LPT+P2+P3
Uniform	$(64,16,4,10^{3})$	99.0%	6.54%	86.0%	6.29%	0%	172
Uniform (wide)	$(64,16,4,10^{4})$	100%	8.67%	93.5%	3.42%	0%	187
Causal	$(64,8,8,200)$	93.5%	2.51%	53.5%	1.39%	0%	107
Uniform †	$(128,16,8,200)$	85.0%	1.08%	48.0%	0.79%	0%	48

Appendix

Table 3: IPO phase ablation across representative CP configurations. Subopt : fraction of random instances not matching the optimum. Max gap : worst-case relative excess over optimal makespan. P3 wins : cases where Phase 3 improves upon Phase 2. All configurations use 200 random instances (100 for

N\!=\!128

Figure 9: Single-device throughput across mask types and sequence lengths. FlashMask-3 consistently outperforms all baselines across all 12 mask patterns.

Baseline	Scope	Min	Max	Mean
FlexAttention	All	+33.6%	+70.5%	+55.3%
	8K	+33.6%	+62.9%	+50.6%
	128K	+51.9%	+70.0%	+58.9%
Magi-Attention	All	+2.0%	+44.1%	+15.0%
	8K	+12.0%	+44.1%	+25.0%
	128K	+2.0%	+8.9%	+5.9%

Appendix

Table 4: FlashMask-3 throughput improvement over baselines (forward+backward TFLOPs/s). Reported as

(T_{\text{ours}}/T_{\text{baseline}}-1)\times 100\%

Figure 10: Q Head = 64, KV Head = 1. CP-scale throughput comparison (TFLOPs/s) across mask types and CP degrees. Columns: forward, backward, and total (forward+backward). Rows: mask types.

vs. Baseline	Phase	Min	Max	Geo. Mean ± std
MagiAttention	FWD	1.04 $\times$	2.39 $\times$	1.56 ± 0.53 $\times$
	BWD	1.10 $\times$	1.19 $\times$	1.15 ± 0.03 $\times$
	Total	1.09 $\times$	1.46 $\times$	1.26 ± 0.15 $\times$
Megatron Ring	FWD	1.62 $\times$	4.12 $\times$	2.50 ± 0.97 $\times$
	BWD	1.86 $\times$	3.57 $\times$	2.38 ± 0.56 $\times$
	Total	1.80 $\times$	3.71 $\times$	2.42 ± 0.65 $\times$

Appendix

Table 5: FlashMask-3 speedup over distributed baselines (aggregated over all masks

\times

CP sizes).

Figure 11: Comparison of compute load distributions under different scheduling strategies. Dark regions indicate work blocks assigned to each rank. The inset (upper-right of each subplot) shows per-rank compute load, with the dashed line denoting the average. The “peak/avg” in each title reports the peak-to-average load ratio across all 8 ranks; lower values indicate more balanced distributions.

Figure 12: Convergence under context parallelism. (a) Training loss of the two runs; faint lines are per-step values and bold lines an exponential moving average (

\alpha{=}0.1

), and the inset magnifies the raw curves over the last 60 steps. (b) Per-step loss difference between the two runs and its moving average, on the same step axis.

Seq	Mask	FA3	FlashMask-3	FlexAttention	MagiAttention
8K	Full	604.43 ± 1.27	569.84 ± 1.88 (0.94)	394.15 ± 1.34 (0.65)	513.66 ± 8.95 (0.85)
8K	Causal	579.73 ± 0.35	520.85 ± 0.54 (0.90)	341.68 ± 1.45 (0.59)	482.17 ± 1.71 (0.83)
8K	Sliding Window	407.23 ± 1.54	290.30 ± 1.04 (0.71)	218.93 ± 1.73 (0.54)	166.22 ± 0.15 (0.41)
8K	Causal Document	478.96 ± 27.20	377.86 ± 35.26 (0.79)	264.06 ± 18.37 (0.55)	247.92 ± 47.43 (0.52)
8K	Document	546.29 ± 17.93	460.49 ± 31.77 (0.84)	299.46 ± 26.68 (0.55)	352.82 ± 48.32 (0.65)
32K	Full	616.65 ± 0.65	603.63 ± 1.48 (0.98)	408.22 ± 0.02 (0.66)	593.40 ± 1.30 (0.96)

Appendix

Table 6: Single-GPU throughput on five masks supported by FA3 (total forward

+

backward TFLOPs/s, mean ± std over 50 iterations). Parentheses give the ratio to FA3.

Seq	Mask	FlashMask-3	FlashMask	FA3-style Hopper port
8K	Full	596.53 ± 2.26	384.88 ± 1.84 (0.65)	214.01 ± 0.23 (0.36)
8K	Causal	553.31 ± 1.13	373.46 ± 1.37 (0.67)	241.86 ± 0.13 (0.44)
8K	Sliding Window	432.15 ± 6.60	256.63 ± 0.07 (0.59)	35.77 ± 0.07 (0.08)
8K	Causal Document	483.97 ± 29.40	302.03 ± 15.27 (0.62)	64.77 ± 15.83 (0.13)
8K	Document	532.26 ± 24.90	287.89 ± 34.54 (0.54)	55.03 ± 14.14 (0.10)
32K	Full	614.19 ± 1.15	388.92 ± 1.84 (0.63)	214.22 ± 0.08 (0.35)

Appendix

Table 7: Isolating the warp specialization from the Hopper port (forward TFLOPs/s, mean ± std). Parentheses give the ratio to FlashMask-3.

Table 20

Context Parallelism

See all 10 papers

Dec 7, 2025

Block Sparse Flash Attention

Modern large language models increasingly require long contexts for reasoning and multi-document tasks, but attention's quadratic complexity creates a severe computational bottleneck. We present Block Sparse Flash Attention (BSFA), a drop-in replacement that accelerates long-context inference while preserving model quality. Unlike methods that predict importance before computing scores, BSFA computes exact query-key similarities to select the top-k most important value blocks for each query. By comparing per-block maximum scores against calibrated thresholds, we skip approximately 50% of the computation and memory transfers for pruned blocks. Our training-free approach requires only a one-time threshold calibration on a small dataset to learn the per-layer and per-head attention score distributions. We provide a CUDA kernel implementation that can be used as a drop-in replacement for FlashAttention. On Llama-3.1-8B, BSFA achieves up to 1.13x end-to-end speedup on LongBench with only a 1.1% accuracy drop, and up to 1.24x on Needle-in-a-Haystack retrieval at a 1% accuracy drop. The attention kernel itself accelerates by up to 1.38x. We compare BSFA against five recent sparse attention baselines (SpargeAttention, MInference, FlexPrefill, XAttention, and BLASST), and verify the method on Qwen2.5-7B and on A6000 and H100 GPUs. The implementation is available at https://github.com/Danielohayon/Block-Sparse-Flash-Attention.

Daniel Ohayon, Itay Lamprecht, Itay Hubara +3

Technion – Israel Institute of Technology, Haifa, Israel · Intel - Habana Labs · Intel – Habana Labs, Tel Aviv, Israel

GPU Acceleration Self-Attention Long-Context Language Model Inference High-Performance Computing LLM Inference Acceleration Sparse Attention+3

Ask PDF

Sep 30, 2026

LampAttention: Look-Ahead Mixed-Precision FlashAttention for Dedicated Accelerators

While most attention logits can be computed in low precision without degrading numerical stability, current attention kernels fail to exploit this phenomenon. We introduce a novel hardware-algorithm co-design in the form of mixed-precision FlashAttention. Our method accumulates key-query products and evaluates their exponentials in 8-bit formats, then adaptively identifies sensitive sub-blocks and recomputes them in 16-bit formats. We propose the specifications for a dedicated accelerator capable of executing this pipeline efficiently. Simulated experiments with Qwen3 and Gemma 3 show that rerouting a selective minority of sub-blocks to high precision is sufficient to recover the baseline model performance.

Stanislav Budzinskiy, Marian Gloser, Tolunay Yilmaz +5

Faculty of Mathematics University of Vienna, Austria · Huawei Heisenberg Research Center, Munich, Germany · Huawei Technologies Co. Ltd

AI Accelerator Inference LLM Inference Acceleration Efficient Attention Attention Mechanisms Neural Network Hardware Acceleration

Ask PDF

May 8, 2026

An Efficient Hybrid Sparse Attention with CPU-GPU Parallelism for Long-Context Inference

Long-context inference increasingly operates over CPU-resident KV caches, either because decoding-time KV states exceed GPU memory capacity or because disaggregated prefill-decode systems place KV data in host memory. Although block-sparse attention reduces attention cost in this setting, sparsity alone is insufficient for end-to-end efficiency. GPU-only designs remain constrained by PCIe bandwidth and metadata memory overhead, while CPU-GPU hybrid designs still suffer from substantial GPU idle time and bottlenecks in CPU-side top-k selection and sparse attention computation. Fluxion is built on three key insights: output-aware KV budgeting, head-specific and granularity-aware sparse configuration, and cross-device coordinated execution for sparse attention over CPU-resident KV caches. Guided by these insights, Fluxion combines a lightweight head-property predictor, a granularity-budget selector, and a priority-based scheduler to jointly optimize budget allocation, sparse configuration, and CPU-GPU execution overlap. This co-design enables hybrid sparse attention to achieve both accuracy and system efficiency in long-context inference. Across 2 models, 3 benchmarks, and 40 tasks, Fluxion preserves quality well -- the worst average degradation is only -0.26 relative to FULL, while delivering 1.5

\times

-3.7

\times

speedup over the strongest fixed sparse hybrid baseline, whose KV budget is only 0.05.

Feiyu Yao, Zhixiong Niu, Xiaqing Li +3

Beijing University of Technology Beijing, China · Microsoft Research Beijing, China · Beijing Jiaotong University Beijing, China

GPU Acceleration Self-Attention Long-Context Language Model Inference KV-Cache Management LLM Inference Acceleration Heterogeneous Computing+2

Ask PDF

Abstract

Context Parallelism

Explore similar work

Block Sparse Flash Attention

LampAttention: Look-Ahead Mixed-Precision FlashAttention for Dedicated Accelerators

An Efficient Hybrid Sparse Attention with CPU-GPU Parallelism for Long-Context Inference