Fine-grained Attention I/O Complexity: Comprehensive  Analysis for Backward Passes

Authors: Xiaoyu Li, Yingyu Liang, Zhenmei Shi, Zhao Song, Yufa Zhou
Primary area: learning theory
Keywords: Attention, I/O Complexity, FlashAttention, Gradient, Backward Pass

Abstract:
Large Language Models (LLMs) have demonstrated remarkable capabilities in processing long-context information. However, the quadratic complexity of attention computation with respect to sequence length poses significant computational challenges, and I/O aware algorithms have been proposed.  
This paper presents a comprehensive analysis of the I/O complexity for attention mechanisms, focusing on backward passes by categorizing into small and large cache scenarios. 
Using the red-blue pebble game framework, we establish tight bounds on I/O complexity across all cache sizes. We confirm that the de facto standard I/O aware algorithm FlashAttention is optimal for both forward and backward passes for the large cache size scenario. For small cache sizes, we provide an algorithm that improves over existing methods and achieves the tight bounds. 
Additionally, we extend our analysis to sparse attention, a mainstream speeding-up approach, deriving fine-grained lower bounds for both forward and backward passes and both small and large caches. 
Our findings complete the theoretical foundation for I/O complexity in attention mechanisms, offering insights for designing efficient algorithms of LLM training and inference.

Decision: reject
Reviewer scores: 8, 5, 6

Reviews and discussion:
<thread>
creation date:
1729194330404

thread_type:
official_review

<thread_object>
creator: reviewer_uwmm

<content>
summary:
The paper comprehensively analyzes the I/O complexity associated with attention mechanisms in large language models. The paper identifies a critical point with respect to the cache size in which the I/O complexity changes and analyzes the I/O complexity upper bounds in both cases. The paper shows that in the small cache case, FlashAttention is not optimal and provides an algorithm that improves existing methods. And for the large cache case, the paper confirms the optimality of FlashAttention. The paper further explores the lower bound of sparse attention forward and backward passes. Overall, this paper provides insights for efficient LLM training and inference.

strengths:
1. Comprehensive study of the I/O complexity for attention mechanism. The separation of cases by the cache size seems appropriate. 
2. The paper completes the theoretical I/O complexity in attention, which is a very important component in transformer architectures.

weaknesses:
1. On modern GPUs such as H100 or A100, it is not clear if it will fall under the case of small cache sizes. With older generations of GPUs, the primary concern might not be the I/O complexity but other hardware constraints such as total GPU memory. 
2. The paper could have an implementation of their algorithm on attention with the small cache case and demonstrate that they can outperform FlashAttention. The paper mentions this as a future work but it is still less convincing. If possible, please provide some preliminary results for your algorithm.

questions:
1. Under what hardware and model will the attention computation fall in different cache cases?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730345011167

thread_type:
official_review

<thread_object>
creator: reviewer_dyya

<content>
summary:
This paper provides a comprehensive analysis of the I/O complexity for attention mechanisms, focusing on dense backward and sparse forward/backward by categorizing them into small and large cache scenarios. The paper confirms that FlashAttention, the standard in industry, performs optimally for large cache sizes but is suboptimal in smaller cache environments. The paper introduces an improved algorithm for small caches, achieving tighter bounds on I/O complexity than existing solutions. Additionally, the paper extends this analysis to sparse attention mechanisms, establishing fine-grained lower bounds for both forward and backward passes. Such analysis could guide the design of efficient algorithms for LLM training and inference.

strengths:
1. The paper is well structured. It clearly defines the problem, notations, and the proofs are clearly explained.
2. The paper provides a concrete algorithm with a computation graph that provides practical guidance for achieving optimality.

weaknesses:
1. The sparse attention only considers the number of nonzero entries, which is inadequate for modeling sparse computation. Therefore, the bounds can potentially be further improved. Other factors like sparsity patterns and data format might also affect the I/O complexity bounds.
2. The dividing point of small and large caches for sparse attention is vaguely defined. L1830 says "It depends on whether all the necessary values for computing each output entry can be stored in the cache during the computation." Based on this definition, the cache is whether small or large also depends on the computation. A more precise mathematical definition might help.

questions:
1. What are the practical metrics that distinguish small and large caches? In the paper's notation, the small cache size grows strictly slower than $d^2$. However, it is unclear that given a specific $d$, what size $M$ should the cache have to be counted as a small cache. Could you provide specific examples of small caches in current AI accelerators and discuss how they relate to typical values of $d$?
2. The paper claims that Algorithm 6 achieves the optimal for small cache sizes. Could you provide some practical evidence for this optimality? The comparison between the IO of Algorithm 6 and the IO of FlashAttention implemented in Triton running on a GPU with a small cache will improve the confidence of the claim.
3. The "slow memory" of a real GPU is an L2-Cache, not high-bandwidth memory (L052 "(e.g., GPU high-bandwidth memory)" is not accurate). Are the theorems of this paper still applicable to this setting? Could you clarify how your cache-memory maps onto real GPU architectures?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731162100235

