FlashMask: Efficient and Rich Mask Extension of FlashAttention

Dear Reviewer,

Thank you for your insightful evaluation and the positive feedback on our work. We are pleased to hear your interest in the potential extension of FlashMask's capabilities.

Weaknesses: As admitted by the authors, this method cannot handle irregular masking patterns within a column of the mask, or completely arbitrary masking patterns.

While it is true that FlashMask does not currently support arbitrary masking patterns, it effectively manages the common mask types utilized in LLM training. We are dedicated to enhancing the expressiveness of FlashMask by addressing some limitations inherent in the current two range-based ($LTS$, $LTE$, $UTS$, $UTE$) sparse representation. Specifically, we aim to improve its capability to handle multi range-based masks and introduce expression-based descriptions, which would remove the necessity for explicit attention mask inputs. For instance, if an attention mask can be described using a formula or expression, it potentially eliminates the need for explicit mask data by allowing the kernel to compute mask positions directly. Furthermore, we are focused on refining FlashMask's implementation to reduce unnecessary memory access and condition checks, particularly in relation to fully masked blocks, thus further enhancing its performance.

Q: Following up on the limitations, I would like to 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).

A: Based on the information we found on the Cerebras website, Cerebras' hardware architecture features large global memory, extensive SRAM, and high memory bandwidth. Theoretically, these characteristics allow it to incorporate arbitrary masking patterns more effectively than traditional GPUs.

Thank you once again for your constructive feedback, which continues to guide our development efforts.

Best regards

Q: How can this technique be integrated with page attention?

A: FlashMask's innovative column-wise mask representation significantly reduces memory complexity by using an index range vector to characterize each column. It classifies tiling blocks into Fully masked, Partially masked, or Unmasked categories, allowing computations to be skipped for Fully masked blocks and optimizing other blocks by eliminating redundant mask applications. On the other hand, Page Attention deals with KV Cache tokens that are managed across discontinuous physical storage based on page size granularity. Since FlashMask reduces memory complexity from $O(N^2)$ to $O(N)$, the necessity for page management is eliminated. These two techniques are orthogonal, facilitating seamless integration. By leveraging FlashMask's capabilities, Page Attention can support a more sophisticated expression of masks.

Q: Can tree-based speculative decoding benefit from this customized attention?

A: Indeed, FlashMask supports tree-based speculative decoding methods, as depicted in SpecInfer's Tree-based Parallel Decoding example (Figure 4), with sequences like $LTS=[1, 2, 3, 4, 5, 6, 7, 8]$, $LTE=[1, 2, 3, 4, 5, 6, 7, 8]$, $UTS=[0, 0, 0, 0, 3, 3, 2, 2]$, and $LTS=[0, 0, 0, 0, 4, 4, 6, 6]$.

Q: Can you report evaluation results on machines such as P100, V100, A10G, and H100? (other than A100)

A: The current version of FlashMask extends from FlashAttention-2 without substantial code restructuring, thus it is not supported on legacy architectures like P100 and V100. Unfortunately, due to the lack of H100 and A10G, we have yet to conduct evaluations on Hopper and other Ampere GPUs. However, the methodology itself is adaptable across hardware platforms.

Comparison of FlashInfer DenseMask (single_prefill_with_kv_cache), FlashInfer SparseMask (BlockSparseAttentionWrapper), and FlashMask on Document Masks at an 128K Sequence Length, with Varying R/C

Although FlashInfer DenseMask and FlashMask do not have specific R/C values, the masks they represent are consistent under different R/C due to the total sequence length and each sub-document's sequence length being divisible by 64. We have included them in the table for comparison with FlashInfer SparseMask, conducting multiple tests to ensure a fair evaluation.

Comparison of FlashInfer DenseMask (single_prefill_with_kv_cache), FlashInfer SparseMask (BlockSparseAttentionWrapper), and FlashMask on Document Masks at an 32K Sequence Length, with Varying R/C

Although FlashInfer DenseMask and FlashMask do not have specific R/C values, the masks they represent are consistent under different R/C due to the total sequence length and each sub-document's sequence length being divisible by 64. We have included them in the table for comparison with FlashInfer SparseMask, conducting multiple tests to ensure a fair evaluation.

Comparison of FlashInfer DenseMask (single_prefill_with_kv_cache), FlashInfer SparseMask (BlockSparseAttentionWrapper), and FlashMask on Document Masks at an 8K Sequence Length, with Varying R/C

Although FlashInfer DenseMask and FlashMask do not have specific R/C values, the masks they represent are consistent under different R/C due to the total sequence length and each sub-document's sequence length being divisible by 64. We have included them in the table for comparison with FlashInfer SparseMask, conducting multiple tests to ensure a fair evaluation.

