FlashMask: Efficient and Rich Mask Extension of FlashAttention
Authors: Guoxia Wang
Jinle Zeng
Xiyuan Xiao
Siming Wu
Jiabin Yang
Lujing Zheng
Zeyu Chen
Jiang Bian
Dianhai Yu
Haifeng Wang
Venue: ICLR 2025 Poster (2025)
Keywords: Attention Mask Efficient Representation
Efficient Attention Computation
Long context
IO complexity
GPUs
LLMs
Primary area: infrastructure, software libraries, hardware, systems, etc.

Abstract: 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 $\mathcal{O}(N^2)$ 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 $\mathcal{O}(N^2)$ memory complexity, leading to inefficiencies. In this paper, we propose \ours{}, 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, \ours{} achieves linear memory complexity $\mathcal{O}(N)$, making it 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 \ours{}'s performance in fine-tuning and alignment training of LLMs such as SFT, LoRA, DPO, and RM. \ours{} 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 \ours{} 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\footnote{\url{https://github.com/PaddlePaddle/Paddle}} and integrated into PaddleNLP\footnote{\url{https://github.com/PaddlePaddle/PaddleNLP}}, supporting models with over 100 billion parameters for contexts extending up to 128K tokens.

Published as a conference paper at ICLR 2025
FLASHMASK: EFFICIENT AND RICH MASK EXTENSION OF FLASHATTENTION
Guoxia Wang. Jinle Zeng. Xiyuan Xiao, Siming Wu, Jiabin Yang, Lujing Zheng. Zeyu Chen, Jiang Bian, Dianhai Yu! Haifeng Wang
Baidu Inc.
(wangguoxia, zengjinle, yudianhai] @baidu. com
ABSTRACT
The computational and memory demands vanilla attention scale quadratically with the sequence length N, posing significant challenges for processing long sequences Transformer models. FlashAttention alleviates these challenges by eliminating the O memory dependency and reducing attention latency through 10-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(N²) memory complexity, leading to inefficiencies. In this paper, propose FLASHMASK, an extension of FlashAttention that introduces column-wise sparse representation of attention masks. This approach efficiently represents wide range of mask types and facilitates the development of optimized kernel implementa- tions. By adopting this novel representation, FLASIMASS achieves linear memory complexity O(N), making it suitable for modeling long contect sequences. More- over, 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 im- provements, 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% 1% to 60.7% in terms of kernel TFLOP//s, achieving 37.8% to 62.3% of the theoretical maximum FLOPs// on the A100 GPU. The code is open-sourced PaddlePaddld and integrated into PaddleNLE supporting models with over 100 billion parameters for contexts extending up to 128K tokens.
INTRODUCTION
The Transformer architecture Vaswani et al. (2017) has become foundational model in wide range of tasks across natural language processing (NLP), computer vision (CV), and multimodal applications. Central to its effectiveness is the attention mechanism, which enables the model to focus on relevant parts of the input data. the vanilla attention mechanism, the attention weights computed as scaled dot-product between query and key vectors, shown in Equation (1). To implement complex logic, the mask M can be added to the QK" term before applying the softmax function, controlling token visibility by setting certain elements -c0
'Equal contribution.
Comresooning author.

---

Published as a conference paper at ICLR 2025
QKT Vdx
Attention( K, V, .M = Softmax
+ M
In NLP tasks, large language model (LLM) training can generally be categorized into two main stages: PreTraining and PostTraining. During PreTraining, different masking strategies are employed to guide the model's learning process. For instance, GPT-style Radford (2018) models use unidirectional causal masking, as illustrated in Figure (a)(1,, while [ et al. (2020) models utilize a combination of unidirectional and bidirectional masking, as shown in Figure (a)().. PostTraining, which typically includes Supervised Fine-Tuning (SFT) [SFT)Jyer& et al. 2022) : Chung et al. (2024); Hu
al. (202T); Liu al. (2024a); Chen et al.] (2023), Direct Preference Optimization (DPO) Rafailov et al. 2024)) Dong et al. (2023): Liu et al. (2023); Ethayarajh et al. (0024); Liu et al. (2024b),
and Reward Model (RM) training within Reinforcement Learning from Human Feedback (RLHF)
Schulman et al. (2017); Shao et al. (2024), also employs a variety of masking techniques depending
on the specific task, as depicted in Figure (3) and (5).
The increasing complexity and diversity of masking strategies pose significant challenges to the efficiency of the attention mechanism. Specifically, the vanilla attention mechanism suffers from a quadratic increase in computational and memory demands, denoted as O(2²), where N represents the sequence length. As models scale to longer sequences, ranging from 128K to 1M tokens in advanced systems like (Achiam et al. (2023), Claude Anthropic (2024), and Gemini Reid et al. (2024), these quadratic dependencies become prohibitive, necessitating more efficient computational approaches. The memory load for masked attention computations also grows quadratically, further exacerbating the challenge of managing various mask configurations across different tasks.
Recent efforts such as Memory Efficient Attention (MEA) & Staats (2021 and
Dao et al. (2022): Dao (2023) have made strides in addressing these issucs by reducing memory over- head and attention latency. Flash Attentiin,,,,n, in particular, eliminates the O memory dependency and reduces attention latency through 1O-aware memory read/write optimizations. However, while FlashAttention natively supports certain mask types without additional memory overhead, its support for more complex masking requirements remains limited.
In this paper, we introduce FLASHMASK, extension of FlashAttention that leverages novel column- wise representation of attention masks. This approach allows for the efficient handling of broader range of mask types without compromising computational accuracy. FLASHMASK achieves linear memory complexity while enabling kernel optimizations that reduce unnecessary computations, resulting in significant computational speedups and enhanced training efficiency.
Our contributions are threefold:
1. We introduce a novel column- wise sparse mask representation that efficiently accommodates a broader range of mask types, enabling more flexible attention mechanisms.
2. We extend Fltetintention's 's masking capabilities by integrating optimized kernel implemen- tations, ensuring high computational efficiency without sacrificing computational accuracy.
3. We demonstrate the effectiveness of FLASHMASK across various attention mask types and models, underscoring its versatility and robustness in large-scale LLM training. FLASH- MASK notably reduces both computational and memory overheads, significantly enhancing its suitability for long-context modeling.
2 BACKGROUND
2.1 ATTENTION MASK TYPES
Transformer-basee models have demonstrated exceptional versatility across a variety of tasks, each benefiting from different attention mask types, as shown in Figure Causal Mask is predominantly used in autoregressive models to predict the next token in a sequence, ensuring that each token only attends to previous tokens and avoids information leakage from future tokens Vaswani et al. (2017). Sliding Window Mask captures local context by allowing tokens to attend to a fixed-size window of neighboring tokens, balancing computational efficiency with the ability to capture local dependencies Beltagy et al. (2020). Causal Document Mask, employed in methods like efficient sequence packing and in-batch/in-tokens techniques, accelerates large language models without
2

---

Published as a conference paper at ICLR 2025
Plac Gamparee Masi.
(a)
(b)
Figure 1: Overview of FLASHMASK. (a) Types of Masks Supported by FLASHMASK, (b) Column- Wise Sparse Representation in FLASHMASK, (c) Efficient Implementation of FLASHMASK.
performance degradation by ensuring tokens attend only to previous tokens within the same document Krell et al. 2 [yer et al. (2022): Dubey et al. (2024). Document Mask, or bi-directional attention, permits tokens attend to all other tokens within the same document, facilitating context learning from both directions and is widely used models like BERT and vision transformers like NaViT
Devlin et al. 2018): Dehghani et al. (2024).
Shared Question Mask is utilized Reward Models (RM) and Direct Preference Optimization (DPO) models, allowing multiple answers share single question, thus eliminating redundant computations and speeding up training Ouyang et al. (2022). The Global + Sliding Window Mask combines global attention with sliding window attention, where global tokens attend all tokens while others use a sliding window mask, effectively handling tasks requiring both global context and local details Zaheer al. 2020).
Causal Wlie Mask, primarily used in in-context learning, divides sequences into blocks, where demonstrations only attend nearby examples within small blocks, while the test example can attend to all demonstrations, allowing the study of model performance improvements in long-context tasks Bertsch al. (022)). Prefix LM Causal Mask is tailored for language modeling tasks, allowing prefix to attend to all tokens to generate coherent text based on the prefix Raffel et al. (2020) Prefix Document Mask extends this concept to multiple documents, where a prefix each document attends to all tokens within that document but not across documents.
OK- OK-Sparse Mask optimizes self-attention by sparsifying query-key pairs, reducing computational load while maintaining performance, which is particularly beneficial for large-scale models Kitaev et al. (000)) Hash-Sparse Mask employs Decality-sensitive hashing to partition sequences into smaller chunks, enabling efficient sparse attention for long sequences Kitaev al. (2020). Lastly, Random Eviction Mask introduces randomness by randomly masking out tokens during training, aiding in generalization and simulating Key-Value Value (KV) cache eviction processes to handle long sequences without memory overflow Chen et al. (2024).
2.2 ATTENTION MASK SUPPORTED
Attention mechanisms fundamental to transformer-basee models, with various mask types en- abling different attention patterns. The vanilla attention mechanism, as shown in Equation 2 supports arbitrary mask types through a dense mask matrix:
S= OKT VdTe
ERNNN
P Softmax (S + M) E NXN, O=
(2)
where Q. K, V E Nxd are input sequences, M E NNN is the attention mask, N is the sequence length, and is the head dimension. The mask M modulates token visibility through element-wise addition with S. While this approach supports arbitrary mask types, it incurs a memory complexity of O(N2), limiting its scalability for long sequences.
3

