FlashMask: Reducing the Complexity of Attention Computation through Sparse Mask Representation

Thanks for your review.

For each weakness you mentioned:

1. It is crucial to highlight the advantages of this method over related work to help readers fully understand its significance. However, in the subsection "Attention Optimization Techniques," the authors only mention the drawback of FlashAttention and discuss its relationship to their work. The introduction of other related works is confusing and makes it difficult to comprehend their relevance to this paper. The overall conclusion, "Both of the previously discussed solutions either compromise precision or yield only marginal enhancements in efficiency. Conversely, our proposed FlashMask is capable of delivering exact computations." is general and non-specific. It is unclear which methods compromise precision and which ones only offer marginal improvements.

Reply:

Vanilla Attention:

• Advantages: Simple implementation and easy to understand. It was initially proposed in the seminal paper Attention is All You Need.
• Disadvantages: Slow computation speed and high memory usage, with a quadratic relationship to sequence length, resulting in a complexity of $O(N^2)$.

Memory Efficient Attention:

• Advantages: Lower memory usage with a complexity of $O(\sqrt{N})$, significantly reducing resource requirements.
• Disadvantages: Slower computation speed compared to IO-aware FlashAttention. See the kernel latency in Figure 3 of the attached PDF file.

FlashAttention:

• Advantages: Fast computation speed with no loss in accuracy compared to Vanilla Attention.
• Disadvantages: The official implementation only supports limited mask lacks the mask functionalities required for downstream tasks such as SFT, DPO, and RM.

FlashAttention with DenseMask:

• Advantages: Third-party implementations support DenseMask in FlashAttention, enabling various masks for downstream tasks like SFT, DPO, and RM.
• Disadvantages: High memory usage with a quadratic relationship to sequence length, resulting in a complexity of $O(N^2)$.

FlashAttention VarLen:

• Advantages: Fast computation speed and support for variable-length sequences, suitable for training tasks such as SFT.
• Disadvantages: Limited support for sparse computation modes, unable to support training tasks like DPO and RM.

Other Approximate Algorithms (e.g., Reformer and Linformer):

• Advantages: Achieve lower memory usage and faster computation speed through approximate sparse attention calculations.
• Disadvantages: Model convergence performance cannot match that of Full Attention, leading to reduced accuracy.

2. In the experiments, the baseline algorithms are limited to Vanilla Attention and FlashAttention. Are there more efficient Transformer algorithms that could be used for comparison? If not, the authors should explain the rationale behind the selection of these specific baselines.

Reply: As illustrated in the previous reply, Vanilla Attention, FlashAttention with DenseMask are the proper baselines for comparision. We add other efficient Transformer algorithm called Memory Efficient Attention (MEA) in Figure 3 in the attached PDF. FlashAttention-DenseMask is the most competitive baseline for comparison. FlashAttention with VarLen is faster than FlashAttention-DenseMask, but it can be only used in SFT and the performance of FlashMask is the same with it.

3. As a non-expert in this field, I found the writing of this paper confusing. For instance, the initialism "HBN" is introduced without any explanation or context.

Reply: We supposed that the "HBN" you mentioned is "HBM". Thanks for your suggestions. HBM is shorted for High Bandwidth Memory, which is the global memory of the GPU.

For each question you mentioned:

1.Has there been any prior work on efficient attention mask computation?

Reply: please refer to our replies of the first and the second weakness.

2.Why was the FA-Varlen method not included in the DPO and RM scenarios? Could the authors provide an explanation in their paper?

Reply: Thank you for your suggestions. FA-Varlen cannot represent the sparse attention mask in the DPO and RM scenarios as mentioned in our reply of the first weakness. We will provide a comprehensive explanation in the Camera-Ready version.

3.In Figure 4, the latency of both FA-Window and FlashMask is almost identical. If the authors aim to demonstrate the efficiency of FlashMask, could they explain the experimental results more clearly?

Reply: Thank you for advices. Figure 4 is used to explain that FlashMask can be not only speed up the SFT, DPO and RM scenarios, but also speed up the sliding window sparse attention mask training scenarios. FlashMask is a more general method than the existing techniques, which can be used in many other NLP training tasks.

4.In Figures 3 and 5, it appears that the performance and efficiency of FA-Varlen are comparable to FlashMask in the SFT setting. Could the authors clarify this comparison?

Reply: In fact, if the sparse attention mask can be represented by either the FlashMask or FA-Varlen method, the efficiency of these two methods are almost the same as shown in Figure 3 and 5. However, FA-Varlen failed to represents the sparse attention masks in DPO and RM. Therefore, the FlashMask method is a more general method in almost all of the NLP tasks, while the usage of FA-Varlen method is limited to SFT.

