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](https://github.com/flashinfer-ai/flashinfer/blob/v0.1.6/include/flashinfer/attention/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