Dear Reviewer,

Thank you for engaging with our work and for highlighting the importance of extending evaluation to include FlashInfer as a baseline. We greatly appreciate your insights, which aid in advancing our research on FlashMask.

Response to Weaknesses and Questions:

We acknowledge the value that FlashInfer brings, bridging training and inference applications. Should our paper be accepted at ICLR 2025, we plan to reference FlashInfer in the camera-ready version and include the following experimental results in the appendix. Our primary focus with FlashMask has been addressing the extensive high bandwidth memory (HBM) requirements posed by diverse attention mask types in large-scale model training. Nevertheless, FlashMask's capabilities extend into inference stages as well, warranting consideration among leading attention mechanisms.

We defined the "mask block size" based on FlashInfer's description, utilizing the BSR API's parameters R and C. The "tiling block size" was designed to match the operational dimensions of the kernel. Our experiments were conducted on an A100-SXM 80G GPU using the official version of FlashInfer v0.1.6, with CUDA 12.1, PyTorch 2.4, and BF16 as the data type. The configuration settings included batch_size = 1, num_qo_heads = 32, num_kv_heads = 8, and head_dim = 128. We compared FlashMask against the dense mask API single_prefill_with_kv_cache and the sparse mask API BlockSparseAttentionWrapper (with varying R/C values) in FlashInfer. Typical attention masks, such as the Causal Document Mask, Document Mask, and Shared Question Mask, were selected for evaluation.

The datasets used were from Section A.5.2, but were slightly modified to ensure that each sub-document sequence length was divisible by 64, allowing for experiments with a FlashInfer sparse mask where C=64. As shown in the experimental results below, FlashMask demonstrated superior TFLOPs/s, effectively addressing the inefficiencies observed with FlashInfer's dense mask API. While FlashInfer with sparse masks shows considerable improvements as the mask block sizes (R and C) increase, especially when R, C ≥ 16, such large mask block sizes are rarely required in practice.

Within FlashInfer’s single_prefill_with_kv_cache DenseMask implementation in prefill.cuh#L1234-L1241, the use of token-by-token dense masks results in substantial inefficiencies, especially since calculations for fully masked blocks are entirely unnecessary. Regarding FlashInfer's BlockSparseAttentionWrapper, the mask block column $C$ functions as the page size. Smaller mask block sizes lead to a marked increase in the padded batch size nblks(padded_batch_size, 1, num_kv_heads), which adversely affects performance because of non-optimal kernel hyper-parameter tuning. In contrast, with larger mask block sizes, FlashInfer maximizes the advantages of BSR's sparse representation by calculating only the necessary tiling blocks and avoiding excess calculations for fully masked blocks, thereby achieving high TFLOPs/s.

The hyper-parameter tuning results of FlashInfer are as follows:

R1C1: request_idx=8180, packed_qo_len=4, kv_len=2496, qo_chunk_size=16, kv_chunk_size=2496, num_tiles_q=1, num_tiles_kv=1
R2C2: request_idx=4094, packed_qo_len=8, kv_len=960, qo_chunk_size=16, kv_chunk_size=1344, num_tiles_q=1, num_tiles_kv=1
R4C4: request_idx=2047, packed_qo_len=16, kv_len=720, qo_chunk_size=16, kv_chunk_size=784, num_tiles_q=1, num_tiles_kv=1
R8C8: request_idx=1023, packed_qo_len=32, kv_len=360, qo_chunk_size=64, kv_chunk_size=392, num_tiles_q=1, num_tiles_kv=1
R16C16: request_idx=511, packed_qo_len=64, kv_len=180, qo_chunk_size=64, kv_chunk_size=196, num_tiles_q=1, num_tiles_kv=1
R32C32: request_idx=255, packed_qo_len=128, kv_len=90, qo_chunk_size=128, kv_chunk_size=98, num_tiles_q=1, num_tiles_kv=1
R64C64: request_idx=127, packed_qo_len=256, kv_len=45, qo_chunk_size=128, kv_chunk_size=49, num_tiles_q=2, num_tiles_kv=1