Thanks for your review.

We were sorry that we had not clarified our key points previously. In the SFT/DPO/RM training scenarios, the sparsity of the attention mask is usually natural, but not intended to speed up training while sacrificing accuracy. For examples, as illustrated in Figure 1, the sparsity of the attention mask may be from: (1) padding mask, (2) InToken mask, and (3) question and answering mask. We were not discussing how to design a new sparse attention mask, balancing the model accuracy and training efficiency. Instead, we took advantage of the sparsity attention property in these NLP tasks itself to speed up the training process. The key points of our paper include: (1) we proposed an efficient sparse mask representation method which can be used in SFT, DPO, RM and many other scenarios; (2) based on our sparse mask representation, we proposed an efficient kernel implementation to speed up the training process. Since we did not introduce extra approximate calculations (the sparsity comes from the NLP task itself), the model accuracy should be exactly the same with or without our method. We also conducted some extra experiments evaluating model quality in the attached PDF file (Figure 1 and Table 1).

Due to time constraints, we have only supplemented the experiments in SFT and DPO scenarios, and the benchmark tests were conducted on a limited set of datasets to demonstrate the accuracy preservation of our method. We used the Huggingface LLaMA2-7B pre-trained model and conducted SFT training using Vanilla Attention, FlashAttention-DenseMask, and FlashMask, followed by DPO training. SFT and DPO use Packing/InToken data training strategy. The SFT phase was performed on the "allenai/tulu-v2-sft-mixture" dataset, using the AdamW optimizer with beta1=0.9, beta2=0.999, a learning rate of 2e-05, an end learning rate of 1e-07, weight decay of 0.0, a total training step count of 12000, a warmup step count of 360, and a global train batch size of 16. The DPO phase was conducted on the "HuggingFaceH4/ultrafeedback_binarized" dataset, using the AdamW optimizer with beta1=0.9, beta2=0.999, a learning rate of 1e-06, an end learning rate of 1e-07, weight decay of 0.0, a total training step count of 1600, a warmup step count of 10, and a global train batch size of 16. We will provide a comprehensive comparison in the Camera-Ready version. It showed that the training loss and the benchmark results are the same between our speedup method and the baselines.

The detailed replies are as follows:

1. The scaling ability of the proposed method deserves further verified on large scale datasets.

Reply: In Section 4.3 and 4.4, we had conducted some experiments to show the speedup and memory saving of our proposed method. The testing dataset included the LongBench (an open-source benchmark for long context understanding) and our synthetic data. As we mentioned before, our method is used to speed up the training process without sacrificing accuracy. We believed that the scaling ability of our method should be good on large scale datasets.

2. While the paper demonstrates FlashMask's effectiveness in specific scenarios, it may lack broader evidence on how it performs across different types of NLP tasks or diverse datasets.

Reply: Our proposed method performed extact computation when the NLP tasks used sparse attention mask itself. Therefore, different datasets would not affect the performance of our proposed methods. The SFT, DPO and RM are almost the most important LLM downstream training tasks, and we believed that our proposed method is a general speedup method in most of the NLP tasks.

3. The paper could provide more detailed insights into how FlashMask handles different sparsity levels and the impact on various model sizes and complexities.

Reply: Our proposed method performed extact computation when the NLP tasks used sparse attention mask itself. Therefore, the model accuracy after using our method would not be changed, regardless of the sparsity levels, model sizes and complexities.

Thanks for your review.

1. Even though FlashMask achieves significant improvement in the memory efficiency of sparse attention, the key idea is similar to FlashAttention, but it is just for sparse attention mechanisms. Based on this fact, the novelty of this paper is not strong. I recommend the authors explain why the red part in the algorithm is designed and why it is unique for sparse attention.

Reply: There are two contributions in this paper: (1) we proposed a column-based sparse attention mask representation, instead of using a large dense tensor to represent the sparse attention mask. (2) we proposed an efficient kernel implementation using our sparse attention mask representation based on FlashAttention. We think that the novelty of this paper is that we found the common sparse attention mask representation in most of the NLP tasks, and speeded up the attention phase combining this representation and FlashAttention. Our sparse attention mask representation is general in most of the NLP tasks, including SFT, DPO, RM, etc.

As shown in Figure 1 (b) of the main text, the Attention Mask required for a long sequence composed of three sequences in the SFT scenario, with sequence lengths [4,3,3]. The FMS values are [4,4,4,4,7,7,7,10,10,10], and FME does not need to be set, defaulting to the maximum number of rows. For example, in the second column, the FMS value is 4, indicating that the elements in rows $[4,10)$ are masked.

