# FlashMask: Efficient and Rich Mask Extension of FlashAttention

arXiv: https://arxiv.org/abs/2410.01359 · HTML: https://mlsanity.com/paper/2410.01359/

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 38 references that are in the map.
- First posted on arXiv 2 Oct 2024.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [DoRA: Weight-Decomposed Low-Rank Adaptation](https://mlsanity.com/paper/2402.09353/) (2024)
- [RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment](https://mlsanity.com/paper/2304.06767/) (2023)
- [Statistical Rejection Sampling Improves Preference Optimization](https://mlsanity.com/paper/2309.06657/) (2023)
- [LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models](https://mlsanity.com/paper/2309.12307/) (2023)
- [OPT-IML: Scaling Language Model Instruction Meta Learning through the Lens of Generalization](https://mlsanity.com/paper/2212.12017/) (2022)
- [Patch n' Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution](https://mlsanity.com/paper/2307.06304/) (2023)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [In-Context Learning with Long-Context Models: An In-Depth Exploration](https://mlsanity.com/paper/2405.00200/) (2024)
- [LiPO: Listwise Preference Optimization through Learning-to-Rank](https://mlsanity.com/paper/2402.01878/) (2024)
- [NACL: A General and Effective KV Cache Eviction Framework for LLMs at Inference Time](https://mlsanity.com/paper/2408.03675/) (2024)
- [Faster Causal Attention Over Large Sequences Through Sparse Flash Attention](https://mlsanity.com/paper/2306.01160/) (2023)

## Built on by (most cited first, 20 of 20)

- [MAGI-1: Autoregressive Video Generation at Scale](https://mlsanity.com/paper/2505.13211/) (2025)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [Inference-Time Hyper-Scaling with KV Cache Compression](https://mlsanity.com/paper/2506.05345/) (2025)
- [FlashBias: Fast Computation of Attention with Bias](https://mlsanity.com/paper/2505.12044/) (2025)
- [ERNIE 5.0 Technical Report](https://mlsanity.com/paper/2602.04705/) (2026)
- [DCP: Addressing Input Dynamism In Long-Context Training via Dynamic Context Parallelism](https://mlsanity.com/paper/2510.10620/) (2025)
- [Toward Native Multimodal Modeling: A Roadmap](https://mlsanity.com/paper/2605.25343/) (2026)
- [Tree Training: Accelerating Agentic LLMs Training via Shared Prefix Reuse](https://mlsanity.com/paper/2511.00413/) (2025)
- [Attention in Diffusion Model: A Survey](https://mlsanity.com/paper/2504.03738/) (2025)
- [FG-Attn: Leveraging Fine-Grained Sparse Attention in Video Diffusion Models](https://mlsanity.com/paper/2509.16518/) (2025)
- [MaxCode: A Max-Reward Reinforcement Learning Framework for Automated Code Optimization](https://mlsanity.com/paper/2601.05475/) (2026)
- [Unifying and Enhancing Graph Transformers via a Hierarchical Mask Framework](https://mlsanity.com/paper/2510.18825/) (2025)
- [CuBridge: An LLM-Based Framework for Understanding and Reconstructing High-Performance Attention Kernels](https://mlsanity.com/paper/2605.05023/) (2026)
- [A Scalable Distributed Framework for Multimodal GigaVoxel Image Registration](https://mlsanity.com/paper/2509.25044/) (2025)
- [SAFE-DiT: Semantics-Aware Fast-path Execution for High-Resolution Diffusion Transformers](https://mlsanity.com/paper/2606.29360/) (2026)
- [Accelerating Sparse Transformer Inference on GPU](https://mlsanity.com/paper/2506.06095/) (2025)
- [Graphical einops: bridging tensor networks and computation graphs](https://mlsanity.com/paper/2605.31485/) (2026)
- [SPADE: An Input-Adaptive Sparse Attention Engine for Fast Video Diffusion Models Inference](https://mlsanity.com/paper/2608.03335/) (2026)
- [Competition Domains Matter: Selective Head Globalization for Training-Free Sparse Attention in Video Diffusion Transformers](https://mlsanity.com/paper/2610.14672/) (2026)
- [FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention](https://mlsanity.com/paper/2610.17901/) (2026)

## Abstract (arXiv)

The computational and memory demands of vanilla attention scale quadratically with the sequence length N, posing significant challenges for processing long sequences in Transformer models. FlashAttention alleviates these challenges by eliminating the O(N2) memory dependency and reducing attention latency through IO-aware memory optimizations. However, its native support for certain attention mask types is limited, and it does not inherently accommodate more complex masking requirements. Previous approaches resort to using dense masks with O(N2) memory complexity, leading to inefficiencies. In this paper, we propose FlashMask, an extension of FlashAttention that introduces a column-wise sparse representation of attention masks. This approach efficiently represents a wide range of mask types and facilitates the development of optimized kernel implementations. By adopting this novel representation, FlashMask achieves linear memory complexity O(N), suitable for modeling long-context sequences. Moreover, this representation enables kernel optimizations that eliminate unnecessary computations by leveraging sparsity in the attention mask, without sacrificing computational accuracy, resulting in higher computational efficiency. We evaluate FlashMask's performance in fine-tuning and alignment training of LLMs such as SFT, LoRA, DPO, and RM. FlashMask achieves significant throughput improvements, with end-to-end speedups ranging from 1.65x to 3.22x compared to existing FlashAttention dense method. Additionally, our kernel-level comparisons demonstrate that FlashMask surpasses the latest counterpart, FlexAttention, by 12.1% to 60.7% in terms of kernel TFLOPs/s, achieving 37.8% to 62.3% of the theoretical maximum FLOPs/s on the A100 GPU. The code is open-sourced on PaddlePaddle and integrated into PaddleNLP, supporting models with over 100 billion parameters for contexts up to 128K tokens.