---

Published as a conference paper at ICLR 2025
addresses this limitation through I0 -aware read/write operations and tiling techniques, eliminating the need for the intermediate E ERR NxN and explicit mask M. However, FlashAttention only supports predetermined mask patterns within its kernel, such as causal, sliding window, causal document, and document masks, as shown in Figure
FlashAttention Dao et al. (2022) Dao (203))
xFormers Lefaudeux al. (2022) extends FlashAttention's 's capabilities, offering support for masks with diagonal offsets. It represents document masks using cumulative sequence lengths, achieving a memory complexity of ONN).
Flex Attention He et al. (202) introduces a more flexible mask description method based on deep learning compiler techniques. By combining block masks with expression-based descriptions, it can support arbitrary mask types. While this approach significantly reduces memory overhead through block-based processing, its memory complexity remains O N2
B.
Our proposed method, FLASHMASK, extends FlashAttention's 's mask support capabilities. It intro- duces a flexible, column- -wise sparse mask representation that covers the majority of mainstream Transformer modeling requirements. As illustrated in Figure Tb), FLASHMASK expresses which intervals need to be masked on a per-column basis, achieving a memory complexity of O(N). This approach bridges the gap between mask flexibility and computational efficiency, offering a more versatile solution for attention mechanisms in large-scale transformer models.
2.3 ATTENTION OPTIMIZATION TECHNIQUES
The attention mechanism, as formulated in Equation presents significant computational and memory challenges, particularly in the computation of QK" As the sequence length N increases, the resultant attention scores matrix grows quadratically, leading to complexity of O To address this scalability issue, researchers have proposed various optimization techniques, focusing on both memory efficiency and computational speed.
Memory Efficient Attention (MEA) Rabe & Staats (021) marks a notable advancement in model training optimizations. By leveraging Online Softmax Milakov & Gimelshein (2018) and chunking techniques, MEA reduces memory requirements from O(N²) (VV),, enabling the use of larger models extended sequence lengths within existing hardware constraints. Building upon this foundation, FlashAttention Dao et al.] (2022); Dao (2023) focuses on reducing attention latency through 10-aware memory read/write optimizations. Utiizing tiling techniques during computation, FlashAttention achieves memory overhead of ONN), proving particularly effective tasks without custom masking requirements. Furthermore, FlashAttention extends to Block-Sparse FlashAttention, introducing a two-dimensional block mask matrix representation to indicate masked tiling blocks. This innovation allows for the skipping of computations for masked blocks, thereby accelerating the process.
For scenarios requiring specific attention masks, several tailored solutions have emerged. Sparse Causal Flash Attention (SCFA) Pagliardini et al. (2023) extends FlashAttention to optimize QK- Sparse and Hash-Sparse scenarios in causal attention structures. SCFA employs indices of queries and keys in the original uncompressed tensors to describe masks, enabling the omission of computations for masked blocks and enhancing computational efficiency. FlexAttention He et al. (2024) leverages compiler techniques to simplify mask attention implementations, exploiting sparsity in the attention mask to skip certain masked blocks and achieve improved speed. However, there remains room for optimization, particularly for complex masking patterns.
Our proposed method, FLASHMASK, builds upon these advancements to support customized complex attention masks. FLASHMASK reduces memory complexity from O (N2) O while leveraging sparsity in the attention mask to skip masked blocks. Through rigorous engineering optimizations, FLASHMASK achieves superior computational speed compared Attentinn, particularly in tasks with complex masking requirements. By synthesizing the strengths of existing approaches with novel optimization techniques, FLASHMASK represents significant advancement in attention mechanism efficiency, addressing both memory constraints and computational speed in large-scale transformers.
4

---

Published as a conference paper at ICLR 2025
3 OBSERVATION
In the current paradigm of training Transformer-based models, attention mechanisms can be broadly categorized based on their causality, as introduced in Section 2.1 and illustrated in Figure I(a). These representations encompass the majority of mask types encountered in training scenarios. Consider the attention score matrix S, where each element Sy represents the attention of the i-th query token to the j-th key token. From the perspective of the key tokens, we observe a critical pattern in how query tokens attend to each key token.
Our key observation is that the inability of query tokens to attend to certain key tokens exhibits a continuous nature. This continuity allows us transform the two-dimensional dense mask M into more compact, one-dimensional representation using row index intervals, as depicted in Figure (). Formally, can express this transformation as:
Mj [startj, endj). Vj E (I.....N)
where M, represents the interval of row indices that are masked for the j-th key token, and N is the sequence length.
While this column-wiss one one-dimessional interval representation may not capture arbitrary mask patterns, it effectively covers the predominant mask types encountered in practice. Moreover, this representation offers a significant advantage: it facilitates a straightforward conversion to masked blocks in tiling-based computations. This conversion enables the elimination of unnecessary calculations, thereby enhancing the computational efficiency of the attention kernel.
This concept of interval representation can be generalized to other forms of continuous intervals. For instance, by transposing the attention matrix, we obtain a row-wise representation using column index intervals. The flexibility of this approach allows for efficient handling of various attention patterns while maintaining compact representation that is conducive optimized computation.
4 FLASHMASK: ALGORITHM AND ANALYSIS
In this section, we introduce the novel column-wise mask representation of FLASHMASK, extend FlashAttention to support complex mask patterns, and provide comprehensive complexity analysis of our approach.
4.1 COLUMN- WISE -WISE
MASK REPRESENTATION
To efficiently handle complex mask patterns in both causal and bidirectional attention scenarios, propose a novel column-wise sparse representation for FLASHMASK. The attention score matrix is partitioned into lower-left and upper-right triangular sections relative to the diagonal. FLASHIMASK expresses the mask using four one mee-imensional vectors:
LTS:: Lower Triangular Start the starting row of the mask in the lower-left triangle.
LTT:: Lower Triangular End the ending row of the mask in the lower-left triangle.
UTS:: Upper Triangular Start the starting row of the mask in the upper-right triangle.
UUTE: Upper Triangular End the ending row of the mask in the upper-right triangle.
The indices of rows to be masked in the lower triangular section are given by [LTS, LTE), and in the upper triangular section by [UTS, UTE). Specifically, each column described by two mask intervals. For the j-th token, tokens within the intervals [LTS). LTE,) U [UTS). UTE,) cannot attend to it. For example, as illustrated in Figure ((b) for the fifth column, [LTS5, LTEs) (UTS5, UTEs) [7, 10) u 4) indicates that rows 2 to 4 and 7 to 9 are masked.
This representation offers several advantages:
1. Compactness: It reduces dense 2D mask to more efficient 1D representation.
2. Flexibility: can capture wide range of practical mask patterns, including causal, bidirec- tional, and more complex attention mechanisms.
3. Computational Efficiency: It facilitates easy conversion to masked blocks in tiling- based computations, enabling the elimination of unnecessary calculations.
5

---

Published as a conference paper at ICLR 2025
4.2 EXTENDING FLASHATTENTION FOR COMPLEX MASKS
We integrate the column- wise mask representation of FLASHMASK into the FlashAttention-2 al- gorithm, extending its mask support capabilities. The high-performance kernel implementation of FLASIMASK consists of two key steps:
Preprocessing: Given the input column-wist sparse mask vectors, we first partition them into Te blocks along the column dimension in banddbhn memory (HBM). For each mask vector, we compute the maximum and minimum values within each block. This results in 8 intermediate vectors: maar Startminn End LTEnd Starrtt UT Startmin UT End and UT UTEndd each of size Te.
Real-time Block Skip Computation: Using these min-max vectors, we can classify each tiling block of attention score matrix into three categories during kernel computation. The block mask type Thlock determined as follows:
if BlockRow min 2 Starrmax and RlockRowwie max S End Partially masked, elif RlockRowmin End'max and Block Row max Start (4) otherwise
Fully masked,
Thlock
Unmasked,
This classification allows us to skip fully masked blocks, reduce computation for unmasked blocks, and apply element-wise masking only for partially masked blocks. Figure T illustrates the entire kernel computation process for a causal scenario with LTS and LTE in the lower-left triangle. Algorithm details the forward computation process FLASHMASK extended from FlashAttention- 2, with blue-shaded parts indicating FL computations.
For the backward pass, FLASHMASK'S column-sparse representation particularly advantageous. The computations of dK and dV are column- column-parallel, allowing efficient loading of maximum and minimum values into registers for extensive data reuse during block computations, as shown in Algorithm 2 in the Appendix.
4.3 COMPLEXITY ANALYSIS
We define block sparsity in attention mask as p
where By, Be are block sizes, and is Bc
B,
the number of completely masked blocks.
Space Complexity:
The dense mask requires O(N2) space, while FLASHMASK uses O(N) space
for LTS, LTE, UTS, UTE e RN and 8 precomputed min-max vectors R de This significant reduction in memory usage enables training on longer sequences.
Memory Access Complexity: The dense mask requires O memory accesses on HBM. FLASH-
MASK reads the LTS, LTE, UTS, UTE RN vectors from HBM as shown in lines 16 and 19 of Algorithm with each Q reading the entire LTS, LTE, UTS, UTE, totaling 4 x T, x N memory N² &.
accesses. This reduces the memory access approximately 4x7,xX XN 4 significantly boosting performance. Furthermore, FLASSMASK'S compact representation allows preloading of mask vectors into SRAM, further enbancing memory access efficiency.
While the standard attention computation has a complexity of O FLASIMASK leverages sparsity in the attention mask to reduce it to O((1 - -p)TTT)) by skipping entirely masked blocks.
Computational Complexity:
These improvements in space, memory access, and computational complexities contribute to FLASH- MASK's superior performance and efficiency handling complex attention patterns.
4.4 CORRECTNESS ANALYSIS
As shown in Equation? 2 the computation of the attention matrix P Softmax SS + M) involves augmenting the attention scores with mask M, where the masked elements are set to -00, This operation ensures that the softmax outputs at these masked positions are zero, effectively omitting them from attention. Consequently, if an entire block is fully masked, the resulting output for that block will be all zeros. FLASHMASK exploits sparsity in the attention mask by skipping computations involving these entirely masked blocks, thus reducing computational overhead without altering the