As shown in Figure 1 (c) of the main text, the Attention Mask required for a long sequence composed of three sequences in the bidirectional SFT scenario, with sequence lengths [4,3,3]. FlashMask uses two pairs of FMS and FME to describe the bidirectional scenario: the lower left part is denoted as $FMS_1$ and $FME_1$, and the upper right part as $FMS_2$ and $FME_2$. $FMS_1$ is [4,4,4,4,7,7,7,10,10,10], and $FME_1$ does not need to be set, defaulting to the maximum number of rows; $FMS_2$ does not need to be set, defaulting to 0; $FME_2$ is [0,0,0,0,4,4,4,7,7,7]. For example, in the fourth column, the $FMS_1$ value is 7, indicating that the elements in rows $[7,10)$ are masked, and the $FME_2$ value is 4, indicating that the elements in rows $[0,4)$ are masked.

As shown in Figure 1 (d) of the main text, the Attention Mask required for one Query and two Answers in the DPO scenario, with the Query length being 4, Answer1 length being 3, and Answer2 length being 3. The FMS values are [10,10,10,10,7,7,7,10,10,10]. For example, in the zeroth column, the FMS value is 10, indicating that the elements in rows $[10,10)$ are masked.

2. The authors only present optimization for column-based sparse attention. The performance for other types of sparse attention is unknown. If the proposed approach can be applied to all sparse attention, the contribution of this paper is extremely great. However, the existing version is not comprehensive.

Reply: Thanks for you suggestions. In fact, it is hard to handle all kinds of the sparse attention masks. If we use a large dense tensor to represent the sparse attention mask, the costs of the memory usages and HBM access are not acceptable; if we use a sparse-coo liked method to represent the sparse attention mask, it is hard to implement an efficient CUDA kernel as well. In this paper, we found that the sparse attention mask in most of the NLP tasks can be represented using a column-based way. Therefore, we can use FMS and FME in our paper to represent the mask efficiently. We believed that our method is general to most of the NLP tasks like SFT, DPO, RM, etc.

3. Based on the experiments, the practical latency is not significantly reduced as compared to other methods, even though the theoretical complexity is from N^2 to N. Besides, the authors do not provide results for accuracy.

Reply: The theoretical complexity $O(N^2)$ corresponds to the FA-DenseMask and Vanilla Attention in Figure 3. We provided a new curve in the attached PDF file (Figure 3) comparing the kernel latency of the FA-DenseMask, Vanilla Attention and our proposed FlashMask method. It showed that the theoretical complexity is from $O(N^2)$ (FA-DenseMask, Vanilla Attention) to $O(N)$ (our FlashMask method). We also conducted some extra experiments evaluating model quality in the attached PDF file (Figure 1 and Table 1).

Due to time constraints, we have only supplemented the experiments in SFT and DPO scenarios, and the benchmark tests were conducted on a limited set of datasets to demonstrate the accuracy preservation of our method. We used the Huggingface LLaMA2-7B pre-trained model and conducted SFT training using Vanilla Attention, FlashAttention-DenseMask, and FlashMask, followed by DPO training. SFT and DPO use Packing/InToken data training strategy. The SFT phase was performed on the "allenai/tulu-v2-sft-mixture" dataset, using the AdamW optimizer with beta1=0.9, beta2=0.999, a learning rate of 2e-05, an end learning rate of 1e-07, weight decay of 0.0, a total training step count of 12000, a warmup step count of 360, and a global train batch size of 16. The DPO phase was conducted on the "HuggingFaceH4/ultrafeedback_binarized" dataset, using the AdamW optimizer with beta1=0.9, beta2=0.999, a learning rate of 1e-06, an end learning rate of 1e-07, weight decay of 0.0, a total training step count of 1600, a warmup step count of 10, and a global train batch size of 16. We will provide a comprehensive comparison in the Camera-Ready version. The model accuracy results showed that the model accuracy using our method is the same with that of the FA-DenseMask and Vanilla Attention.

As we near the author-reviewer discussion deadline, we seek your feedback on our rebuttal. Your insights have been crucial in improving our work, and we're grateful for the time and effort you've dedicated to our manuscript.

Thank you for your valuable guidance. We look forward to your response and are ready to make further adjustments if needed.

Best regards,

Authors

Thanks for your review.

1. The paper seems to make claims that are unsubstantiated by experiments. In the abstract and introduction, the paper claims speedup without sacrificing model quality. However, there is no experiment evaluating model quality in the experiments. This is a critical flaw.

