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

W1: 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.

We appreciate this thoughtful observation about modern GPU architectures. We want to clarify several points:

• The theoretical distinction between small cache $M = o(d^2)$ and large cache $M = \Omega(d^2)$ regimes is important regardless of specific hardware, as it characterizes a fundamental phase transition in I/O complexity behavior. While modern GPUs like H100/A100 may typically operate in the large cache regime, our analysis proves FlashAttention's optimality in this regime, providing theoretical justification for its widespread adoption.
• The small cache analysis remains valuable since our theory suggests that if, in the future, the commonly used hidden dimension size increases while some commercial GPU cache sizes remain insufficiently large, our algorithm designed for small cache sizes would become relevant and useful. Hence, our work provides theoretical insights that could guide future developments in attention mechanisms tailored to evolving hardware limitations.

Our goal is to provide a comprehensive theoretical foundation for understanding attention's I/O complexity across all possible regimes, which can inform both current implementations and future hardware/algorithm design.

W2: If possible, please provide some preliminary results for your algorithm.

We thank the reviewer’s valuable suggestion. We refer the reviewer to the global response Implementation part.

Q1: Under what hardware and model will the attention computation fall in different cache cases?

The current network architectures usually set $d = 128$. In this case, the dividing point is approximately $d^2 \\times \\text{size\\_of(data type)}$, e.g., $16,384 \\times 32$-bit = $65.5$ KB for float32. For NVIDIA A100 GPU, the size of each streaming multiprocessor (SM/L1 cache) is 192 KB, so we should choose FlashAttention. However, for old GPUs such as NVIDIA GTX1060, the size of each SM is 48 KB, so the algorithm for the small cache size is suitable.

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

Our theory is applicable in this situation, with GPU bandwidth serving as an illustrative example. In practical applications, there are various memory hierarchies; however, as long as they adhere to a two-layer architecture, our theory remains relevant. Our two-level memory hierarchy is a simplified representation of real-world GPUs, such as the NVIDIA A100, which is also the case discussed in Section 2.1 of FlashAttention [10]. According to link, an NVIDIA A100 GPU includes 108 (192 KB) streaming multiprocessors (SM/L1 cache), a 40 MB L2 cache, and 80 GB of HBM2 memory. The SMs execute arithmetic and other instructions, while data and code are accessed from DRAM through the L2 cache. In our theoretical model, we simplify this structure by treating the data transfer as between the L1 cache and HBM2 memory. As for L2 cache, we omit it for theoretical simplicity since it serves as an intermediate role in I/O exchanges.

[1] Sun, M., Liu, Z., Bair, A., & Kolter, J. Z. A Simple and Effective Pruning Approach for Large Language Models. ICLR 2024

[2] Clarkson, K. L., & Woodruff, D. P. Low-rank approximation and regression in input sparsity time. JACM 2017

[3] Bubeck, S., Cohen, M. B., Lee, Y. T., & Li, Y. An homotopy method for lp regression provably beyond self-concordance and in input-sparsity time. STOC 2018

[4] Li, Y., & Woodruff, D. Input-sparsity low rank approximation in schatten norm. ICML 2020.

[5] Musco, C., & Woodruff, D. Is input sparsity time possible for kernel low-rank approximation?. NeurIPS 2017

[6] Woodruff, D., & Zandieh, A. Near input sparsity time kernel embeddings via adaptive sampling. ICML 2020

[7] Diao, H., Song, Z., Woodruff, D., & Yang, X. Total least squares regression in input sparsity time. NeurIPS 2019

[8] Jayaram, R., Samadian, A., Woodruff, D., & Ye, P. In-database regression in input sparsity time. ICML 2021

[9] Pagh, R., & Stöckel, M. The input/output complexity of sparse matrix multiplication. ESA 2014

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

W1: Sparse attention lower bound only considers number of nonzero entries.

Thank you for your comments! We appreciate the reviewer’s insight regarding additional factors, such as sparsity patterns and data formats, which may impact I/O complexity. We agree that the specific sparsity patterns, e.g., structured N:M sparsity [1], and row/column/block-wise sparsity, are important and widely used in modern GPU and CUDA implementation. However, we want to emphasize that:

• Our sparsity problem setting is standard for studying sparse instances and our lower bound for sparse attention is based on input sparsity (i.e., the number of nonzero entries), which is widely studied and applied in prior work, including linear regression [2, 3], low-rank approximation [3, 4, 5], kernel methods [5,6], total least squares regression [7], and in-database regression [8]. We believe such bounds are also useful in attention mechanisms and suggest the potential algorithm design.
• Our sparsity problem setting is more general than the setting of specific sparsity patterns, thus providing foundations for studying the more special sparsity pattern setting. Once we make additional assumptions like sparse patterns, then we restrict to more specific problem instances, and then the IO complexity might change, e.g., exploiting the sparsity pattern can lead to a smaller IO that breaks the general lower bound in Thm 4.5.
• We believe our I/O complexity bound in this problem setting is tight, in the sense that our sparsity lower bound matches the upper bound from the general case (Thm 1.1). We give the lower bound for sparse attention computation which matches the lower bound for standard attention computation when matrices are dense. Our bound is tight considering the family of problem instances with sparse patterns (but not so on the smaller family of instances like those with special sparse patterns).

We will certainly consider these aspects in our future research, specifically exploring how sparsity patterns and data formats impact the theoretical and computational aspects of attention computation.

W2: Dividing point of small and large caches for sparse attention is vaguely defined.

Thanks for your suggestion. We revised our paper to add the discussion of the dividing point for sparse attention in Line 152-153 & 420-425. We provide a precise mathematical analysis for the dividing point between small and large caches for sparse attention. The dividing point is $M = Z^2\_\mathrm{input} / Z\_\mathrm{QK}$, which also matches the dense case.

Q1: 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$?

Here, we provide a specific example. The most network architectures set $d = 128$. In this case, the dividing point is approximately $d^2 \\times \\text{size\\_of(data type)}$, e.g., $16,384 \\times 32$-bit = $65.5$-KB for float32. For the currently used modern GPUs, it is generally more practical to employ FlashAttention (the algorithm suited for the large-cache case). However, our theory suggests that if, in the future, the hidden dimension size of the network architecture increases while some commercial GPU cache sizes remain insufficiently large, our algorithm designed for small cache sizes would become highly relevant and useful. Hence, our work provides theoretical insights that could guide future developments in attention mechanisms tailored to evolving hardware limitations.

Q2: T​​he paper claims that Algorithm 6 achieves the optimal for small cache sizes. Could you provide some practical evidence for this optimality?

We thank the reviewer’s valuable suggestion. We refer the reviewer to the global response Implementation part.

Thanks for your appreciation. We are so glad that you like our work and comment on our work as a solid theory paper. For the implementation, we refer the reviewer to the global response Implementation part.