Research map / Architectures and efficiency / Attention and transformers
FlashMask: Efficient and Rich Mask Extension of FlashAttention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FlashMask, a column-wise sparse (four interval vectors) attention mask representation integrated into FlashAttention-2 kernels, giving O(N) mask memory and skipping fully masked blocks, with end-to-end and kernel speedups over FlashAttention dense masks and FlexAttention.
Problem
Vanilla and dense-mask attention require O(N^2) memory for masks, and FlashAttention natively supports only a few fixed mask patterns, while complex masks used in SFT, LoRA, DPO, RM and long-context training need efficient support.
Approach
Represent each column's masked row intervals with four 1D vectors (LTS, LTE, UTS, UTE) and precompute per-block min/max values so that tiling-based FlashAttention-2 kernels can classify blocks as fully masked, partially masked, or unmasked and skip the fully masked ones, with bit-level numerical equivalence to dense-mask FlashAttention.
Key results
- End-to-end speedups of 1.65x to 3.22x over FlashAttention dense mask across Llama-2 7B/13B/70B tasks.
- Kernel TFLOPs/s exceeds FlexAttention by 12.1% to 60.7%, reaching 37.8% to 62.3% of A100 theoretical maximum FLOPs/s.
- Llama-2 7B LoRA supports sequence lengths up to 544K versus 64K for dense-mask methods.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace FlashAttention dense masks with FlashMask column-wise sparse masks | Llama-2 fine-tuning and alignment training · SFT, LoRA, DPO, and RM | throughput: 1.65x to 3.22x | Helped |
| Use FlashMask instead of FlexAttention | Attention kernels · 12 mask cases | TFLOPs/s: 12.1% to 60.7% | Helped |
| Use FlashMask instead of FlashAttention dense masks | Llama 3.1 8B training · SFT, LoRA, DPO, and RM | loss convergence | No clear effect |
| Use FlashMask rather than dense mask methods | Llama-2 7B LoRA training · 544K sequence length | maximum supported sequence length: 544K vs 64K | Helped |
Surprising findings
- Mask memory footprint is negligible relative to activations, so peak memory still grows rapidly with sequence length despite linear mask storage.
Limitations
- Column-wise representation cannot express arbitrary masks with irregular masked regions within a column.
- Completely random masks are hard to represent and compute efficiently.
- Current implementation is tied to PaddlePaddle and not yet ported to other frameworks; no Hopper-specific optimization.
Experimental setup
Datasets: allenai/tulu-v2-sft-mixture HuggingFaceH4/ultrafeedback_binarized synthetic packed-document sequences
Benchmarks: end-to-end training throughput (SFT, LoRA, DPO, RM) kernel TFLOPs/s on 12 mask types loss convergence curves
Metrics: throughput TFLOPs/s latency (ms) peak memory training loss
Baselines: FlashAttention dense mask FlexAttention Vanilla attention xFormers
Builds on: FlashAttention-2 FlexAttention xFormers Sparse Causal Flash Attention
Closest papers by meaning
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR 2024 · similarity 0.90)
- MonarchAttention: Zero-Shot Conversion to Fast, Hardware-Aware Structured Attention (NeurIPS 2025 · similarity 0.89)
- FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores (ICLR 2024 · similarity 0.88)
- FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (NeurIPS 2024 · similarity 0.88)
- Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels (NeurIPS 2025 · similarity 0.87)
- XAttention: Block Sparse Attention with Antidiagonal Scoring (ICML 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.