Reply: We were sorry that we had not clarified our key points previously. In the SFT/DPO/RM training scenarios, the sparsity of the attention mask is usually natural, but not intended to speed up training while sacrificing accuracy. For examples, as illustrated in Figure 1, the sparsity of the attention mask may be from: (1) padding mask, (2) InToken mask, and (3) question and answering mask. We mainly focused on how to speed up the training process when using these kinds of natural sparse attention masks. Our baselines are the non-speedup methods using the same sparse attention mask. Therefore, we claimed speedup without sacrificing model accuracy.

Due to time constraints, we have only supplemented the experiments in SFT and DPO scenarios, and the benchmark tests were conducted on a limited set of datasets to demonstrate the accuracy preservation of our method. We will provide a comprehensive comparison in the Camera-Ready version. We used the Huggingface LLaMA2-7B pre-trained model and conducted SFT training using Vanilla Attention, FlashAttention-DenseMask, and FlashMask, followed by DPO training. SFT and DPO use Packing/InToken data training strategy. The SFT phase was performed on the "allenai/tulu-v2-sft-mixture" dataset, using the AdamW optimizer with beta1=0.9, beta2=0.999, a learning rate of 2e-05, an end learning rate of 1e-07, weight decay of 0.0, a total training step count of 12000, a warmup step count of 360, and a global train batch size of 16. The DPO phase was conducted on the "HuggingFaceH4/ultrafeedback_binarized" dataset, using the AdamW optimizer with beta1=0.9, beta2=0.999, a learning rate of 1e-06, an end learning rate of 1e-07, weight decay of 0.0, a total training step count of 1600, a warmup step count of 10, and a global train batch size of 16.

From the train loss curves in Figure 1 of the attached PDF, it is evident that FlashMask ensures convergence accuracy comparable to the baselines (Vanilla Attention and FlashAttention-DenseMask). Additionally, FlashMask and FlashAttention-DenseMask exhibit identical loss values. The evaluation metrics from the benchmark tables (Table 1 in attached PDF) indicate that FlashMask achieves the same accuracy on par with FlashAttention-DenseMask. Therefore, FlashMask, in comparison to FlashAttention-DenseMask, is an exact algorithm with no loss in accuracy.

2. Further, the contribution of the paper is unclear. Block-sparsity is already supported in FlashAttention (see section 3.3 of FlashAttention). It is unclear how this paper is different. There are also more recent works such as "Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention" (NeurIPS 2023), which seem to be strictly more expressive in features than this paper."

Reply: SFT, DPO and RM are the very important scenarios in the NLP downstream training tasks. Although block-sparsity is already supported in FlashAttention, the block-sparsity does not match the sparse attention mask pattern in SFT, DPO and RM shown in Figure 1. The paper you mentioned ("Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention") contributed to speed up the training process when using the QK-Sparse attention and Hash-Sparse attention. But it also failed to represent the sparse mask attention in SFT, DPO and RM shown in Figure 1. We proposed an efficient sparse mask representation method which can be used in SFT, DPO, RM and many other scenarios, even including the QK-Sparse/Hash-Sparse attention in the mentioned paper. Thanks to our efficient representation, we proposed an efficient kernel implementation to speed up the sparse attention phase. Our approach is more general than the existing research works and can be used in most of the NLP training scenarios.

For example, as shown in Figure 2 of the attached PDF, FlashMask can represent QK-sparse and Hash-sparse masks using FlashMask. In the figure, $S_1$ denotes the starting row index of the mask in the lower left triangle, and $E_2$ denotes the ending row index of the mask in the upper right triangle. The diagonal elements should be considered part of the lower left triangle. In FlashMask, $E_1$ (representing the ending row index of the mask in the lower left triangle) does not need to be set and defaults to the maximum number of rows, while $S_2$ (representing the starting row index of the mask in the upper right triangle) defaults to 0. Here, the $S_1$ and $S_2$ represent the $FMS$ in Algorithm 1, and the $E_1$ and $E_2$ represent the $FME$ in Algorithm 1.

For example, in Figure 2(a) (in the attached PDF file) for QK-sparse, in the second column, $S_1$ is 6, indicating that the elements in the rows $[6,6)$ are masked, and $E_2$ is 3, indicating that the elements in the rows $[0,3)$ are masked. In Figure 2(b) (in the attached PDF file) for Hash-sparse, in the first column, $S_1$ is 1, indicating that the elements in the rows $[1,8)$ are masked, and $E_2$ is 1, indicating that the elements in the rows $[0,1)$ are masked. In Figure 2(b) (in the attached PDF file) for Hash-sparse, in the second column, $S_1$ is 5, indicating that the elements in the rows $[5,8)$ are masked, and $E_2$ is 2, indicating that the elements in the rows $[0,2)$ are masked.