DuSA: Fast and Accurate Dual-Stage Sparse Attention Mechanism Accelerating Both Training and Inference

We would like to thank the reviewer for the constructive comments provided to help us improve the quality of the paper. We have addressed the issues raised in the comments point-by-point as follows:

Response to W1 and W2: We will move the section about upper bound analysis of the approximation error to the appendix and move more experimental discussions and ablation studies to the main body in the next version.

Response to W3: The current version of DuSA is a static sparse attention mechanism, which cannot preserve the accuracy of VSA without loss in a plug-and-play application. It is designed for accelerating both the training and inference processes. After training, it can match or even outperform VSA. Inspired by XAttention [1], which is an excellent training-free dynamic sparse attention SOTA for accelerating the prefilling stage of LLMs, we have provided a dynamic sparse attention version of DuSA. The dynamic version of DuSA can match the performance of VSA in plug-and-play applications and provides a significant acceleration effect. Details are available in Response to Q1.

Response to Q1: The answer is yes. DuSA can be applied to accelerate the prefilling stage. Based on XAttention [1], we introduce DuSA2 (Strategy 2: Rolling indices-based attention scores fusion) to search antidiagonal/diagonal patterns of attention matrices to filter important or unimportant blocks for further dynamic sparse attention performed within important blocks only. As the following Table R4-1 shows, compared to XAttention, the performance of DuSA2 is closer to that of full attention. As Tables R4-1 and R4-2 show, DuSA2 can use a larger stride (larger sparsity) to achieve similar or higher performance than XAttention, further accelerating full attention. By introducing antidiagonal patterns of XAttention, the performance of DuSA2 can be further improved.

Table R4-1. The performance comparison of different attention mechanisms on LongBench using the Llama-3.1-8B-Instruct model. Note: Higher scores mean better performance.

Table R4-2. Speedup relative to FlashAttention (FlashInfer’s implementation) obtained by different attention mechanisms across different sequence lengths using Llama-3.1-8B-Instruct on a single H20 (96GB) GPU.