R2C2: partition_kv=0, padded_batch_size=4096, num_warps_x=1, num_warps_z=4, num_frags_x=1, num_frags_y=8, num_frags_z=2
R4C4: partition_kv=0, padded_batch_size=2048, num_warps_x=1, num_warps_z=4, num_frags_x=1, num_frags_y=8, num_frags_z=2
R8C8: partition_kv=0, padded_batch_size=1024, num_warps_x=4, num_warps_z=1, num_frags_x=1, num_frags_y=8, num_frags_z=8
R16C16: partition_kv=0, padded_batch_size=512, num_warps_x=4, num_warps_z=1, num_frags_x=1, num_frags_y=8, num_frags_z=8
R32C32: partition_kv=0, padded_batch_size=256, num_warps_x=4, num_warps_z=1, num_frags_x=2, num_frags_y=8, num_frags_z=4
R64C64: partition_kv=0, padded_batch_size=256, num_warps_x=4, num_warps_z=1, num_frags_x=2, num_frags_y=8, num_frags_z=4

We hope this response clarifies our approach and experimental scope, and we look forward to any additional feedback you might provide.

Best regards

Dear Reviewer gr4R,

I hope this message finds you well. I am writing to kindly follow up on my response to your feedback. I have addressed the questions and comments you raised, including conducting the additional experiments and providing the requested results. I sincerely appreciate the time and effort you have dedicated to reviewing my work, and I am grateful for your constructive feedback.

If there is anything further you would like me to clarify or expand upon, please do not hesitate to let me know. Additionally, if the rebuttal has satisfactorily addressed your concerns, I would kindly appreciate your consideration in reflecting this in the final scoring.

Thank you once again for your thoughtful feedback and your time. I look forward to your response at your earliest convenience.

Best regards,

The Authors

Q: What would be the block size supported by FlashMask? Namely, what would be the granularity of mask/unmask chunks?

A: In response to your question, we define the "block size" or "mask/unmask chunks" as "mask block size." FlashMask represents masks with token-level granularity, thereby setting the mask block size to 1. This minimal granularity supports arbitrary mask block size. As detailed in sections 4.1 and 4.2, FlashMask divides tiling blocks into Fully masked, Partially masked, and Unmasked categories by representing each column with a range of indices, thereby bypassing computation on Fully masked blocks while applying masks only to Partially masked blocks.

Q: How does different block/chunk size affect the speed-up, in different mask types?

A: Again, defining "block size" or "chunk size" as "mask block size" and differentiating from the computation tiling block size in FlashAttention Kernels, FlashMask maintains a mask block size of 1. Thus, it supports any mask block size. The tiling block size remains unchanged from FlashAttention-2. FlashMask primarily focuses on reducing memory complexity via the novel column-wise mask representation. Our experiments (see Section 5.3, Figure 4a) reveal a strong correlation between block sparsity and processing latency—more structured attention masks result in fewer partially masked blocks, thereby decreasing the computational load and enhancing performance. It should be noted that FlashMask is not yet optimized for maximum performance, as there are additional memory accesses and condition checks that have not been fully optimized. In the future, we plan to further optimize these aspects.

Q: 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?

A: The current FlashMask implementation inherits from FlashAttention 2, assigning each thread block to handle calculations for a given batch, head, and tiling row. Despite potential workload imbalances due to varying mask types within tiling blocks (Fully masked, Partially masked, Unmasked), the SMs (Streaming Multiprocessors) remain fully utilized as multiple thread blocks can be assigned to each SM. Thus, mask diversity does not negatively impact overall performance. That said, future optimizations will focus on reducing unnecessary memory access and conditions, as depicted in Algorithm 1, lines 9-14.

Q: Can we have a comparison between the theoretical FLOPs reduction and wall-clock speed-up for different mask types?

A: As stated in Section 4.3, the theoretical FLOPs reduction is linearly related to the block sparsity ($\rho$) in the attention mask and is independent of the specific mask type. This indicates that kernel latency should be proportional to $\mathcal{O}((1-\rho)T_rT_c)$. The actual speed-up observed in the experiments detailed in Figure 4(a) further supports this relationship, showing a strong correlation between the theoretical and empirical results.

Q: How does tensor parallelism and pipeline parallelism affect the speed-up?

A: FlashMask is an optimization at the operator level, and its effects are orthogonal to those of tensor parallelism and pipeline parallelism. These parallel strategies distribute computation equally across different devices. Therefore, we believe that different parallel strategies have minimal impact on the performance gains provided by FlashMask.

Dear Reviewer,

Thank you for the valuable feedback and questions on our paper. We appreciate the chance to elaborate on our work and address your concerns.

Weaknesses 1: Implementation limited to PaddlePaddle.

Our current implementation of FlashMask is indeed conducted within the PaddlePaddle framework. However, it is integrated as a third-party module in PaddlePaddle and decouples from framework, so it's easy to be integrated into PyTorch or other frameworks. We aim to expand our validation efforts to ensure broader compatibility and adoption across platforms.