thread_type:
official_review

<thread_object>
creator: reviewer_dt7y

<content>
summary:
In this paper the authors consider the problem of optimizing the I/O complexity of implementing attention mechanisms in LLMs. They consider both forward and backward passes and offer matching lower and upper bounds. All the analyses are done using the red-blue pebble game framework.

strengths:
This is a solid theory paper that settles the I/O complexity for attention mechanisms. They fill some of the gaps that existed in the literature with regard to this I/O complexity. They confirm that FalshAttention, a prior algorithm, is optimal for both forward and backward passes for large cache sizes. For small cache sizes they fill a gap by presenting new algorithms that are optimal. The authors also offer lower bounds on the I/O complexity of sparse attention in the context of forward and backward passes and for both large and small cache sizes.

weaknesses:
Some practical implementation results could enhance the paper.

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1731736669011

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
Thanks for the reviews' valuable comments.

We appreciate the reviewers' recognition of the theoretical rigor and insights of our analysis. Reviewers DT7Y and uwMM agreed that we provided a comprehensive study of the I/O complexity for the attention mechanism, addressing gaps in the literature on the theoretical understanding of attention computation's I/O complexity. Furthermore, reviewers DT7Y and dYYa highlighted our introduction of a new algorithm for small cache sizes that achieves optimal I/O complexity. Reviewer DT7Y also noted our contribution of lower bounds for sparse attention computation, while reviewer dYYa praised the clear structure and presentation of our paper. We thank reviewers dYYa and uwMM for acknowledging the importance of our results and their potential to offer practical guidance for transformer architectures.

We have updated a **revision** for our draft and put all major updates in brown. Note that all line numbers in the rebuttal are corresponding to the revised version.
- Line 152-153 & 420-425: add discussion of the dividing point for sparse attention

Here, we address the common question, and we will respond to other questions individually. 

### Implementation
We sincerely appreciate and agree with the reviewers' suggestion regarding the importance of empirical validation. Implementing FlashAttention [1] requires a sufficiently large cache size, making it unsuitable for comparison with scenarios in small cache sizes. Consequently, it is likely necessary to rewrite the CUDA/Triton functions for both FlashAttention and our proposed algorithm. While empirical validation would be interesting, we admire that implementing and measuring exact I/O operations on modern GPUs presents some technical challenges: 
- Current CUDA implementations combine various optimizations beyond pure I/O considerations.
- GPU hardware abstracts away many low-level memory operations, making it difficult to precisely measure theoretical I/O complexity.
- The theoretical model considers an idealized two-level memory hierarchy, while real GPU memory hierarchies are more complex.

Furthermore, we would like to clarify that this work is a pure theoretical contribution focusing on establishing fundamental mathematical bounds on I/O complexity for attention mechanisms. The key contributions are:
- Providing rigorous mathematical proofs for tight bounds on I/O complexity
- Identifying the critical point $M=d^2$ where I/O complexity behavior changes
- Establishing theoretical optimality algorithms in different cache regimes
- Extending analysis to sparse attention with fine-grained bounds

Our work provides theoretical foundations that may be able to guide future practical implementations. We agree that bridging theory and practice is important and have noted this as future work in our conclusion. We hope that our theoretical bounds remain valid and valuable independent of specific hardware implementations. 

[1] Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. Flashattention: Fast and memory-efficient exact attention with io-awareness. NeurIPS 2022.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper presents a theoretical analysis of I/O complexity for attention mechanisms in large language models, focusing on backward passes and sparse attention. The authors establish tight bounds on I/O complexity for different cache size regimes using the red-blue pebble game framework, confirming FlashAttention's optimality for large cache sizes while proposing an improved algorithm for small cache sizes.

The paper's primary strength is its rigorous theoretical foundation for understanding attention I/O complexity, providing a comprehensive analysis across different scenarios. The work is well-structured with clear proofs and has potential to guide future algorithm design for attention mechanisms.

However, the paper's main weakness is its lack of empirical validation or implementation results, which significantly limits its practical impact. The analysis of sparse attention considers only the number of non-zero entries, potentially overlooking real-world sparsity patterns. Furthermore, there is insufficient discussion on how the theoretical findings translate to current GPU architectures and LLM training scenarios.

The primary reason for rejection is the paper's limited practical applicability to current LLM training. While the theoretical contribution is sound, the focus on small cache sizes, which are atypical in modern GPUs used for LLM training, reduces the paper's immediate relevance. The absence of implementation results makes it difficult to assess real-world performance improvements. To strengthen this work, the authors should bridge the gap between theory and practice, demonstrating how these theoretical insights can improve current attention mechanisms in LLMs.