Response to Q2: XAttention uses the sum of antidiagonal/diagonal values of the attention matrix as the metric for ranking block importance. While in stage 2, DuSA uses the diagonal values of the attention matrix for block-level mixture to introduce long-range dependencies from other diagonal blocks from stage 1. We take the version of XAttention that uses diagonal patterns for comparison. For the input $**Q**$ and $**K**$ with a size of $(B, H, M, C)$, where $B$ is the batch size, $H$ is the number of heads, $M$ is the sequence length, and $C$ is the embedding dimension of each head, the version of XAttention using diagonal patterns reshapes the $**Q**$ and $**K**$ to a shape like $(B, H, M//stride, stride\times C)$ and performs attention computation using softmax on the final two dimensions to obtain diagonal values of the attention matrices. While DuSA1 (DuSA using Strategy 1) reshapes the $**Q**$ and $**K**$ to a shape of $(B, H, M//stride, stride, C)$ and transposes the third and fourth dimensions to a shape like $(B, H, stride, M//stride, C)$. DuSA1 also performs attention computation using softmax on the final two dimensions to obtain diagonal values of the attention matrices (For satisfying the input requirements of some memory-efficient methods, $(B, H, stride, M//stride, C)$ may be further reshaped to $(B, H\times stride, M//stride, C)$. The softmax operation is still performed on the final two dimensions, hence they give the same results.). The difference between DuSA1 and XAttention is that DuSA1 introduces a multi-stride computation to obtain diagonal values of the attention matrices. And DuSA2 (DuSA using Strategy 2) further sums the stride dimension to obtain $**A**_2$ as Figure 3 of the original manuscript shows.

References:

[1] Xu, R., Xiao, G., Huang, H., Guo, J., & Han, S. (2025). Xattention: Block sparse attention with antidiagonal scoring. arXiv preprint arXiv:2503.16428.

We would greatly appreciate it if you could tell us whether our rebuttal addresses your concerns. If there are no remaining concerns, we would greatly appreciate it if you could consider updating your evaluation of our paper. We welcome any further questions you may have about our work.

We are looking forward to your response.

Thank you for your time and effort in reviewing our paper!

Best regards,

Authors

We would like to thank the reviewer for the constructive comments provided to help us improve the quality of the paper. We have addressed the issues raised in the comments point-by-point as follows:

Response to Q1: We have applied DuSA2 (DuSA using Strategy 2) to the simple vanilla ViT architectures (ViT-T/16 and ViT-L/16). Table R2-1 shows the performance comparison of vanilla ViT architectures and the versions using DuSA2 to replace VSA. The sequence length $m$ of vanilla ViT architecture is relatively short; for the resolution of $224^2$, its maximum sequence length is only $196+1$, and its four levels have sequence lengths: $197$-$197$-$197$-$197$ with a constant embedding dimension $c=64$ for each head. Even for the resolution of $384^2$, its maximum sequence length is only $576+1$, and its four levels have sequence lengths: $577$-$577$-$577$-$577$ with a constant embedding dimension $c=64$. The maximum sequence length of Swin and CSWin is $m=3136$/$9216$ for the resolution of $224^2$/$384^2$ and is significantly larger than $197$/$577$. Their four levels have sequence lengths: $3136$-$784$-$196$-$49$/$9216$-$2304$-$576$-$144$ for the resolution of $224^2$/$384^2$ with a constant embedding dimension $c=32$ for each head. As the ablation study of Table A2 of the original manuscript shows, if the sequence length $m$ is not significantly larger than the head embedding dimension $c$, the acceleration effect will be limited, which cannot completely show the superiority of our method for acceleration. Although the acceleration effect of DuSA on the ViT series is limited due to the short sequence length, DuSA achieves a much better acceleration effect than FlashAttention-2 (FA2) on other monolithic transformers, such as CogvideoX-2B, which uses 30 attention blocks with a constant sequence length $m=17776$ and head embedding dimension $c=64$.

Table R2-1. The performance comparison of vanilla ViT architectures and the versions using DuSA2 to replace VSA. All models are trained from scratch on ImageNet-1K. Block size $b$ is 14 for the tiny model and 24 for the large model. Padding is used to make the sequence length divisible by the block size.

Response to Q2: We have applied DuSA2 for small models like ViT-T/16, which only has 5.8M parameters, and large models like ViT-L/16, which has 307M parameters. Details are available in the Response to Q1. For larger models, please see the acceleration of using DuSA2 on CogvideoX-2B, as Table 5 of the original manuscript shows, and using DuSA2 on the acceleration of the prefilling stage of the Llama-3.1-8B-Instruct model, as Tables R2-2 and R2-3 show.

Table R2-2. The performance comparison of different attention methods on LongBench using the Llama-3.1-8B-Instruct model. Note: Higher scores mean better performance.

Table R2-3. Speedup relative to FlashAttention (FlashInfer’s implementation) obtained by different attention mechanisms across different sequence lengths using Llama-3.1-8B-Instruct on a single H20 (96GB) GPU.

Response to Q3: We will revise the final version of the paper according to reviewer’s suggestion to shorten the bound derivation and use the space instead to provide further insights to the reader in terms of the method, e.g. by including some of the visualizations/analyses currently placed in the appendix – e.g. Fig. A2, or the swapped execution of stages (Table A4), or similar.

Response to Further Comment 1: We will clarify the use of Strategies 1 and 2, move more visualizations from the appendix to the main body, and shorten the bound derivation.

Response to Further Comment 2: We will revise all “hard-aware” to “hardware-aware”.

Response to Further Comment 3: Indeed, as the reviewer pointed out, the nature of the scaled dot product attention calculation leads to a complexity of $O(m^2\times c)$. We will make the statement more rigorous. For vanilla ViT, its sequence length is $m=(196+1)$/$(576+1)$ for a resolution of $224^2$/$384^2$, and its embedding dimension is $c=64$ for each head for both $224^2$ and $384^2$ resolutions. Hence, it still satisfies $m > c$. And in the current stage, a long sequence length $m$ is preferred. Take Swin and CSWin for example, their embedding dimension is $c=32$ for each head for all four levels in both $224^2$ and $384^2$ resolutions. Their four levels have sequence lengths: $3136$-$784$-$196$-$49$ (for the input resolution of $224^2$)/$9216$-$2304$-$576$-$144$ (for the input resolution of $384^2$). CogvideoX-2B uses a sequence length $m=17776$ with $30$ heads. Each head has an embedding dimension $c=64$. Llama-3.1-8B-Instruct has $32$ heads. Each head has an embedding dimension $c=128$. Llama-3.1-8B-Instruct supports a maximum sequence length of 128K. $m$ is also significantly larger than $c$. Our method is designed for the acceleration for the case $m \gg c$. The larger $m$ is than $c$, the more obvious the acceleration effect.

Response to Limitations: We will add some additional limitations (for example, the acceleration effect of DuSA may be limited when $m$ is not significantly larger than $c$) inherent to our DuSA in terms of applicability or other aspects that might be of interest to the reader or practitioner.

Response to Paper Formatting Concerns: We will revise the checklist and provide proper justifications as well as references to the sections where these questions are answered.

There is a small typo in Table R2-1: The term "GFLOPS" should be "GFLOPs". FPS and Inference Memory in Table R2-1 are obtained on a single H20 (96GB) GPU using a batch size of 32 for large (L) models and 512 for tiny (T) models. We are sorry for these unclear points.

Another two points to clarify is: The number of attention layers in ViT-T/16 is 12 and the number of attention layers in ViT-L/16 is 24. We just wrote it to a 2-2-6-2 (layers number for each level) architecture form (for ViT-T/16) / a 2-2-18-2 (layers number for each level) architecture form (for ViT-L/16) for comparison with the four level architecture of Swin-T/Swin-B.

We would like to thank the reviewer for the constructive comments provided to help us improve the quality of the paper. We have addressed the issues raised in the comments point-by-point as follows:

Response to W1 and Q1: The theoretical upper bound is used to provide the worst-case guarantee. A larger bound represents a larger confidence interval, and it may be larger than its real approximation error. Hence, as we mentioned in the limitation, a tighter bound is needed that can more accurately represent the real approximation error.

Response to W2: Tables R3-1 and R3-2 show the memory comparison results. Even compared to SOTAs (ELFATT and FlashAttention-2 (FA2)), DuSA2 (DuSA using Strategy 2) still has the lowest memory usage, and its memory usage is significantly lower than that of VSA.

Table R3-1. Peak memory usage comparison on ImageNet-1K using a single H20 (96GB) GPU. The batch size used is 32 for both inference and training. OOM denotes out of memory.

Table R3-2. Peak memory usage comparison on ADE20K and COCO using a single H20 (96GB) GPU. The batch size used is 1 for inference and 2 for training. OOM denotes out of memory.

Response to W3: The current version of DuSA is a static sparse attention mechanism, which cannot preserve the accuracy of VSA without loss in a plug-and-play application. It is designed for accelerating both the training and inference processes. After training, it can match or even outperform VSA. And we have provided a dynamic sparse attention version, which is training-free and maintains higher performance. The dynamic version of DuSA can match the performance of VSA in a plug-and-play application and provides a significant acceleration effect. We introduce DuSA2 (Strategy 2: Rolling indices-based attention scores fusion) to search antidiagonal or diagonal patterns of attention matrices to filter important or unimportant blocks for further dynamic sparse attention only performed within important blocks. As the following Table R3-3 shows, compared to XAttention [1], the performance of DuSA2 is closer to that of full attention. As Tables R3-3 and R3-4 show, DuSA2 can use a larger stride (larger sparsity) to achieve similar or higher performance than XAttention, further accelerating full attention. By introducing antidiagonal patterns of XAttention, the performance of DuSA2 can be further improved.

Table R3-3. The performance comparison of different attention mechanisms on LongBench using the Llama-3.1-8B-Instruct model. Note: Higher scores mean better performance.

Table R3-4. Speedup relative to FlashAttention (FlashInfer’s implementation) obtained by different attention mechanisms across different sequence lengths using Llama-3.1-8B-Instruct on a single H20 (96GB) GPU.

Response to Q2: We have applied DuSA2 (DuSA using Strategy 2) to the simple vanilla ViT architectures (ViT-T/16 and ViT-L/16). Table R3-5 shows the performance comparison of vanilla ViT architectures and the versions using DuSA2 to replace VSA. The maximum sequence length of the vanilla ViT architecture is relatively short. For the resolution $224^2$, its maximum sequence length is only $196+1$, and its four levels have sequence lengths: $197$-$197$-$197$-$197$ with a constant embedding dimension $c=64$ for each head. Even for the resolution $384^2$, its maximum sequence length is only $576+1$, and its four levels have sequence lengths: $577$-$577$-$577$-$577$ with a constant embedding dimension $c=64$ for each head. While the maximum sequence length of Swin and CSWin is $3136$/$9216$ for the resolution of $224^2$/$384^2$, and their four levels have sequence lengths: $3136$-$784$-$196$-$49$/$9216$-$2304$-$576$-$144$ for the resolution of $224^2$/$384^2$ with a constant embedding dimension $c=32$. As the ablation study of Table A2 of the original manuscript shows, if the sequence length $m$ is not significantly larger than the head embedding dimension $c$, the acceleration effect will be limited, which cannot completely show the superiority of our method for acceleration. Although the acceleration effect of DuSA on the ViT series is limited due to the short sequence length, DuSA achieves a much better acceleration effect than FlashAttention-2 (FA2) on other general transformers, such as CogvideoX-2B, which uses 30 attention blocks with a constant sequence length $m=17776$ and head embedding dimension $c=64$.

Table R3-5. The performance comparison of vanilla ViT architectures and the versions using DuSA2 to replace VSA. All models are trained from scratch on ImageNet-1K. Block size $b$ is 14 for the tiny model and 24 for the large model. Padding is used to make the sequence length divisible by the block size.

References:

[1] Xu, R., Xiao, G., Huang, H., Guo, J., & Han, S. (2025). Xattention: Block sparse attention with antidiagonal scoring. arXiv preprint arXiv:2503.16428.

There is a small typo in Table R3-5: The term "GFLOPS" should be "GFLOPs". FPS and Inference Memory in Table R3-5 are obtained on a single H20 (96GB) GPU using a batch size of 32 for large (L) models and 512 for tiny (T) models. We are sorry for these unclear points.

Another two points to clarify is: The number of attention layers in ViT-T/16 is 12 and the number of attention layers in ViT-L/16 is 24. We just wrote it to a 2-2-6-2 (layers number for each level) architecture form (for ViT-T/16) / a 2-2-18-2 (layers number for each level) architecture form (for ViT-L/16) for comparison with the four level architecture of Swin-T/Swin-B.

The comparison of DuSA2, FlashAttention-2 (FA2), and SpargeAttn [2] on the acceleration of CogvideoX-2B is available in Section 4.6 and Table 5 of the original manuscript.

References:

[2] Zhang, J., Xiang, C., Huang, H., Wei, J., Xi, H., Zhu, J., & Chen, J. (2025). Spargeattn: Accurate sparse attention accelerating any model inference. arXiv preprint arXiv:2502.18137.

We would greatly appreciate it if you could tell us whether our rebuttal addresses your concerns. If there are no remaining concerns, we would greatly appreciate it if you could consider updating your evaluation of our paper. We welcome any further questions you may have about our work.

We are looking forward to your response.

Thank you for your time and effort in reviewing our paper!

Best regards,

Authors

Thanks the authors for the response. It solves some of my major concerns. Would be good to see more analysis in future versions.

Thank you for your response! We will add more analysis including but not limited to highlight and enrich the analysis of experimentally shown good results, add the analysis on training memory consumption patterns in the above rebuttal to the manuscript, and restructure the manuscript according to your and other reviewers’ constructive comments to make the key design decisions and results of the paper more informative in the future version. If you have any remaining concerns or other questions about our work please feel free to tell us.

Thank you again for your time and effort in reviewing our paper!

Best regards,

Authors

We would like to thank the reviewer for the constructive comments provided to help us improve the quality of the paper. We have addressed the issues raised in the comments point-by-point as follows:

Response to W1: ELFATT [1] is a SOTA for ViT acceleration, which was just online several months before DuSA. The FLOPs of DuSA are lower than ELFATT. You may have confused ELFATT with EFFATT [2]. EFFATT is a kernelized linear method, and its theoretical complexity is indeed lower than that of DuSA. But its performance is significantly inferior to DuSA. Its real speed doesn’t show a significant superiority compared to DuSA. Because both stages of DuSA can be further sped up by memory-efficient methods.

Response to Q1: The number of floating point operations (FLOPs) doesn’t represent the real speed of the method, as pointed out by [1,3,4]. Even though some methods have lower FLOPs, their real speed could be slower. Hence, we provide the inference throughput of all methods, which will be more accurate to represent the real speed of each method.

Response to Q2: Yes, it can. XAttention [5] has shown that the diagonal pattern-based attention mechanisms can be used in accelerating the prefilling stage for recent auto-regressive transformer-based (decoder-only) LLMs. Based on XAttention, we introduce DuSA2 (Strategy 2: Rolling indices-based attention scores fusion) to search antidiagonal or diagonal patterns of attention matrices to filter important or unimportant blocks for further dynamic sparse attention only performed within important blocks. As Table R1-1 shows, compared to XAttention, the performance of DuSA2 is closer to that of full attention. As Tables R1-1 and R1-2 show, DuSA2 can use a larger stride (larger sparsity) to achieve similar or higher performance than XAttention, further accelerating full attention. By introducing antidiagonal patterns, the performance of DuSA2 can be further improved.

Table R1-1. The performance comparison of different attention mechanisms on LongBench using the Llama-3.1-8B-Instruct model. Note: Higher scores mean better performance.

Table R1-2. Speedup relative to FlashAttention (FlashInfer’s implementation) obtained by different attention mechanisms across different sequence lengths using Llama-3.1-8B-Instruct on a single H20 (96GB) GPU.

References:

[1] Wu, C., Che, M., Xu, R., Ran, Z., & Yan, H. (2025). ELFATT: Efficient linear fast attention for vision transformers. arXiv preprint arXiv:2501.06098.

[2] Shen, Z., Zhang, M., Zhao, H., Yi, S., & Li, H. (2021). Efficient attention: Attention with linear complexities. In Proceedings of the IEEE/CVF winter conference on applications of computer vision (pp. 3531-3539).

[3] Chen, J., Kao, S. H., He, H., Zhuo, W., Wen, S., Lee, C. H., & Chan, S. H. G. (2023). Run, don't walk: chasing higher FLOPS for faster neural networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 12021-12031).

[4] Zhang, J., Xiang, C., Huang, H., Wei, J., Xi, H., Zhu, J., & Chen, J. (2025). Spargeattn: Accurate sparse attention accelerating any model inference. arXiv preprint arXiv:2502.18137.

[5] Xu, R., Xiao, G., Huang, H., Guo, J., & Han, S. (2025). Xattention: Block sparse attention with antidiagonal scoring. arXiv preprint arXiv:2503.16428.

Thank you for submitting the final rating! We would greatly appreciate it if you could tell us whether our rebuttal addresses your concerns. If there are no remaining concerns, we would greatly appreciate it if you could consider updating your evaluation of our paper. We welcome any further questions you may have about our work.

We are looking forward to your response.

Thank you for your time and effort in reviewing our paper!

Best regards,

Authors