Weaknesses 2: Inference support is missing.

FlashMask can certainly be extended to inference applications, like Page Attention and tree-based speculative decoding. Although this paper primarily focuses on training aspects, we acknowledge the importance of discussing inference use cases and plan to explore these in future work.

Weaknesses 3: Missing citation to [1].

We appreciate your suggestion regarding [1] (S2-Attention: Hardware-Aware Context Sharding Among Attention Heads). It is indeed a commendable work, and we intended to cite it in the camera-ready version of our submission. S2-Attention employs a Block-Sparse Attention strategy focusing on approximate attention sparsity calculations at both per-head and per-context-range levels for long-sequence modeling. In contrast, FlashMask emphasizes supporting a wide variety of attention masks used in transformer training through precise computation rather than approximate methods. Our approach reduces memory complexity to $O(N)$ by introducing a novel column-wise mask representation. This enables handling longer sequences efficiently while supporting mainstream attention masks precisely.

Thank you once again for your thoughtful review and questions, and we look forward to refining FlashMask further.

Best regards

Dear TH6Z,

I hope you’re doing well. I wanted to follow up on my rebuttal, where I’ve addressed your questions and comments. I sincerely appreciate your time and constructive feedback.

If there’s anything else you’d like me to clarify or expand on, please let me know. Your insights are invaluable in improving my submission.

Thank you again for your thoughtful review. I look forward to your response.

Best regards,

The Authors

Dear Reviewer,

Thank you for your insightful feedback and questions regarding our paper. We appreciate the opportunity to address your concerns and provide additional clarifications.

Q: 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?

A: Similar to FlexAttention, FlashMask also capitalizes on the sparsity of attention masks. The key distinction is that FlashMask is built upon FlashAttention2, inheriting all its manual and hyperparameter optimizations. As demonstrated in the performance chart on the official FlexAttention blog, the current FlexAttention is slower than FlashAttention2. For more details, you can refer to the blog post at https://pytorch.org/blog/flexattention/#performance. FlexAttention, however, is still in development, utilizing Triton-based compiler technology, and has not yet reached its full performance potential.

Moreover, although FlashMask's current performance is not optimal, as it only achieves 37.8% to 62.3% of the theoretical maximum FLOPs/s on A100 GPUs, we plan to enhance it by reconstructing the traversal method, further improving throughput for all mask types.

Q: Would the Block-Sparse FlashAttention handle the mask types described in Fig 1(a)? If yes, it should be used instead of the DenseMask variant for throughput comparisons across the paper. If not, please mention why.

A: No. As explained in Section 2.3, Block-Sparse FlashAttention represents masks with tiling-level granularity, which makes it unsuitable for handling token-level masks like DenseMask. Previously, DenseMask was the only comprehensive method capable of representing the attention masks shown in Fig 1(a). Therefore, we used DenseMask as a performance baseline in our paper.

Furthermore, Block-Sparse FlashAttention was developed as part of a new approximate attention algorithm. In contrast, our FlashMask is designed to accelerate the attention module when the sparsity of the attention mask naturally arises from the problem itself. For instance, the causal document mask, which is inherently sparse, is commonly used in SFT training. Unlike the approximate computations in Block-Sparse FlashAttention, our approach leverages the inherent sparsity of the mask to achieve acceleration without any loss of precision.

Q: Why is the FlexAttention’s memory utilization lower than that of FlashMask for sequence length lower than 16K in Fig 4b?

A: As discussed in section 2.2, the memory complexity for FlexAttention is $O\left(\frac{N^2}{{B_c}{B_r}}\right)$, where $B_r=128$ and $B_c=128$. FlashMask, on the other hand, employs an innovative column-wise mask representation, maintaining a linear memory complexity of $O(N)$. When the sequence length is less than 16K, $O\left(\frac{N^2}{{B_c}{B_r}}\right)$ is actually less than $O(N)$. Although the FlexAttention blog suggests potential memory savings by increasing the block sizes $B_r$ and $B_c$, its complexity still remains quadratic with respect to the sequence length $N$. Conversely, when the sequence length exceeds 16K, FlashMask's memory usage is lower than that of FlexAttention.

We appreciate your thorough scrutiny and hope these responses satisfactorily address your concerns.

Best regards

Dear Uyji,

I hope you’re doing well. I wanted to follow up on my rebuttal, where I’ve addressed your questions and comments. I sincerely appreciate your time and constructive feedback.

If there’s anything else you’d like me to clarify or expand on, please let me know. Your insights are invaluable in improving my submission.

Thank you again for your thoughtful review. I look forward to your response.

Best regards,

The Authors