---

Published as a conference paper at ICLR 2025
Algorithm 1 FlashAttention-22 Forward Pass Extended with FLASIIMASK
Require: Matrices K. E gxx in HBM, block sizes Be, By. vectors LTS, LTE, UTS, UTE E RN Divide Q into T, =
1:
B,
N
blocks QI...., Qr, of size By x d each, and divide K, V to Tc of size Be x d each.
B,
N
blocks
Kj....,K Kr, and V1 I...... VI,,
2: Divide the output O E 2 into T, blocks OI...., Or, of size B, x d each, and divide the logsumexp into T, blocks Li...., Lr, of size B, each.
3:
Divide LTS, LTE. UTS, UTE into Tę blocks LTS .IS... LTE ,TE,, UTS UTS UTE,..... UTE of size B, each.
4:
Precompute the min and max row index LT min(LTS)). LTStaat max(LTS)).
LTEnd min(LTEj), LTEndjax = max(LTE)), UTStami) = min(UTS;), UT max(UUT))) Und min(UTR;). Ennduu max(UTE)). Vje .l.....Tech write HBM.
5: for iis T, do
6: Load Q; from HBM to on-chip SRAM.
7: On chip, initialize ; B,) E BB.ddd ,((0) (0)B, e ER .m;
(0) = (-00) E ERR
8:
for 1 S Te do
9:
if (i - 1) x B, 2 LT and ix B, s LTEnd then
10:
Continue // lower triangular skip calculation of masked block
11:
end if
12:
if (i 1) 2 Star and i B, UTEndflin then
13:
Continue // upper triangular skip calculation of masked block end if
14:
Load Kj. y, from HBM to on-chip SRAM.
16:
On chip, compute 5 K.K. c BB,.B
17:
if x By > LTStaar and (i - 1) x B, < LLTnndaas then Load LTS, and LTE, from HBM to on-chip SRAM.
18:
19:
On chip, apply mask: 5 [x][y]
y, such that LTS,DI ((-1)) (i-l)xB, LTE))
20:
end if
21:
B, > UTarta and (i 1) x B, < UT UTEndaa then Load UTS, and UTE, from HBM to on-chip SRAM.
23:
chip, apply mask: 5//) [x][x] = -oo Vx. such that UTS & s (i end if
24:
On chip, compute m; (1) m (j-1) rowmax E B, p( cxp(s -m mm
25:
R BB (pointwise), e(0 m! -m(" ((-1) rrowsum (()) E
On chip, compute 0 diag(e" mm 1-- 0 -!!-)) p(DV Vj. end for
26:
27:
28: On chip, compute O; (Te)
29:
On chip, compute I, = m (Tee) + log( (Te)
30:
Write O; HBM the ith block of O.
31: Write I, to HBM as the i-th block L.
32: end for
33: Return the output O and the logsumexp L.
outcome. Importantly, FLASHMASK maintains bit-level numerical equivalence with the computations performed using dense mask in FlashAttention, ensuring that there is no loss in precision. This exactness corroborated in our experimental evaluations, where we verify that the loss convergence curves from end-to-end training align precisely at the bit level (see Section 5.2
5 EXPERIMENTS
In this section, we evaluate the performance of FLASHMASK through series of experiments designed to demonstrate its end-to-end acceleration and memory efficiency under different model scales and sequence lengths, training convergence practical scenarios, its relationship with block sparsity in the attention mask, and its effectiveness across various attention mask patterns. All experiments were conducted on machines equipped with NVIDIA A100-SXM S0G GPUs, Intel(R) Xcon(R)
7

---

Published as a conference paper at ICLR 2025
Varilb Anamiin
Tladhbok
LoRA
LLaMA- 22-700
Rpodity
Syeed
Ldwe
Suquanee (emph
Suqaies Lengla (K)
Scceener Langhh
Supuso Leggh
MaMA--21338
2004
16
124 296
16
128
124
30 Segensio Leaght
1:
22 Sepuonel Lenglh
121
SegumoeLeegg.000
Segaese Lengln
794
LJaMA-2-7B8
Sood Us
Speed Up
128 25d
16
12
126
256
54c
10
Figure 2: End-to-end training throughput was assessed across varying sequence lengths for different Llama2 model scales in four downstream training tasks: SFT, LoRA, DPO, and RM.
Platinum 8350C CPUs, CUDA 12.0, and driver version $25.125.06. Due to space limitations, detailed information about datasets and hyperparameters settings is provided in the appendix
5.1 END-TO-END TRAINING THROUGHPUT
To showcase the practical effectiveness of FLASHMASK, we evaluated the end-to-end training throughput Llama-2 models of three scales (7B, 13B, and 70B) across four downstream tasks involving the fine-tuning and alignment training (SFT, LoRA, DPO, and RM) with varying sequence lengths. We compared FLASHMASK with two dense mask methods. The experimental results presented in Figure? 2 The results lead to two key conclusion. Higher Throughput: FLASIIMASK achieves higher throughput compared to dense mask methods with quadratic memory complexity. Specifically, FLASHMASK attains a 1.65x to 3.22x improvement over the maximum sequence length supported by FlashAttention dense mask. This demonstrates that, in practical applications, FLASH- MASK can significantly enhance training throughput of large language models, thereby reducing training costs. Linear Memory Overhead: FL.SHM to support longer sequence lengths. In the Ilama-2 7B LoRA training, FLASHMASK supports sequence lengths up to 544K, whereas other methods are limited to 64K. At a sequence length 64K, the memory overhead for dense mask methods amounts to SGB, Figure (b) depicts the memory overhead curve, highlighting the efficiency of FLASHMASK in terms of memory consumption.
linear memory overhead enables
5.2 END-TO-END TRAINING CONVERGENCE VERIFICATION
The core innovation of FLASHMASK lies in introducing column-wise sparse mask representation, which leverages sparsity in the attention mask to skip computations on fully masked blocks, thereby enhancing speed without altering the algorithm's precision. To verify that FLASHMASK does not compromise convergence accuracy, we conducted end-to-end training experiments on the Llama 3.1 Dubey et (2024) 8B model across four fine-tuning and alignment training tasks of LLMs.
is important to note that the backward computation of dQ in the CUDA kemel implementation may introduce randomness due the accumulation order (see line 27 of Algorithm 2). Therefore.
8

---

Published as a conference paper at ICLR 2025
SFT Trin with Deteminntic OIT
LokA Thin Line Detemiinidi Off
DPO Thin with Detemmmmtic Off
Tnin Deteminsticc Off
FadAtarden Duss Fabkiak
1.2
Fabltaak
1.0
1008
0.6
0.24
0.4
0.3
0.4
0.8-
12500
255
2009
1000
SET Zrais Losa with Determiiiitti On
LoRA Train Losa with Detemmnnitto On
D20 Trais Losa with Detamminiati On
RM Trais Losa with Determiiiatioo On
Faberamiod
Deac
1.8
1.00
0.8-
0.3
1.00
0.1
0,04
0:20
9.4
0.8
0.8
2100
1:00
19/0
2000
17:00
2.50
1200
19:00
12.
17400
10.00
1200
2009
1.00
Skrp
Stip
Skyg
Figure 3: End-to-end training loss.
Canoal Decemant Madk
Docrmment Mask
Stare Queetion Mark
Memoyy Joage (G8) us Siqgunton Lengt)
800
Figure 4: (a) Kernel execution latency at different sparsity levels, (b) Memory usage, Y-axis uses a base-10 10 logarithmic scale.
we performed convergence experiments under conditions with and without deterministic control.
The loss curves are shown in Figure 3 When deterministic control is enabled, the loss curves of FLASHMASK and FlashAttention dense mask align precisely, demonstrating identical numerical behavior. When deterministic control is disabled, both methods exhibit the same loss convergence trends. These results conclusively prove that FLASHMASK is an exact algorithm that preserves convergence accuracy.
5.3 SPARSITY-RELTEED EXPERIMENTS
FLASHMASK leverages block sparsity in the attention mask to skip computations on fully masked blocks, resulting in computational complexity proportional to O((1 p)T,T,). To verify this relationship, we performed experiments on three different mask cases under the configuration of 6 data type, sequence length of 32K, head dimension of 128, and 32 heads. We sampled data with varying sparsity levels for testing. Figure (a) illustrates the kernel execution latency at different sparsity levels, demonstrating a linear relationship between latency and sparsity.
5.4 KERNEL PERFORMANCE COMPARISON
To thoroughly evaluate the expressiveness and computational efficiency of FLASHMASK under common attention mask patterns, we conducted kernel-level comparisons with FlexAttention. The experiments were carried out across 12 different mask cases, with sequence lengths of 8K, 32K, and 128K, and head dimensions of 64 and 128, using 16 data type. The total number of tokens was fixed at 128K; varying sequence lengths yielded corresponding batch sizes, and fixed hidden size of 4096 allowed us to adjust the number of heads by changing the head dimension. Both FLASHMASK and FlexAttention exploit sparsity in the attention mask. We measured kernel speed using the TFLOPs/s metric. As shown in Figure FLASHMASK outperforms FlexAttention in terms of total TFLOPs/s for both forward and backward passes across all cases, with improvements ranging

---

Published as a conference paper at ICLR 2025
Segunso kegth SK,
Sugunee Sough
Sequence lengh
P41
1G
1414
Sakin Midde
203
12.5
Diczzet Nark
114
Qentton
140
Deta LM Diczzen Nack
1033
1-24
(nule We
(400 012
Phita
@foppar Nark Harden bruckee
U14
200 PH
36 106
Spund EIPOOMO
Spead (TILOre)))
Figure 5: Kernel forward and backward speed (head dim 128, BF16) on A100-SXM 80G GPU. FlexAttention using PyTorch 2w.0.dew00244220cculllo from 12.1% % 60.7%. FLASHMASK achieves 37.8% to 62.3% of the theoretical maximum FLOPs/s on the A100 GPU.
LIMITATIONS AND FUTURE DIRECTIONS
While FLASHMASK significantly advances the efficiency of attention mechanisms for long sequences, it has limitations. The column-wise mask representation reduces memory complexity from O(N2) to ONN), offering substantial memory savings for long sequence training and effectively capturing the most common mask patterns. However, cannot represent arbitrary masks, particularly those with irregular masked regions within a single column. Extreme cases, such as completely random masks, pose challenges for both representation and efficient computation. Future research should focus on developing more sophisticated sparse representations that simultaneously maximize expressiveness and computational efficiency, particularly those amenable tiling techniques for figh-performance kernels. Extending FLASIMMASK to leverage features of newer architectures, such as NVIDIA's Hopper, could further enhance performance. Additionally, while our current implementation is based on the PaddlePaddle Ma et al. (2019) framework, integrating FLASHMASK into other popular deep learning frameworks could broaden its impact and accessibility. These efforts aim to address current limitations while expanding FL.SHHMSKK applicability across broader range of tasks, contributing the ongoing evolution of efficient transformer models for long sequence processing.
CONCLUSION
In this paper, we introduced FLASHMASK, an innovative extension of the FlashAttention algorithm that introduces a column-wise sparse mask representation to efficiently handle a wide spectrum of attention mask patterns in Transformer models. Our approach reduces the memory complexity from O(N²) to (N), enabling the processing of significantly longer sequences, which is crucial for modern large language models. By integrating this representation into the FlashAttention algorithm and implementing optimized kernels, FLASHMASK leverages sparsity in the attention mask to skip computations fully masked blocks without sacrificing computational accuracy. This strategic approach achieves notable computational speedups, with observed end-to-end enhancements ranging from 1.65x 3.22x during fine-tuning and alignment training of large language models, compared to the existing FlashAttention dense method. Furthermore, FLASHMASK significantly decreases the memory overhead associated with attention mask storage, thereby extending support for even longer sequence modeling. Additionally. FLASHMASK outperforms the latest counterpart, FlexAttention, by 12.1 1 % 60.7% kernel TFLOPs/s, achieving 37.8% to 62.3% of the theoretical maximum FL.PP//s on the A100 GPU. Our approach has been validated on downstream tasks of large language models, and we anticipate its widespread adoption in the industry.
ACKNOWLEDGMENTS
This work was supported by the Beijing Municipal Science and Technology Project (No. Z231 120100000032)002.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0877/2 2023.
Anthropic.
Introducing claude.
https::/www.anthhroop....cooo/e
Introducing-claudee 2024. Accessed: May 20, 2024.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/ 2020.
Amanda Bertsch, Maor Ivgi, Uri Alon, Jonathan Berant, Matthew R Gormley, and Graham Neu- big. In-context learning with long-context models: An in-depth exploration. arXiv preprini arXiv:2405.002200 2024.
Yilong Chen, Guoxia Wang, Junyuan Shang, Shiyao Cui, Zhenyu Zhang, Tingwen Liu, Shuo- huan Wang. Yu Sun, Dianhai Yu, and Hua Wu. NACL: A general and effective KV cache eviction framework for LLM at inference time. In Proceedings of the 62nd Annual Meet- ing of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 7913- 7926. Bangkok, Thailand, August 2024. Association for Computational Linguisties. URL ht /ittps://aclanteee www.tnnhhhnoy..o.ooooooooooooo.m a0t.2222.2..lllo 428
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309. 12307, 2023.
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetunel language models. Journal of Machine Learning Research, 25(70):1-53, 2024.
Tri Dao. Flashattention-22: Faster attention with better parallelism and work partitioning. arXiv 2023.
preprint
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16355 16359, 2022.
Mostafa Dehghani, Basil Mustafa, Josip Djolonga, Jonathan Heek, Matthias Minderer, Mathilde Caron, Andreas Steiner, Joan Puigcerver, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Patch pack: Navit, a vision transformer for any aspect ratio and resolution. Advances in Neural Information Processing Systems, 36, 2024.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 04805, 2018.
Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, and Tong Zhang. Raft: Reward ranked finetuning for generative foundation model alignment. arXiv preprint urXiv:2304.0677 2023.
Abhimanyu Dubey, Abhinay Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad AI-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21770 2024.
Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. Kto: Model alignment as prospect theoretic optimization. arXiv preprint arXiv:2402.01306, 2024.
Horace He, Driss Guessous, Yanbo Liang, and Joy Dong. FlexAttention: The Flexibility of PyTorch with the Performance FlashAttention. El Flexattentions August 2024. Accessed August 07, 2024.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint 2rXiv:2106.000 09685, 2021.
11

---

Published as a conference paper at ICLR 2025
Hamish Ivison, Yizhong Wang. Valentina Pyatkin, Nathan Lambert, Matthew Peters, Pradeep Dasigi, Joel Jang, David Wadden, Noah A. Smith, Iz Beltagy, and Hannanch Hajishirzi. Camels changing climate: Enhancing im adaptation with tulu 2, 2023.
Srinivasan Iyer, Xi Victoria Lin, Ramakanth Pasunuru, Todor Mihaylov, Daniel Simig, Ping Yu, Kurt Shuster, Tianlu Wang, Qing Liu, Pumit Singh Koura, al. Opt-iml: Scaling language model instruction meta learning through the lens of generalization. arXiv preprint arXiv:22 2022.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint #rXiv:2001.0445/ 2020.
Mario Michael Krell, Matej Kosec, Sergio P Perez, and Andrew Fitzgibbon. Efficient sequence packing without cross-comtamination: Accelerating large language models without impacting performance. arXiv preprint arXiv:2107.02027, 2021.
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov. xformers: A modular and hackable transformer modelling library.
http://hhh...coocccccoooooeeee wttp://sseearccooooooeeeeeee 2022.
Shih- Yang Liu, Chien-Yo Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang- Ting Cheng, and Min-Hung Chen. Dora: Weight-decomposed low-rank adaptation. arXiv preprint arXiv:2402.09335. 2024a.
Shih-Yang
Tianqi Liu, Yao Zhao, Rishabh Joshi, Misha Khalman, Mohammad Salch, Peter J Liu, and Jialu Liu. Statistical rejection sampling improves preference optimization. arXiv preprint arXiv:2309.0666 2023.
Tianqi Liu, Zhen Qin, Junru Wu, Jiaming Shen, Misha Khalman, Rishabh Joshi, Yao Zhao, Moham- mad Saleh, Simon Baumgartner, Jialu Liu, al. Lipo: Listwise preference optimization through learning-to toriinggtt-rank.... arXiv preprint arXiv:2402.01878 2024b.
Yanjun Ma, Dianhai Yu, Tian Wu, and Haifeng Wang. Paddlepaddle: An open-source deep learning platform from industrial practice. Frontiers of Data and Domputing, 1(1):105-1155 2019. 2019
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv: 1805::8805.0280 2018.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730- 27744, 2022.
M. Pagliardini, D. Paliotta, M. Jaggi, and F. Fleuret. Faster causal attention over large sequences through sparse flash attention. In Proceedings of the international conference on Neural Information Processing Systems (NeurIPS), pp. 198808598311 2023. URL htttps:// org/ Sapers/pacy http:pppppaaaaii--------eeeeeeee
.d pdf
Markus Rabe and Charles Staats. Self-attention does not need o (n2) memory. arXiv preprint arXiv:2112.05682. 2021.
Alec Radford. Improving language understanding by generative pre-training. 2018.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36, 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal machine learning research, 21(140): 166,, 2020.
12

---

Published as a conference paper at ICLR 2025
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, lean-baptiste Alayrac. Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530, 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, YK Li, Y Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300 2024.
Lewis Tunstall, Edward Beeching, Nathan Lambert, Nazneen Rajani, Kashif Rasul, Younes Belkada, Shengyi Huang, Leandro von Werra, Clémentine Fourrier, Nathan Habib, Nathan Sarrazin, Omar Sansevicro, Alexander M. Rush, and Thomas Wolf. Zephyr: Direct distillation of Im alignment, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Iliia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, et al. Flashinfer: Efficient and customizable attention engine for Ilm inference serving. arXiv preprint arXiv:2501.01000 2025.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17297, 2020.
A APPENDIX
A.I BACKWARD PASS ALGORITIM DETAILS
The detailed implementation of the FLASHMASK backward pass is presented in Algorithm 2 Similar to the forward pass, we precompute the maximum and minimum values of LTS, LTE, UTS, and UTE. These precomputed values LT Starthhr,,, LtarStarr L.7' L.T'Endrrin LT L.EEndpaxx UT Startii UTSrtrimax UT UTEndd UT Endfaax can be directly loaded into registers and kept resident because the backward computation operates in a column-paralla mode. Additionally, LTS,, LTE,, UTS,, and U can be loaded into SRAM outside of the inner loop (lines 10-11), thereby reducing the number of accesses to HBM 4 x N. Within the inner loop, the computation logic of FLASHMASK remains identical to that of the forward pass.
A.2 END-TO-END TRAINING THROUGHPUT
Recent models such Llama 3.1 and GPT-4, the Claude series, and Google's Gemini support sequence modeling beyond 128K tokens. FLASHMASK, with its reduced memory overhead, facilitates training with even longer contexts. However, existing public DPO and RM datasets lack training data for scenarios exceeding 128K tokens. To comprehensively evaluate FLASHMASK, we constructed synthetic data simulate long-sequence training and verify end-to-end throughput improvements. We validated our method across four downstream tasks involving the fine-tuning and alignment training of large language models: SFT, LoRA, DPO, and RM.
A.2.1 DATA CONSTRUCTION METHON
For end-to-end training, to realistically simulate real dataset distributions, we needed to distinguish between source tokens and target tokens within a document's sequence length. Additionally, the data construction method differed from that used in the kernel experiments. Given a maximum training sequence length N and document count range [1, 10], first randomly sampled the number of documents, then sampled each document's sequence length such that the total sequence length equaled N. The last document was considered padding. For the RM training, special constraints
13

---

Published as a conference paper at ICLR 2025
Algorithm 2 FlashAttention-2 Backward Pass Extended with FLASHIMSSK
Require: Matrices Q.K. V, E RNx in HBM, vector L E RN in HBM, block sizes Bes Br LTS, LTE. UTS, UTS. ERV.
1: Divide Q into
blocks Q1.-... Qr, of size B, x each, and divide K, V in to KI.....Ky and V...... Vr,- size B. each.
N
blacks
Be
2: Divide O into T, blocks O...... Or, of size B, x cach, divide dO into T, blocks d dO,, of size Br x d cach, and divide L into T, blocks Li.....Ly Lr, of size By cach.
3: Initialize dQ= (0) xd in HBM and divide it into T, blocks dQ...... .Q.., of size B, d cach. Divide
dV., dK, dV E R Nxd in to Te blocks j..... dKrr, and dV .VV.. of size Be x d cach.
4: Compute D = rowsum dO O) Rd (pointwise multiply), write D to HBM and divide it into T, blocks
DI...., Dr, of size B, each.
5: Divide LTS, L.TE, UTS, UTE into Te blocks LTS.,..... LTST. LTE,.... LTTT, ITS!..... UTST.
UTE UTET, of size Bg cach.
6: Precompute the min and max row index L.Tttarrhin min(LTSj), LT'Starruas max(LTSs)) LT I.TEnd minn( minT.. L.T L.TEnd max(L.TE,), UTStaart) min(UTS))) UT UT max maxx(USS)).... (UTSTU).TTTA"""""" UTEnd min(UTE)). UT UnEndaax max((UTE)). Vj (l....,Te). write HBM.
7:
for I S Tc do
8:
Load Ky. y from HBM to on-chip SRAM. Initialize dky (0) dV, B.) SRAM.
9:
10:
Load LTS, and LTE, from HBM to on-chip SRAM.
11:
Load UTS/ and UTE, from HBM on-chip SRAM.
12:
for s si T, do
13:
if (i- 1) z LT and x B, LTEndfjin then
14:
Continue // lower triangular skip calculation of masked block
15:
end if
16:
if (i- 1) B, 2 Utarn and i B, UTEnd then
17:
Continue // upper triangular skip calculation masked block end if
18:
19:
Load Q; O, dO, dQ,, Li D; from HBM to on-chip SRAM. On chip, compute s(") = Q,K7 E BB..B.
20:
21:
i x B, > LTStarri and (i - 1) x B, LTEndmaax then
22:
On chip, apply mask: s [x][y] -o0, Vx, y, such that LTS, s (i-1)x LT end if if B, Starm and (i 1) x B, UT Endmaxx then
23:
24:
25:
On chip, apply mask: s [x][y] 00, Vx, y, such that UTS (i-1)xB, LTF end if
26:
27:
On chip, compute p(d) exp(Si) - L,) B.
28:
On chip, compute dV; < dV, (P! T dO; Bexd
29:
On chip. compute dp 1 = do,v] B.
On chip. compute as / p(f) (ap Dí) E R B...
30:
31:
Load dQ; from HBM SRAM, then on chip, update dQ; - dQ, dQ, Kj R B. and write back to HBM.
32:
On chip. compute dk, dk, + as Q B.xd
33:
end for
34: Write j dV, to HBM.
35:
end for
36:
Return dQ. dK, dv.
were applied: n E [1,3] for N e (0, (0,4096 and n E [1.4] for N e (4096, 8192]. During sampling, set the minimum document length to 128 for SFT, LoRA, and DPO, and 512 for RM. Padding lengths did not exceed 128 for SFT, LoRA, DPO, and 512 for RM.
Assuming document had sequence length L, it was further divided into source tokens and target tokens. SFT and LoRA were represented as (Question, Answer) pairs, DPO as (Question, Answer Answer,,nnwwey Answerz) with two answers, and RM, having to 6 answers, was standard-
14

---

Published as a conference paper at ICLR 2025
SPT4K
SPT
SFT 18K
SFT 32K
SPT 64K
SFT 128K
DPD
BPO
12K
1.8
5.8
1.
86
0.1
96
16
1890
67
1390
50
0.1
1.0
RM
64K
128K
Scart
Sparchy
Spardit
Spartiy
Sparety
Qpardhy
Figure 6: Sparsity distribution of synthetic dataset for end-to-end training throughput testing.
ized to have answers: (Question,Answerr.... Answers). . and k answers based on the training task, with k being for SFT and LoRA, 2 for DPO, and 6 for 0.1% 0.27.
Thus, was partitioned into query
RM. The length of each answer was randomly determined from the range T40.1K 1+0.2k making each answer approximately 10% 20% of the query length. Consequently, the query length was calculated as L minus the total answer lengths. For each sequence length N, we collected 240 valid samples and categorized them into bins by sparsity p, illustrated Figure
Table Hyperparameters and distributed configurations for various scales of Llama2 models.
Model
LLama2-7B LLama2-13B
LLama2-70B
Batch Size
AceSteps
Sharding Stagel Degree
Tensor Parallel Degree
4
Pipel.ine ine Parallel Degree
2
Sequence Parallel
A.2.2 EXPERIMENTAL CONFIGURATION AND DISTRIBUTED STRATEGY
We evaluated different model scales of Llama2 (7B. 13B, 70B) across sequence lengths ranging from 4K to 544K, comparing against two dense methods: Vanilla Attention and FlashAttention DenseMask. All end-to-end throughput experiments with eight NVIDIA A800-SXM 80G GPUs, totaling 32 GPUs. The objective was not to optimize peak performance for each configuration but to assess scalability with longer sequences; thus, uniformly enabled full recomputation. Model parameters and computations utilized the BFloat data type, while gradient accumulation and communication employed Float32. The hyperparameters and distributed strategies for different scales are detailed in Table I
conducted four servers, each equipped
A.2.3 END-TO-END TRAINING MEMORY CONSUMPTION
Figure 2 in the main paper reports the end-to-end training throughput. We also recorded the peak memory consumption, presented in Figure 7 Notably, the memory usage of FLASHMASK increases significantly slower than that of dense methods. However, the figure also indicates that FLASHMASK'S memory consumption still escalates rapidly with longer sequence lengths, primarily due to increased activation memory from longer sequences, as shown in Table 2 The Param & Opt State column indicates the memory consumption for parameters, gradients, and optimizer states, with sharding stage 1 applied. Activations refers to the memory consumed by the inputs of the 32 decoder layers. Peak Mem One Layer represents the peak memory usage when full recompute is employed. Total denotes the overall memory consumption for FlashAttentio without the attention mask.
15

---

Published as a conference paper at ICLR 2025
SFT
LoRA
DPO
RM
Liama 2-78
120
244
56
114
256
544
30
120
244
Usir Haans-a 32
78
136 324
125
Liane 7:770
Figure 7: End-to-end training peak memory consumption across varying sequence lengths for different
Llama2 model scales four downstream training tasks: SFT, LoRA, DPO, and RM.
Table 2: Memory consumption comparison between FlashAttention without attention mask and FLASHMASK on the Llama-2 7B model. The observed differences in total memory footprint are attributed to memory fragmentation effects.
Sequence Length (K)
Param & Opt State
Activations
Peak Mem One Layer
Total
FLASHMASK
13.12
0.00
0.73
13.86
13.14
1.29
13.73
16
13.12
1,00
2.30
16.63
16.01
32
64
13.12
4,00
9.89
27.02
25.57
128
256
13.12
16.00
39.56
68.69
68.81
A.3 END-TO-END TRAINING CONVERGENCE VERIFICATION
We selected the Llama 3.1 8B model to verify convergence across four downstream tasks involving the fine fine-tuning and alignment training of large language models: SFT, LoRA, DPO, and RM. SFT and LoRA utilized the same dataset, validated using
-tt.r-ittttmiiiuuue-------- Ivison al. (003) For DPO and RM, which both employ (Question, Answer) data formats, we used the Hugg winggaacett /ul fults aalt:ee t. feedbaeecccc ck ____bba___ binnari i zed Tunstal et al. (2023) dataset for validation. We consistently applied linear learning rate decay strategy, with warm-up steps set to 3% of the total training steps. The Adam W optimizer was used with Bi = 0.9 and Br 0.999. Model parameters and computations utilized the BF data type, while gradient accumulation and communication employed Float32. The maximum training sequence length was set to 8K. Distributed parallelism combined sharding and tensor parallelism. Additional hyperparameters are listed Table
A.4 SPARSITY- RELATEE EXPERIMENTS
Prior to each run, we perform 10 warm-up iterations, followed by 100 runs of the kernel computation, recording the average execution time (in milliseconds) using CUDA Events. We aimed to verify that the computational complexity of FLASHIMASK scales linearly with block sparsity in the attention mask. We report the total latency for the kernel's 's forward and backward passes. This validation was performed on sequences of length 32K for three common mask types: Causal Document Mask, Share Question Mask, and Document Mask corresponding to downstream training tasks large language models such SFT, DPO/RM, and pre-training of vision models like NaVIT, respectively.
16

---

Published as a conference paper at ICLR 2025
Table The configuration of end-to-end training loss convergence verification.
Task
Dattset
Leaming Rate
Training Step
Epochs
Busch Size
AccSceps
GPUs
Sharding Degree
TP Degree
SFT
1ulu-22-dfl 2-e----mmmmmitee
2e-05
12000
16
32
12 stt
00002
12000
32
DIP
HugginggFaeHHane Fachha hinarieed
HuggingFaodd Saggiigggaollllllllneeeeeeoo fewlee/eeeeeeeecccccccc.. binurized
bead din
Soqurnce
S3aa Wader
1o0a
141
Mack
U14A
DM5
Quodke
S3dyy Wader
516
Euc do
Nata LM Mark
1058
12 svenant
@Kopan Mark
Spead
Figure 8: Kernel forward and backward speed (head dim 64, BF16) on A100-SXM 80G GPU. Flex Attention using PyTorch (6.0ddew022499204cu1122
A.4.1 DATA CONSTRUCTION METHON
Causal Document Mask and Share Question Mask are causal attention types with block sparsity values in the range [0.5, 1.0], while the Document Mask is bidirectional attention type with block sparsity values in [0.0, 1.0]. We partitioned the sparsity values into buckets: buckets for causal types and 20 buckets for bidirectional types, each with intervals of 0.05, ensuring the number of samples per bucket ranged between 10 and 20.
For the Causal Document Mask, given maximum sequence length, we limited the number of documents to [2, 20). We randomly sampled the number documents and then sampled the length of each document such that the total sequence length equaled the maximum sequence length. Following each sample, we calculated the block sparsity and assigned it to the corresponding bucket until each bucket met the required number of samples. The data sampling process for the Document Mask was similar; to ensure coverage of all sparsity levels, the number of documents was limited to [2, 10].
The data sampling for the Share Question Mask differed slightly. The number of documents was limited to [1, 5]. We first sampled the length of each document to sum up to the given maximum sequence length. Each document was then partitioned into a Question and Answers, ensuring there was one Question and 2 to 6 Answers. As before, after each sample, we calculated the block sparsity and allocated it to the appropriate bucket until all buckets were adequately populated.
In total, we sampled 182 samples for the Causal Document Mask, 175 for the Share Question Mask, and 374 for the Document Mask.
A.5 KERNEL PERFORMANCE COMPARISON
A.5.1 TESTING METHON
Both FLASHMASK and FlexAttention exploit sparsity in the attention mask to skip fully masked blocks, thereby reducing redundant computations. To provide intuitive comparison, we employ the TFLOPs/s metric for evaluation. For each test case, we assess both the forward and backward computations. Prior to each run, perform 10 warm-up iterations, followed by 100 runs of the kernel computation, recording the average execution time (in milliseconds) using CUDA Events. Based on the block sparsity in the attention mask, calculate the FLOPSs for single and subsequently compute the TFLOPs/s.
A.5.2 DATA CONSTRUCTION METHOD
We conducted detailed comparisons across varying batch sizes and sequence lengths (8K, 32K, 128K), different head dimensions (64, 128), and numbers of heads. We fixed the total number of tokens at
17

---

Published as a conference paper at ICLR 2025
Table 4: Kernel speed details (8K, head dim 128, BF16) On A100-SXM 80G GPU.
Operaion
Speniy
11:45
116:1
$5.11
11803
755.4
0.95
11:33
19165
134
Makin Widene
10:30
4135
33.0
1355
111
1.64
131.3
18:54
18:33
000
Dearieer Madi
21,1
123 00
134.00
4.54
11.34
1533
141.11
106.0
115.17
6.74
Techtaiion
(lubal Striig Wandep
3166
7953
16141
2.0
733
13:13
13334
0047
18.35
055
Prils Ltl Docenec Miak
Masi
1838
122.15
5452
3.0
0.22
18:12
122.35
14.0
100.3
Dretk Molk
218.4
28.12
134.8
5.17
201.55
5.36
153
123.00
0.50
Red ie
07:3
715.66
7016
3.5%
110:5
D-9
FLE
5630
224.14
11.50
13.08
62.55
331.5
106.0
30153
Dre
10101
21.21
193.1
Saking Mintice
235
1025
123
3.35
118.4
002
Caael Doceenc
242
4.5
16480
006
2512
13.66
PLASEMASK
Shan Querdin
3010
53.0
166
435
091
Caasi Biedovis Mas
3677
10.04
115.00
2775
34.77
13.12
166.00
111.73
BAT
Protic Lat Chasal Mick
55.30
18333
211.1
11.6
27.08
19.13
25541
19842
119.0
BST
Raniin bnicten Masi
2005
155
154.00
3.55
22.35
11333
19217
165.6
04
0.59
Table 5: Kernel speed details (32K, head dim 128, BF16) A100-SXM 80G GPU.
Opendios
Spanty
Casal
31..0
663 11
081.50
14.13
123.03
154.4
113.35
140.36
0.50
Slu8 Windin
11:33
15.11
151:55
111 50
Madk
25/0
$7.10
113.20
116
113.07
0195
Deounce Maik
39.20
126.39
160.0
2722
14.23
23.00
141.1
111.00
110
000
Clubei Slidng Wadda
61.55
15882
26036
120
23.33
12.55
151.01
119.4
124.0
EXTT
Pesic Ltl Docren Viok
51.4
115
155
(05
Peeti Maak
3900
06618
33.00
11855
130 43
0.2
Raddon Exisos Maic
30150
53844
10600
11535
0.50
PLE
36.25
15813
116488
115.02
240.25
231.25
204.0
211.4
Slaking Windin
36.0
13.55
455
1554
156.0
150 35
Mal
41.0
1934
3.7
13894
134
(lubal Stiing Waddey
66.56
16855
396.1
6.50
13.55
155.1
057
Canal dooto
0..0
13116
11.58
171
6ST
Philk Lat Cuava Muck
20730
00829
104.0
4610
130
134.17
202.7
18.00
kudien Redienlskii Bakiiin Macc
190.0
401.0
467.6
06.33
65.30
133
18174
668
128K; by varying the sequence length, we computed the corresponding batch size. With the hidden size fixed at 4096, varying the head dimension allowed to determine the number of heads.
To encompass a broader range of block sparsity cases in the attention mask for a given sequence length, we utilized constructed data for testing. Given a test sequence length, we defined the document count range as n E [Docmin- Do We first sampled the number of documents and then sampled the length of each document such that the total length equaled the test sequence length. For the Share Question Mask type, we further partitioned each document into one Question and 2 to 6 Answers. The document count ranges were [3, 7] for 8K, [10, 14] for 32K, and [11, 15] for 128K. For each sequence length, we generated five test data samples.
Table 6: Kernel speed details (128K, head dim 128, BF16) on A100-SXM 80G GPU.
Makull
Operaiia
Thee
(pa)
TOTAL Teac
T7LOX
TOTTAL TTLOO.
TTLOON
BW TILOOW
TOTAL IELO0
Spalty
FrBl
5311 06
090
3.1.15
153 c4
055.15
15383
111.5
143 34
Dro
Sloinn Windin
10: 50
140.4
245.90
11:14
11:33
13.15
163.1
133
115 54
Decunee Mai..
159.4
40431
634.38
25.05
54.21
13.39
151.4
125.20
138.81
Glubal Stikig Wandey
216.0
664.5
182.00
122.00
151.33
131.1
136.55
Canel Ri Recdoose Madk
20127
611.5
000
133 60
19153
153 56
04
000
Phets LMl Cuasal Mack
00205
30010
40010
110
413
013.0)
15401
1525%
140.4
Rudion Exictio Mare
1077.00
50
125770
14055
152 34
163.6
135.2
106.1
115.30
6.0
PLE
126.27
140116
411-33
31148
125.09
955.10
231.4
26.73
213.3
Canel
63112
1619 12
251044
14858
152.00
45310 10
325.1
019.0
1:0
Canel Docenec Mirk
K746
11535
20130
1130
12855
4437
145.7
11310
16780
Shave Quericc Mart
6266
10850
173.4
7.41
13
25.00
11801
14..1
150
Camni Recdovie Mask
210-2
123.0
12523
133.0
Prekk
Misk
45.2
112.1
14813
LMd
Chare bllk
291.00
433.5
197.33
037
Qi-quane Mark
20000
136.66
415.55
135.3
154.4
022
Maic
72699
163
A.5.3 EXPERIMENTAL RESULTS
Figure g
presents the comparison of total TFLOPs/s for forward and backward passes between FLASHMASK and FlexAttention when the head dimension is 64. Similar results are illustrated in Figure in the main paper for a head dimension of 128. Across all cases, FLASHMASK outperforms Flex Attention in terms of total TFLOPs/s for both forward and backward passes, with improvements
18

---

Published as a conference paper at ICLR 2025
Table 7: Kernel speed details (8K, head dim 64, BF16) on A100-SXM S0G GPU.
Operaiia
Speniy
118.3
Canal
S6.11
756.11
3.95
11:33
153.0
Makin Wineie
1000
40:13
51.43
135
1.3
1.64
131.11
K1:11
18.00
000
Dearieer Madi
51.7
10233
133
4.54
11.34
1333
143.0
100.09
116.3
6.74
(lubal Striig Wandep
2383
5350
105
2.0
733
13:13
131.43
13.47
05.00
055
Prits Ltl Docnnee Maak
Masi
18.0
5464
3.0
0.22
28:12
19.52
100.3
Preti Molk
28.12
132.8
5:01
151.3
5.36
151.5
123
0.50
2905
3.5%
00:75
D-9
FLE
67M
345.6
55..5
11.50
13.08
62.55
323.3
15.10
155 13
Dre
21.21
Saikng Winlice
R030
34.00
133
3.25
12614
03.33
194.4
002
Casel Doceuac
4.3
12831
006
2011
01.8
113.50
PLASEMASK
Shan Querdin
1760
160
333
091
Caasi Biddoiie Mas
34.22
10823
134.4
2775
34.37
23:12
187.35
142.0
140.17
Protic Lat Caara Mick
5665
18222
24180
11.6
27.06
15.13
185447
15177
166.0
8.5
Ranien bnicten Masi
51.35
136-3
214.5
3.55
22.33
11.33
13515
14241
163.0
0.59
ranging from 4.2% to 53.6%. FLASHMASK achieves 33.6% to 55.1 % of the theoretical maximum FLOPs/s on the A100 GPU. Tables 9 detail, for each test mask case, the sparsity, and the forward and backward computation latency, TFLOPs, and TFLOPs/s.
Table 8: Kernel speed details (32K, head dim 64, BF16) on A100-SXM 80G GPU.
Mihha
Opereiiie
PW Trme
Bo Tiee
197AA)
Teancee
FM
111.0P
19TERS
TOTTL
TRECEN
PM THEIRS
BW
TRLOPU
TOTAL İHH OPL
Spunty
Causl
215.33
011.5
356.1
13.35
$5.30
151.53
131.1
138.3
0.00
Caoa Dweree Misk
N:17
50.05
116
116
144
17.56
181.33
100.0 55
110.50
Star Quesiiin Mart
16.25
18:15
14
3.53
3.43
1:33
135.33
55.33
01.0
(59)
Clubal
344.3
114.8
Ricdooia Mark
0111
19721
268.2
3.00
33.35
11.30
125.48
087
Prati 12d Docanee Mask
3691
11720
3.0
13.43
100.00
114.5
095
OK-qpine Stak
20:4
648.2
364.0
33.00
$4.07
11855
19227
15855
131.2
122
Casl
12541
550.0
076.31
$5.30
123.0
152.5
0.50
Slul Wiuline
113.7
Mack
3609
08111
3.66
4.39
13.6
155.25
Dvoumen Maik
66.99
10.14
16100
2.75
14.02
13.19
12h.3
134.12
PLASIMAS
(kuee Skidng Wendiwe
0566
18542
25020
020
23.23
1253
13576
16290
148.6
100
Protic but Doceme Wisk
26-3
13.44
100
155
4.5
13.00
153
110.1
135.3
Qi-qpane Mall
185.00
78508
3409
83.23
119.3
194.99
163.77
168.0
0,2
Randan Ericdos Mase
223.0
318.22
20.08
15.30
123.0
195.00
125.66
134.15
0.50
Table 9: Kernel speed details (128K, head dim 64, BF16) on A100-SXM 80G GPU,
Mokal
Operition
121145
TPLOPW
Spunty
PLB
177977
563453
7354-72
251.48
133 03
183.15
155155
125:71
136.74
Mlakng Waddow
16.00
343
-41.8
17:11
43.37
03.25
15114
125
134.06
300.0
22.05
1535
25
(00
15331
11:.0
(Mde) Wendew
22.79
085.1
12110
132.35
133.54
Casal dooto Mart
33696
135.0
EST
LMd Misk
1557
121
Prdiic Cress Mark
(06
615
90
118
6.7
Raidon Ericsio Masi
100755
363.60
405120
14008
152
403.6
135.55
15.00
160.4
0.50
PLB
L80.21
302831
315412
201.4
125.09
105.1
196.4
Casal Window
102800
301148
19855
151
13155
11.25
13.55
176
11
1167
141.5
69
Stav Quodicc Masa
72.72
125.4
10730
1.40
15.5
13.00
351.31
14577
138
107
Canel Hi Rivboi Mark
73mn
538
765
133.00
155
163
53
057
Peets LM Coasal Mick
95312
90611
112.0
613.0
031
Randon Zokcis Maie
990.22
226575
3155/0
148.8
453.6
153.5
155.4
156.0
0.50
B FLASHMASK APPLICATION IN INFERENCE
In the main body of our paper, we focused on the application of FLASHMASK during the training phase of large-scale models. However, it is important to highlight that FLASHMASK is equally effective during the inference stage. In this appendix, we provide detailed experimental results demonstrating the efficacy of FLASHMASK in inference, comparing it with state-of-the-an attention implementations, including FlashInfer Ye et al. (2025).
B.1 EXPERIMENTAL SETUP
Our experiments were conducted on an NVIDIA A100-SXM 80G GPU using FlashInfer version 0.1.6, CUDA 12.1, PyTorch 2.4, and BF16 data type. We set the batch size to 1, with 32 query/output
19

---

Published as a conference paper at ICLR 2025
Table 10: Performance Comparison on Causal Document Mask 8K, 32K, and 128K Sequence
Lengths
Method
Seq Length
Sparsity
FW Time (ms)
FW TFLOPs
FW TFLOPs/s
FlashInfer SparseMask
8,192
0.8806
9.33
0.1313
13.95
FlashInfer DenseMask
8,192
0.8806
11.93
0,1313
11.01
FLASHMASK
8,192
0.8806
0.96
0.1313
135.07
FlashInfer SparseMask
32,768
0,9532
54.77
0.8233
15.01
FlashInfer DenseMask
32,768
0,9532
184.20
0.8233
4.47
FLASHMASK
32,768
0.9532
5.99
0.8233
137.09
FlashInfer SparseMask
131,072
0.9558
788.94
12.4435
15.77
FlashInfer DenseMask
131,072
0.9558
2,948.23
12.4435
4.22
FLASHMASK
131,072
0.9558
84.13
12.4435
147.58
Table 11: Performance Comparison on Shared Question Mask at 8K, 32K, and 128K Sequence
Lengths
Method
Sey Length
Sparsity
FW Time (ms)
FW TFLOPs
FW TFLOPs/s
FlashInfer SparseMask
8,192
0.9324
6.12
0.0743
11.76
FlashInfer DenseMask
8,192
0.9324
11.94
0.0743
6.23
FLASHMASK
8,192
0.9324
0.73
0.0743
98.49
FlashInfer SparseMask
32,768
0.9742
32.87
0.4537
13.74
FlashInfer DenseMask
32.768
0,9742
184.40
0,4537
2.46
FLASHMASK
32,768
0.9742
4.59
0.4537
98.26
FlashInfer SparseMask
131,072
0,9751
443.80
7.0146
15.80
FlashInfer DenseMask
131,072
0.9751
2,948.89
7.0146
2.38
FLASHMASK
131,072
0.9751
61.57
7.0146
113.21
heads and 8 key/value heads, each with head dimension of 128. The evaluation included typical attention masks such as the Causal Document Mask, Document Mask, and Shared Question Mask.
To ensure compatibility with FlashInfer's 's sparse mask representation (where the mask block size C 64), we adapted the datasets from Section |A.522 Section[A.5.21 so that each sub-document sequence length is divisible by 64. We defined the mask block size based FlashInfer's Block Sparse Row (BSR) API parameters and C, and matched the filing block size to the kernel's operational dimensions.
B.2 COMPARISON WITH FLASHINFER
We compared FLASHMASK
with FlashInfer's dense
mask and sparse mask various
API
Liingle_rrefili with la_kkv_caaa kv_cache) (BlockSparseAttent ionMrapper) across
API
sequence lengths (8K, 32K, and
128K tokens). The results are summarized in Tables 10 to 14
Efficiency Analysis: FLASHMASS ASK consistently outperformed both the dense and sparse implemen-
tations of FlashInfer in terms of TFLOPs/s, particularly addressing the inefficiencies observed with FlashInfer's ' dense mask API. While FlashInfer with sparse masks showed performance gains with increasing mask block sizes (R, C 2 16), such large block sizes are seldom practical due to the nature of attention patterns in real-world applications.
In the FlashInfer s lines 1234-1241) p. token-by-token dense masks lead to sigmificant inefficiencies by per- forming unnecessary computations on fully masked blocks. Additionally, in Flashinfer's Bl rockkparreetttent ioWrrapprr, smaller mask block sizes increase the padded batch size (nblks ks (padded bat___baahhh__eeeeee _iii,,,,,,,ee 1, num_kv_hee (e__heads)) negatively impacting performance due to suboptimal kernel hyper-parameter tuning. In contrast, FLASHMASK efficiently computes only the required tiling blocks, avoiding redundant calculations, and thus achieves superior TFLOPs/s.
sarefillwwtth_ae
implementation (see prefill. cuh
20

---

Published as a conference paper at ICLR 2025
B.3 IMPACT OF MASK BLOCK SIZE
For the Document Mask, we investigated the effect of varying the mask block size (R/C) perfor- mance. Although FlashInfer DenseMask and FLASIMAASK do not utilize specific R/C values, we included them the comparison for completeness. The results for sequence lengths of 8K, 32K, and 128K tokens are presented in Tables 12 and 14 respectively.
Table 12: Performance on Document Mask at 8K Sequence Length with Varying Mask Block Sizes
Method
R/C
Sparsity
FW Time (ms)
FW TFLOPs
FW TFLOPs/s
FlashInfer SparseMask
0.7868
15.39
0.2344
15.19
FlashInfer SparseMask
2
0.7613
8.57
0.2624
30.48
FlashInfer SparseMask
0.7613
4.31
0.2624
60.57
FlashInfer SparseMask
0.7613
3.23
0.2624
80.97
FlashInfer SparseMask
0.7613
1.65
0.2624
158.55
FlashInfer SparseMask
32
0.7613
1.51
0.2624
172.61
FlashInfer SparseMask
64
0.7613
1.51
0,2624
172.82
FlashInfer DenseMask
0.7613
11.91
0.2624
22.03
FLASHMASK
0.7613
1.66
0.2624
156.82
Table 13: Performance on Document Mask at 32K Sequence Length with Varying Mask Block Sizes
Method
R/C
Sparsity
FW Time (ms)
FW TFLOPs
FW TFLOPs/s
FlashInfer SasrseMask
0.9064
104.65
1.6460
15.73
FlashInfer SparseMask
2
0.9064
52.46
1.6460
31.36
FlashInfer SaskseMask
0.9064
25.96
1.6460
63.47
FlashInfer SparseMask
8
0.9064
19.68
1.6460
83.59
FlashInfer SparseMask
0.9064
9.87
1.6460
166.58
FlashInfer SparseMask
32
0.9064
8.89
1.6460
185.13
FlashInfer SparseMask
64
0.9064
8.89
1.6460
185.16
FlashInfer DenseMask
0.9064
183.99
1.6460
8.95
FLASHMASK
0.9064
11.73
1.6460
139.84
Table 14: Performance on Document Mask at 128K Sequence Length with Varying Mask Block
Sizes
Method
R/C
Sparsity
FW Time (ms)
FW TFLOPs
FW TFLOPs/s
FlashInfer SparseMask
0,9116
1,571.12
24.8848
15.84
FlashInfer SpsrseMask
0.9116
783.62
24.8848
31.75
FlashInfer SparseMask
0.9116
391.20
24.8848
63.61
FlashInfer SparseMask
0.9116
288.97
24.8848
86.11
FlashInfer SparseMask
16
0.9116
145.13
24.8848
171.45
FlashInfer SparseMask
32
0.9116
131.31
24.8848
189.50
FlashInfer SparseMask
64
0.9116
131.33
24.8848
189.48
FlashInfer DenseMask
0.9116
2,946.81
24.8848
8.44
FLASHMASK
0.9116
172.81
24.8848
143.68
21


Reviews:
[{"review_id": "tYwrirqQWm", "reviewer": "Reviewer_Uyji", "summary": "The paper introduces a novel compression scheme for the attention mask where only the boundary indices of the masks are stored for every column. For a specific set of attention masks, it is sufficient to store two sets of boundary indices for every column to represent the attention mask. This reduces the memory complexity of attention masks from quadratic on sequence length to linear on sequence length, enabling handling of longer context lengths. The column-wise sparse representation is also used to skip fully masked blocks increasing the overall compute efficiency of the attention mechanism. This technique is augmented with FlashAttention algorithm for efficient computation of the attention mechanism and the modified algorithm for both forward pass and backward pass are presented. The experiments section shows that FlashMask is faster than FlashAttention dense method by up to 3.22x and can achieve up to 60.23% more throughput than FlexAttention. The proposed method also doesn’t alter the convergence during training.", "questions": "You mentioned FlexAttention can also exploit sparsity by skipping computation on fully masked blocks. If that’s the case where’s compute throughput advantage of FlashMask coming from?\nWould the Block-Sparse FlashAttention be able to handle the mask types described in Fig 1(a)? If yes, that should be used instead of the DenseMask variant for the throughput comparisons across the paper. If not, please mention why.\nIn Fig 4b, why is the FlexAttention’s memory utilization lower than that of FlashMask for sequence length lower than 16K?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written and easy to understand. The results section is elaborate with a wide range of benchmarks to demonstrate the advantages of the proposed methods. The appendix section and the analysis with synthetic data to corroborate the claims are very insightful. The compute and memory utilization advantages of FlashMask are well demonstrated. The proposed sparse representation scheme is novel and should be adopted wherever applicable for its memory efficiency and ability to support longer context lengths.", "weaknesses": "While the results section shows that FlashMask achieves higher computational efficiency, I’m not sure if it’s attributable to the proposed columns-wise sparse representation.\nThe computational efficiency of FlashMask comes from skipping computation on entirely masked blocks as discussed in section 4.3. However, this technique is also used in Block-Sparse FlashAttention and FlexAttention. The advantages of FlashMask over Block-Sparse FlashAttention and FlexAttention in terms of computational efficiency is not clear.\nAlso as mentioned in the paper, the idea of column-wise sparse representation used in FlashMask is limited to specific attention patterns. Any other pattern can’t be handled even naively.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hajGyIH9s4", "reviewer": "Reviewer_TH6Z", "summary": "The paper proposes an efficient sparse mask representation by using composition of LT and RT range for expressing complex patterns. The proposed mask is compatible with FlashAttention-2 and can bring speed-up when applied.", "questions": "1. What would be the block size supported by Flashmask? Namely, what would be the granularity of mask/unmask chunks?\n2. How does different block/chunk size affect the speed-up, in different mask types?\n3. When tiling, it seems some mask may lead to different workload among thread blocks, which could hurt the overall performance. Is there any mitigation to this?\n4. Can we have a comparison between the theoretical FLOPs reduction wrt wall-clock speed-up for different mask types?\n5. How does tensor parallel and pipeline parallel affect the speed-up?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper open-sourced a rather general sparse self-attention representation framework, which could facilitate many research and production attempts in the field.\n2. The implementation is practical, shown wall-clock speed-up over FlashAttention-2.", "weaknesses": "1. It seems the implementation is limited to Paddle. It would be good to see if it can also be made more general so that the Torch/Megatron community can also leverage the framework.\n2. Inference support is missing. It would make more sense to discuss how such sparse mask can be put into actual inference/serving. \n3. [1] was published earlier, and also provide a general sparse self-attention training & serving framework. It would be ideal to also cite [1].\n\n[1] S2-Attention: Hardware-Aware Context Sharding Among Attention Heads", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7h4hM00zjC", "reviewer": "Reviewer_gr4R", "summary": "This paper introduces an extension for SDPA that supports different types of masks in an easy-to-understand way. The method is novel due to its sparse representation. Experiments show FlashMask outperforms FlexAttention by a significant gap.", "questions": "1. How can this technique be integrated with page attention?\n2. Can tree-based speculative decoding benefit from this customized attention?\n3. Can you report evaluation results on machines such as P100, V100, A10G, and H100? (other than A100)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Although this representation might be similar to COO/CSR/CSC, this is the first time I have ever seen these techniques used in attention, one of the most important operators in LLMs.", "weaknesses": "This paper lacks two baselines:\n1.  Flashinfer with dense masks;\n2.  Flashinfer sparse mask (https://docs.flashinfer.ai/api/python/sparse.html);\n\nAlthough Flashinfer does not support backward, I believe it is an important baseline for SOTA attention implementation. If this comparison is presented, I will raise my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LXOT1a6Bxh", "reviewer": "Reviewer_gKFL", "summary": "The article presents Flash Mask - a method to incorporate a wide class of attention masks into flash attention. Algorithmically, this means that the  non trivial structuring of the mask is incorporated into the online calculation of the softmax operation involved in self attention without materializing the full mask and paying quadratic in sequence length memory. Furthermore, this algorithm is implemented in a hardware aware fashion, much like flash attention to minimize memory access and data movement while exploiting the thread group level parallelism in GPUs. Empirically, this method is benchmarked against the dense mask of flash attention 2 and also against flex attention - an alternative state of the art method to incorporate structured masks in efficient attention computations and the method presented shows noticeable gains both in inference and in training.", "questions": "Following up on the limitations, I would like the know whether the authors think that arbitrary masking patterns can be incorporated in a GPU friendly manner or if memory efficient implementations of such masking patterns would require alternative hardware architectures (such as those developed by Cerebras).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The strengths of this paper are in the novelty of the incorporation of a class of structured masks into the online softmax calculation involved in memory efficient attention and further into the hardware aware version of the algorithm which is flash attention, and in the  comprehensive validation of the superior efficiency of the algorithm when benchmarked against dense masking in conventional flash attention 2 and in flex attention. \nThe algorithm for the forward and backward pass are presented clearly and the empirical results are presented clearly.", "weaknesses": "As admitted by the authors, this method cannot handle irregular masking patterns within a column of the mask, or completely arbitrary masking patterns.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)