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

## 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.

## 1 Introduction

Large Language Models (LLMs), such as GPT-4 [], Claude [], Llama [], and more recently o1 [] from OpenAI, have demonstrated immense potential to enhance various aspects of our daily lives, including conversational AI [], AI agents [, ], search AI [], AI assistants [, ], and many others.
One of the most emergent abilities of LLMs is dealing with long-context information, which is crucial for processing materials such as academic papers, official reports, and legal documents.
LLMs have proven adept at tackling long-context tasks, such as zero-shot summarization [, ] and maintaining very long-term conversations [, ].
OpenAI’s o1 model [] serves as a significant advancement in this area.
It leverages Chain-of-Thought (CoT) reasoning [, ] and employs Retrieval Augmented Generation (RAG) [, ] to exhibit PhD-level abilities, where both techniques require long context inputs for generation. This proficiency underscores the necessity for developing long-context modeling capabilities within LLMs.

LLMs are primarily based on the Transformer architecture [], whose core component is the self-attention mechanism.
However, the quadratic complexity of attention computation with respect to sequence length dominates the computational FLOPs during long-context training and inference.
To address this issue, FlashAttention [, , ] accelerates attention computation and has become the de facto standard in the industry of LLM training and inference deployment.
The success of FlashAttention lies in its I/O awareness [], accounting for reads and writes to different levels of fast *cache* (e.g., GPU on-chip SRAM) and slow *memory* (e.g., GPU high-bandwidth memory) within the hardware hierarchy.
Leveraging modern hardware design in GPUs, e.g., NVIDIA A100 and H100, efficiently allows FlashAttention to be integrated as a go-to method for LLM training and inference.

For the I/O complexity of exact attention(^1^11In this work, we only consider exact attention computation without any approximation.) forward computation, the theoretical analysis of FlashAttention in DFE^+ [] only provides upper and lower bounds when the cache size $M\in[d,nd]$. Their bounds are only tight in the range of $M=\Theta(nd)$, where $n$ is the input sequence length and $d$ is the hidden dimension.
By fine-grained analysis, a recent work [] provides matching upper and lower I/O complexity bounds of the attention *forward* passes for *any* cache size $M$.
For the I/O complexity of attention *backward* passes, existing work only provides an upper bound for FlashAttention for the cache size $M\in[d,nd]$ [], without known lower bounds.
Thus, the tight bounds for the I/O complexity of attention backward passes are lacking. This raises a natural question:

*What is the optimal I/O complexity of attention backward computations for any cache size?*

In this paper, we address this question and provide matching upper and lower I/O complexity bounds for backward passes of exact attention computation for all cache sizes, completing the picture of I/O complexity for the attention mechanism.

### 1.1 Our Contributions

In this work, we analyze the I/O complexity in the same setting as the existing work of FlashAttention [] and SY []. We consider a two-level memory hierarchy consisting of a small but fast layer called the *cache* and a large but slower layer referred to as *memory*.
The I/O complexity quantifies the data transfer between these two layers, which can be formally defined as a red-blue pebble game [] as in Definition 3.4.
We study the exact attention computation using standard matrix multiplication as the existing work(^2^22Note that there are many fast matrix multiplication methods. We do not study them, as they are hard to be parallelized. Standard matrix multiplication is still the most popular implementation on GPU, e.g., PyTorch. We refer readers to Section 3 for more details.) and focus on backward gradient computation.
We establish matching I/O complexity upper and lower bounds for attention backward computation (formalized in Theorem 1.1 and illustrated in Fig. 1).
Combined with the attention forward results from SY [], this completes the theory of I/O complexity in the attention mechanism.

Our main result is stated as follows:

###### Theorem 1.1 (Main result) .

Let $n$ be the sequence length, $d$ the head dimension, and $M$ the cache size. The I/O complexity of attention backward computation under standard matrix multiplication is

$$ $\displaystyle\Theta\left(\min\left\{\frac{n^{2}d^{2}+nd^{3}}{M},\frac{n^{2}d+ nd^{2}}{\sqrt{M}}\right\}\right).$ $$

To interpret our main result, we categorize the cache size $M$ into two cases: the small cache case where $M=o(d^{2})$ and the large cache case where $M=\Omega(d^{2})$ (see Fig. 1 for illustration).

In the small cache scenario, $M=o(d^{2})$, by computation graph Fig. 2 and Algorithm 6, we show that the upper bound of the I/O complexity is $O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.
In detail, Algorithm 6 explicitly read/write the $n\times n$ attention matrix and other $n\times d$ intermediate matrices from/to memory.
Note that, when $M=o(d^{2})$, our Algorithm 6 has a better upper bound than FlashAttention, whose upper bound is $O(\frac{n^{2}d^{2}+nd^{3}}{M})$.
Furthermore, to establish a lower bound on the I/O complexity, we show that the I/O complexity of attention backward computation is equivalent to the I/O complexity of matrix multiplication when $M=o(d^{2})$, which matches the upper bound of Algorithm 6.

Figure: Figure 1: Attention backward I/O complexity comparison. The $x$-axis is the cache size, and the $y$-axis is the I/O complexity. The red line represents our tight upper/lower bound (Theorem 1.1), and the blue dash denotes the upper bound for FlashAttention []. The cross point is $M=\Theta(d^{2})$, the dividing point of large cache and small cache settings. The results show that FlashAttention is optimal when $M=\Omega(d^{2})$.
Refer to caption: x1.png

In the more practical large cache case, $M=\Omega(d^{2})$, we prove an upper bound $O(\frac{n^{2}d^{2}+nd^{3}}{M})$ on the I/O complexity for the attention backward algorithms (Algorithm 9), which matches that of FlashAttention [, , ]. We prove that this upper bound is tight by providing a matching lower bound for the I/O complexity of attention backward using the red-blue pebble game analysis framework from HK [].

Therefore, we provide the optimal bounds and algorithms for backward passes for all cache sizes. This fully characterizes the I/O complexity of attention forward/backward when combined with existing results on forward passes [].
Notably, we confirm that FlashAttention is optimal for both the forward and backward passes when the cache size is large enough $M=\Omega(d^{2})$.

Moreover, in recent years, sparse attention has become another mainstream method for speeding up the training process of transformer-based models [, , ]. These approaches mainly focus on techniques for sparsifying the attention matrix, thereby reducing the quadratic bottleneck in running time.
However, it remains unknown whether this method can be integrated with I/O-aware algorithms like FlashAttention. Consequently, we further analyze the I/O complexity of sparse attention to provide theoretical guarantees, offering fine-grained lower bounds.

###### Theorem 1.2 (Lower bound for sparse attention forward and backward, informal version of Theorem 4.5 ) .

Let $Z_{\mathrm{input}}$ and $Z_{\mathrm{QK}}$ be the number of nonzero entries of the input matrix and the key-query matrix, respectively. Then any algorithm for both attention forward and backward computation using sparse semi-ring matrix multiplication has I/O complexity

$$ $\displaystyle\Omega\left(\min\left\{\frac{Z_{\mathrm{input}}^{2}}{M},\frac{Z_{ \mathrm{input}}\sqrt{Z_{\mathrm{QK}}}}{\sqrt{M}}\right\}\right).$ $$

Our I/O complexity lower bound for sparse attention recovers the lower bound for both attention forward and backward passes when matrices involved in attention computation are dense, i.e., $Z_{\mathrm{input}}=\Omega(nd),Z_{\mathrm{QK}}=\Omega(n^{2})$. In such case, our lower bound reads as $\Omega(\min\{\frac{n^{2}d^{2}}{M},\frac{n^{2}d}{\sqrt{M}}\})$, matching Theorem 1.1.

**Table 1: Summary of our contributions. We categorize the cache size $M$ into two cases: (1) Large cache $M=\Omega(d^{2})$; (2) Small cache $M=o(d^{2})$. Assume $n\geq d$. We list our contributions for general and sparse attention below. $Z_{\mathrm{input}}$ and $Z_{\mathrm{QK}}$ denote the number of nonzero entries of the input matrix and the key-query matrix, respectively.**
| Attention Algorithm | Large Cache | Reference | Small Cache | Reference |  |
| --- | --- | --- | --- | --- | --- |
| General | Forward Upper | $O(n^{2}d^{2}/M)$ | DFE^+ [] | $O(n^{2}d/\sqrt{M})$ | SY [] |
| Forward Lower | $\Omega(n^{2}d^{2}/M)$ | SY [] | $\Omega(n^{2}d/\sqrt{M})$ | SY [] |  |
| Backward Upper | $O(n^{2}d^{2}/M)$ | DFE^+ [] | $O(n^{2}d/\sqrt{M})$ | Theorem 4.3 |  |
| Backward Lower | $\Omega(n^{2}d^{2}/M)$ | Theorem 4.2 | $\Omega(n^{2}d/\sqrt{M})$ | Theorem 4.4 |  |
| Sparse | Forward Lower | $\Omega(Z_{\mathrm{input}}^{2}/M)$ | Theorem 4.5 | $\Omega(Z_{\mathrm{input}}\sqrt{Z_{\mathrm{QK}}}/\sqrt{M})$ | Theorem 4.5 |
| Backward Lower | $\Omega(Z_{\mathrm{input}}^{2}/M)$ | Theorem 4.5 | $\Omega(Z_{\mathrm{input}}\sqrt{Z_{\mathrm{QK}}}/\sqrt{M})$ | Theorem 4.5 |  |

We summarize our contributions in Table 1 and also conclude as follows:

- •
For small cache sizes $M=o(d^{2})$ in the backward pass, we present optimal upper and lower bounds and propose an algorithm achieving the optimal (Algorithm 6). Notably, FlashAttention is not optimal in this setting, and our algorithm outperforms it.
- •
For large cache sizes $M=\Omega(d^{2})$ in the backward pass, we establish an optimal lower bound that matches the existing upper bound. We also prove the optimal upper bound and introduce an optimal algorithm (Algorithm 9), matching the existing results for FlashAttention but providing a different analysis.
- •
For sparse attention, we offer fine-grained lower bounds for both forward and backward passes and across all cache sizes (Theorem 4.5).

###### Theorem 1.1 (Main result) .

###### Theorem 1.2 (Lower bound for sparse attention forward and backward, informal version of Theorem 4.5 ) .

#### Roadmap.

In Section 2, we review related literature.
In Section 3, we introduce the definitions and background necessary for our study.
We present our main results in Section 4 and discuss the techniques we employed in Section 5.
Section 6 concludes our paper.

## 2 Related Work

#### Attention Computation Acceleration.

The quadratic time complexity of attention computation with respect to the length of the input sequence [] poses significant computational challenges, especially for long sequences. Consequently, accelerating attention computation has become a crucial research area.
From a theoretical standpoint, numerous works focus on approximating the attention matrix to accelerate computation [, , , , , ].
Experimental approaches involve modifying model architectures and optimizing implementations to accelerate inference. Methods such as Mamba [, ], Linearizing Transformers [, ], Hopfield Models [, , , , , , ]
and PolySketchFormer [, ] aim to improve model performance and inference speed.
System-level optimizations, such as FlashAttention [, , ] and block-wise parallel decoding [], address bottlenecks in attention mechanisms and enhance inference speed through efficient implementation strategies. Collectively, these advancements contribute to making attention mechanisms more scalable and efficient, facilitating the deployment of large-scale language models. [] accelerates inference by compressing the input text.

#### Learning with Bounded Memory and I/O Complexity.

A common memory model in computational systems is the two-level memory hierarchy. In this model, there are two layers of memory: a small but fast layer called the *cache*, and a large but slower layer called the *memory*. The I/O (input/output) complexity of an algorithm measures its efficiency based on the number of data transfer operations it performs between the cache and the memory. The early work of HK [] formulated the I/O complexity mathematically using the language of graph theory.
Learning with bounded memory has been studied in various fields in machine learning such as online learning [, , ], convex optimization [, ], active learning [], attention computation [], and continual learning [, ].

#### Sparse Attention.

Over the past few years, there has been extensive research on sparse Transformer/Attention models with weights pruning and inputs pruning, aimed at accelerating computation and training [, , , , , , , , , ]. In practice, the attention matrix is sparse, significantly reducing computational costs. Theoretical studies, such as [], have demonstrated that sparse transformers are expressive enough and can achieve universal approximation properties.

Figure: Figure 2: The computational graph for attention forward and backward. The blue boxes are input matrices, the gray boxes are intermediate matrices, the green box is the forward output, and the orange box is the final gradient matrix. Here, $A_{1},A_{2},A_{3}$ denote the previous inputs, $\mathrm{d}O$ denotes the upstream gradient, and $X,Y$ denote the attention weights. More detailed definitions of each variables can be found in Section 3 and B.
Refer to caption: x2.png

## 3 Preliminary

In this work, we consider using a standard algorithm for matrix multiplication, which means that for any two matrices $A\in\mathbb{R}^{n_{1}\times d},B\in\mathbb{R}^{d\times n_{2}}$, each entry of $AB$ is computed by $(AB)_{i,j}=\sum_{k=1}^{d}A_{i,k}B_{k,j}$ for $i\in[n_{1}],j\in[n_{2}]$. Note that this setting is also used in FlashAttetnion [] and SY [].
Then, we introduce some key concepts needed for this paper.

### 3.1 Key Concept of Attention

Before formally stating our results, we begin by precisely defining the problems we study. We define the following computation of the general Softmax attention forward layer.

###### Definition 3.1 (Attention forward computation) .

Let $n$ be the input length and $d$ be the head dimension.
Let $A_{1},A_{2},A_{3}\in\mathbb{R}^{n\times d}$ be the inputs of previous layer.
Given query, key and value weights matrix $W_{Q},W_{K},W_{V}\in\mathbb{R}^{d\times d}$, we have the Softmax attention forward computation being

$$ $\displaystyle\mathsf{Attn}(A_{1},A_{2},A_{3}):=D^{-1}\exp(A_{1}W_{Q}W_{K}^{ \top}A_{2}^{\top})A_{3}W_{V},$ $$

where (1) $D:=\operatorname{diag}(\exp(A_{1}W_{Q}W_{K}^{\top}A_{2}^{\top})\cdot{\bf 1}_{n})$,
(2) $\exp$ denotes the exponential function and is applied entry-wisely,
(3) $\operatorname{diag}()$ operation takes a vector and outputs a diagonal matrix with the entries of that vector,
and (4) ${\bf 1}_{n}$ denotes the length-$n$ all ones vector.

To simplify and focus more clearly on the core computational aspects of the problem, we set $X=W_{Q}W_{K}^{\top}\in\mathbb{R}^{d\times d}$ and $Y=W_{V}\in\mathbb{R}^{d\times d}$.

Note that, we have $\mathsf{Softmax}(A_{1}XA_{2}^{\top})=D^{-1}\exp(A_{1}XA_{2}^{\top})\in\mathbb{
R}^{n\times n}$, and usually we call it the attention matrix.
The above definition is general and encompasses both self-attention and cross-attention mechanisms in Transformer architectures. Specifically, self-attention occurs when $A_{1}=A_{2}=A_{3}$, meaning that the queries, keys, and values are all derived from the same source. In contrast, cross-attention happens when $A_{2}=A_{3}$, indicating that the keys and values come from one source while the queries come from the other.

Notably, FlashAttention [, , ] and SY [] consider $Q,K,V\in\mathbb{R}^{n\times d}$ after applying the linear layer to the previous inputs, while we consider a more detailed structure as $Q=A_{1}W_{Q},K=A_{2}W_{K},V=A_{3}W_{V}$ (Definition 3.1) explicitly calculating module-wise gradients on attention weights.
This explains why our I/O complexity bound $\Theta(\min\{\frac{n^{2}d^{2}+nd^{3}}{M},\frac{n^{2}d+nd^{2}}{\sqrt{M}}\})$ in Theorem 1.1 has an additional term $nd^{2}$ in the small cache case and $nd^{3}$ in the large cache case. When $n\geq d$, the additional term will disappear.

Mathematically, optimizing the attention computation involves adjusting the attention weight matrices $X$, and $Y$.
Using the previous results on attention gradients from AS24a [] and LSS^+ [], we have the following definition of attention gradient:

###### Definition 3.2 (Attention backward gradient) .

Let $A_{1},A_{2}\in\mathbb{R}^{n\times d}$.
Let $p(X)\in\mathbb{R}^{n\times n}$ be defined in Definition B.9 (see Fig. 2 for an illustration).
Let $L(X)$ be some loss function.
The attention backward gradient for $X\in\mathbb{R}^{d\times d}$ is:

$$ $\displaystyle\frac{\mathrm{d}L(X)}{\mathrm{d}X}=A_{1}^{\top}p(X)A_{2}.$ $$

###### Remark 3.3 .

Since the attention module depends only linearly on $Y$, it is straightforward to incorporate it into an algorithm, and it is not a complexity bottleneck.
Thus, we focus on the case where $X$ is variable and $Y$ is a fixed input.

Figure: Figure 3: This diagram shows a summation tree with $d=2$ in the computational graph for the backward passes of attention using standard matrix multiplication. The orange and green nodes represent the input nodes of the level-$1$ summation tree. The brown nodes, along with the blue nodes (output from the level-$1$ summation tree), serve as inputs for the level-$2$ summation tree. The purple nodes represent the target output. When $d$ gets larger, the summation tree will expand with additional layers, where each new layer introduces intermediate nodes that represent the sums of pairs of nodes from the previous layer, i.e., there will be total $1+\log_{2}d$ layer in total.
Refer to caption: x3.png

###### Definition 3.1 (Attention forward computation) .

###### Definition 3.2 (Attention backward gradient) .

###### Remark 3.3 .

### 3.2 Summation Tree

In this subsection, we need to introduce the computational graph of the attention backward gradient, which is the key concept in our I/O complexity analysis.

In the computational graph shown in Fig. 2, we can first compute $A_{1}X$ and then compute $(A_{1}X)A_{2}^{\top}$, or first compute $XA_{2}^{\top}$ and then compute $A_{1}(XA_{2}^{\top})$. In either case, we perform two matrix multiplications: one between an $n\times d$ matrix and a $d\times d$ matrix, and the other between an $n\times d$ matrix and a $d\times n$ matrix.
Without loss of generality for illustration, we consider the first case. To compute $A_{1}X$, we need to calculate the products $\{(A_{1})_{i,k}X_{k,j}\}$ for all $i\in[n],\,k\in[d],\,j\in[d]$. Each entry $(A_{1}X)_{i,j}$ is then obtained by summing these products over $k$:
$(A_{1}X)_{i,j}=\sum_{k=1}^{d}(A_{1})_{i,k}X_{k,j}.$
In the computational graph, this summation is represented by a summation tree that connects the product nodes $(A_{1})_{i,k}X_{k,j}$ to the sum node $(A_{1}X)_{i,j}$. We define the product nodes $(A_{1})_{i,k}X_{k,j}$, the nodes corresponding to the sums $(A_{1}X)_{i,j}$, and all intermediate nodes in the summation trees as *level-1 nodes*. Similarly, we define *level-2 nodes* as these nodes in the summation trees involved in computing $(A_{1}X)A_{2}^{\top}$. We give an example of the summation tree with $d=2$ in Fig. 3.

### 3.3 I/O Complexity

There are various ways to define the two-level memory hierarchy and the I/O complexity. We state the definition in HK [], which formulates the two-level memory hierarchy as a red-blue pebble game played on a computational graph. Very recently, SY [] proved that the I/O complexity of forward computation of FlashAttention is optimal by analyzing the red-blue pebble game on an attention forward computational graph.

###### Definition 3.4 (Red-blue pebble game [ 49 ] ) .

Consider a game played on a directed acyclic graph that has a limited number of red pebbles and an unlimited number of blue pebbles. Initially, each input node (a node with no parents) is marked with a blue pebble, while all other nodes have no pebbles. The player is allowed to perform the following operations:

- •
Input: Replace a blue pebble on a node with a red pebble.
- •
Output: Replace a red pebble on a node with a blue pebble.
- •
Compute: Place a red pebble on a node if all its parent nodes have red pebbles.
- •
Delete: Remove a pebble from a node.

The objective of the game is to place blue pebbles on all output nodes (i.e., nodes with no children) while minimizing the total number of input and output operations used throughout the process.

In the red-blue pebble game, each node represents a computational task. A red pebble denotes a unit in the small but fast layer known as *cache*, while a blue pebble represents a unit in the large but slower layer called *memory*. A task can only be computed once all its dependent tasks are completed. All computations are assumed to occur within the cache. Hence, efficient use of cache plays a critical role in reducing the I/O operations of an algorithm to minimize the cost associated with data transfer between memory and cache. We can define the I/O complexity by using the red-blue pebble game.

###### Definition 3.5 (I/O complexity [ 49 ] ) .

Consider the red-blue pebble game played on a directed acyclic graph $G$. Let $M$ be a positive integer. The I/O complexity, denoted as $Q(G,M)$, is the minimum number of input and output operations to complete the objective of the game with the restriction that no more than $M$ red pebbles are present on the graph at any time. We omit $G$ when it is clear in the context.

The red-blue pebble game provides insight into cache management by modeling the limited cache size through the number of red pebbles. The maximum number of red pebbles corresponds to the size of the cache, which means that there can be at most $M$ items in the cache at any given time.

###### Definition 3.4 (Red-blue pebble game [ 49 ] ) .

###### Definition 3.5 (I/O complexity [ 49 ] ) .

## 4 Main Results

In Theorem 1.1, we provide matching upper and lower bounds for the I/O complexity of attention gradient computation in the backward passes.
In detail, Theorem 1.1 states that the I/O complexity of the attention gradient computation is $\Theta(\min\{\frac{n^{2}d^{2}+nd^{3}}{M},\frac{n^{2}d+nd^{2}}{\sqrt{M}}\})$, which splits the cache size into two cases: (1) small cache $M=o(d^{2})$; (2) large cache $M=\Omega(d^{2})$.
At the cross point $M=d^{2}$, we have $\frac{n^{2}d^{2}+nd^{3}}{M}=\frac{n^{2}d+nd^{2}}{\sqrt{M}}=n^{2}+nd$.
An intuitive figure of the asymptotic I/O complexity is shown in Fig. 1.

Here we discuss two implications of Theorem 1.1.
First, through the fine-grained analysis, our result identifies a critical point at $M=d^{2}$, where the I/O complexity changes its behavior.
For $M=o(d^{2})$, we establish better upper and lower bounds compared to existing results, demonstrating that FlashAttention is not optimal in this regime.
Second, when $M=\Omega(d^{2})$, Theorem 1.1 provides a tighter lower bound than existing work using red-blue pebble game (Definition 3.4), offering insights of algorithm design.

Second, by combining the results of SY [] with our findings, we provide a more general and tighter I/O complexity characterization of FlashAttention 1/2 [, ].
In the large cache scenario where $M=\Omega(d^{2})$, the attention forward I/O complexity is $\Theta(\frac{n^{2}d^{2}}{M})$, as discussed in Theorem 5.1 of SY [].
Combining this result with our attention backward I/O complexity $\Theta(\frac{n^{2}d^{2}+nd^{3}}{M})$ (Theorem 1.1), we conclude that the overall complexity is $\Theta(\frac{n^{2}d^{2}+nd^{3}}{M})$.
Thus, given the cache size is sufficiently large, i.e. $M=\Omega(d^{2})$, the I/O complexity of the forward and backward computation for FlashAttention 1/2 is optimal.

Our main result Theorem 1.1 is a summary of our results for different cache sizes (Theorem 4.1, 4.2, 4.3, and 4.4), which will be discussed in the later subsections.

### 4.1 Large Cache

The large cache scenario is more interesting and practical.
We now prove an upper bound below.

###### Theorem 4.1 (Large cache upper bound, informal version of Theorem D.5 ) .

Suppose $n$ is the input length, $d$ is the head dimension, and $M=\Omega(d^{2})$ is the cache size.
There is an algorithm (see Algorithm 9) outputs a $d\times d$ matrix $g=\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ (Definition 3.2) with I/O complexity $O(\frac{n^{2}d^{2}+nd^{3}}{M})$.

We then demonstrate that this upper bound is tight by providing a matching lower bound for the I/O complexity of the attention backward passes.
To achieve this, we employ the framework developed in HK [], which shows that executing an algorithm on a machine with a two-level memory hierarchy can be modeled by a red-blue pebble game (Definition 3.4) on a directed acyclic graph.
We present the large cache lower bound below, which shows as long as the cache size $M=\Omega(d^{2})$, the I/O complexity is at least $\Omega(\frac{n^{2}d^{2}+nd^{3}}{M})$.

###### Theorem 4.2 (Large cache lower bound, informal version of Theorem E.9 ) .

Suppose $n$ is the input length and $d$ is the head dimension.
Suppose the cache size $M=\Omega(d^{2})$.
Then the I/O complexity of attention gradient computation using standard matrix multiplication is always $\Omega(\frac{n^{2}d^{2}+nd^{3}}{M})$.

###### Theorem 4.1 (Large cache upper bound, informal version of Theorem D.5 ) .

###### Theorem 4.2 (Large cache lower bound, informal version of Theorem E.9 ) .

### 4.2 Small Cache

In the small cache case, we provide an upper bound below. Notice that this is better than the I/O complexity of FlashAttention which is $O(\frac{n^{2}d^{2}+nd^{3}}{M})>O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$ when $M=o(d^{2})$.

###### Theorem 4.3 (Small cache upper bound, informal version of Theorem C.12 ) .

Suppose $n$ is the input length, $d$ is the head dimension, and $M=o(d^{2})$ is the cache size.
There is an algorithm (see Algorithm 6) outputs a $d\times d$ matrix $g=\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ (Definition 3.2) with I/O complexity $O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$, time complexity $O(n^{2}d+nd^{2})$, and space complexity $O(n^{2}+d^{2})$.

Furthermore, we show that attention gradient computation can be reduced to matrix multiplication, establishing a matching lower bound.

###### Theorem 4.4 (Small cache lower bound, informal version of Theorem E.10 ) .

Suppose $n$ is the input length and $d$ is the head dimension.
Suppose the cache size $M=o(d^{2})$.
Then the I/O complexity of attention gradient computation using standard matrix multiplication is always $\Omega(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.

###### Theorem 4.3 (Small cache upper bound, informal version of Theorem C.12 ) .

###### Theorem 4.4 (Small cache lower bound, informal version of Theorem E.10 ) .

### 4.3 Lower Bound of Sparse Attention Forward and Backward Passes

Sparse attention is a generalization of standard attention and has been popular in practical applications. We refer readers to Section 2 for more discussion. To state our results, we first introduce some notations. For any matrix $A$, we use
$\operatorname{nnz}(A)$ to denote the number of non-zero entries in the matrix $A$. We assume that sparse matrices are stored by listing only their non-zero entries along with their coordinates. We assume sparse semi-ring matrix multiplication, which restricts operations to addition and multiplication of these entries. Each output entry $(AB)_{i,j}$ can only be computed as the sum of products given by $\sum_{k}A_{i,k}B_{k,j}$.

###### Theorem 4.5 (Lower bound for sparse attention forward and backward, formal version of Theorem 1.2 ) .

Suppose $n$ is the input length, $d$ is the head dimension, and $M$ is the cache size. Let $Z_{A}:=\min\{\operatorname{nnz}(A_{1}),\operatorname{nnz}(A_{2})\},Z_{X}:=
\operatorname{nnz}(X),Z_{AX}=\min\{\operatorname{nnz}(A_{1}X),\operatorname{
nnz}(XA_{2}^{\top})\},Z_{AXA}:=\operatorname{nnz}(A_{1}XA_{2}^{\top})$. Then any algorithm for both attention forward and backward computation using sparse semi-ring matrix multiplication has I/O complexity

$$ $\displaystyle\Omega\left(\min\left\{\frac{Z_{A}^{2}+Z_{A}Z_{X}}{M},\frac{Z_{A} \sqrt{Z_{AXA}}+\sqrt{Z_{A}Z_{X}Z_{AXA}}}{\sqrt{M}}\right\}\right).$ $$

###### Remark 4.6 .

When matrices involved in attention computation are dense, i.e., $Z_{A}=\Omega(nd),Z_{X}=\Omega(d^{2}),Z_{AX}=\Omega(nd)$, and $Z_{AXA}=\Omega(n^{2})$. In such case, our lower bound reads as $\Omega(\min\{\frac{n^{2}d^{2}+nd^{3}}{M},$ $\frac{n^{2}d+nd^{2}}{\sqrt{M}}\})$. Hence, it matches the result of lower bounds in the dense case.

###### Theorem 4.5 (Lower bound for sparse attention forward and backward, formal version of Theorem 1.2 ) .

###### Remark 4.6 .

## 5 Technical Overview

Upper Bound of Small Cache.
In Section C, we present algorithms for the backward passes of attention in the small cache case, where $M=o(d^{2})$. We observe that when $M=o(d^{2})$, we have $\frac{n^{2}d^{2}+nd^{3}}{M}>\frac{n^{2}d+nd^{2}}{\sqrt{M}}>n^{2}+nd$. Then we can exploit this to design a better algorithm with I/O complexity better than $\frac{n^{2}d^{2}+nd^{3}}{M}$, by reading/writing the $n\times n$ attention matrix and other $n\times d$ intermediate matrices from/to memory. In detail, our small cache algorithm (Algorithm 6) follows the computational graph in Figure 2 and is divided into four phases. In Phase 1 (Algorithm 2), we compute the attention matrix $f$ (Definition B.5) and write it to memory. In Phase 2 (Algorithm 3), we compute $q$ (Definition B.8), incorporating the information from the upstream gradient $\mathrm{d}O$. Phase 3 (Algorithm 4) computes the gradient component matrix $p$ (Definition B.9). Finally, in Phase 4 (Algorithm 5), we compute the final gradient $g=A_{1}^{\top}pA_{2}$ (Definition 3.2).
At a high level, our algorithm splits the input and output matrices into blocks of size $\sqrt{M}\times\sqrt{M}$.
On the other hand, FlashAttention divides the $n\times d$ input matrices into multiple $k\times d$ matrices, where $k<n$.
Compared to our upper bound, we can see that FlashAttention is not optimal in this case.
Following the computational graph in Figure 2, we perform the backward passes of attention using each $\sqrt{M}\times\sqrt{M}$ block as basic elements in standard matrix multiplication.
Compared to forward passes, the computational graph of backward passes is more complicated and requires more fine-grained analysis, e.g., the four phases mentioned above.
Through a detailed analysis of Algorithm 6, we establish Theorem 4.3.

Upper Bound of Large Cache.
In Section D, we present algorithms for attention backward in the large cache case, where $M=\Omega(d^{2})$. Similar to FlashAttention, the $n\times n$ attention matrix $f$ (Definition B.5) cannot be directly loaded into cache, even though it has been computed and can be stored in memory.
The overall algorithm (Algorithm 9) consists of two phases. In Phase 1 (Algorithm 7), we compute $S=A_{1}X$ and $h=A_{3}Y$, and these two matrices are then passed to Phase 2. In Phase 2 (Algorithm 8), the inputs are matrices $A_{1},A_{2},S,h,O,\mathrm{d}O\in\mathbb{R}^{n\times d}$ (Definitions 3.1, B.6, B.7, and B.8), and vector $l\in\mathbb{R}^{n}$ (Definition B.4). We vertically divide the inputs into row block matrices of size $B_{r}\times d$ or $B_{c}\times d$, where $B_{r}=\min\{\lceil M/4d\rceil,d\}$ and $B_{c}=\lceil M/4d\rceil$. Using these row block matrices as computation units, we follow the computational graph (Fig. 2) and FlashAttention’s procedure.
After accounting for the reads and writes of the overall algorithm (Algorithm 9), we prove Theorem 4.1.
Furthermore, when the cache size is as large as $\Theta(nd)$,
the I/O complexity can be reduced to $O(nd+d^{2})$, which corresponds to the size of the input and output of the algorithm.

Lower Bound of Large Cache and Small Cache.
In Section E, we establish the lower bounds for the I/O complexity of attention gradient computation in both large and small cache cases. Following Definitions 3.4 and 3.5, we analyze the red-blue pebble game on the computational graph of any attention backward algorithm using standard matrix multiplication. More specifically, the key concept is the $M$-partition, which decomposes the graph into subgraphs, ensuring that each subgraph satisfies conditions related to dominator and minimum sets (Definitions E.1, E.2, E.3, E.4, and E.5). Our proofs for the lower bound of backward passes builds upon the lemmas (Lemmas E.7 and E.8), which provide the foundation for relating the number of subgraphs to the I/O operations required.
For the large cache scenario, $M=\Omega(d^{2})$, we demonstrate that the I/O complexity scales with the need to compute matrix products efficiently. In the small cache case, $M=o(d^{2})$, we show that higher I/O complexity is unavoidable due to the data transfers between cache and memory by reducing to the standard matrix multiplication. These analyses are formally established in the proofs of Theorems E.9 and E.10. In particular, our Theorems E.10, the small cache lower bound case, requires a new analysis deviation.

###### Remark 5.1 .

The Softmax in Definition 3.1 can be changed to other non-linear activation functions and our lower bound still holds.
It is because we must compute matrix multiplication of size $n\times d$ and $d\times n$ in non-linear attention.
However, for linear attention, i.e., $A_{1}XA_{2}^{\top}A_{3}Y$, our lower bound is loose, since we can compute $\underbrace{A_{2}^{\top}}_{d\times n}\underbrace{A_{3}}_{n\times d}$ first, and then we have $\underbrace{A_{1}}_{n\times d}\underbrace{X}_{d\times d}\underbrace{A_{2}^{
\top}A_{3}}_{d\times d}\underbrace{Y}_{d\times d}$.

Lower Bound of Sparse Attention Forward and Backward Passes.
In Section F, we establish lower bounds on the I/O complexity of sparse attention computation for both forward and backward passes. Sparse matrix multiplication is considered, where only non-zero entries are stored and used in computations. We derive I/O complexity bounds based on the non-zero counts of input matrices and the I/O operations required for sparse matrix multiplication (Lemma F.1). We further extend these bounds to the matrix products involved in the attention mechanism (Lemma F.2), which requires multiple sparse matrix multiplication analysis. We analyze scenarios where matrices are stored in cache or require intermediate I/Os during computation to obtain the I/O complexity bounds for both forward and backward passes (Theorems F.3 and Theorem F.4), and Theorem 4.5 directly holds as a consequence.

###### Remark 5.1 .

## 6 Conclusion

This work provided a comprehensive analysis of the I/O complexity for attention mechanisms, focusing on backward passes. We established tight bounds on I/O complexity for both small and large caches. Our results confirm that FlashAttention is optimal for both forward and backward on large cache sizes. For small cache sizes, we provided improved upper and lower bounds compared to existing methods.
Additionally, we derived lower bounds for sparse attention for both forward and backward and across cache sizes. Our findings complete the theoretical foundation for I/O complexity in attention mechanisms, offering insights for efficient LLM training and inference. We leave exploring practical implementations leveraging these theoretical insights and investigating I/O complexity for other emerging attention variants as our future work.

## Acknowledgement

Research is partially supported by the National Science Foundation (NSF) Grants 2023239-DMS, CCF-2046710, and Air Force Grant FA9550-18-1-0166.

## Appendix A More Related Work

#### Large Language Models.

The exceptional success of generative large language models (LLMs), such as GPT-4 [], Claude 3 [], Gemini 1.5 [], Llama 3.1 [], Mistral Nemo [], Phi 3.5 [], is fundamentally attributed to the transformer architecture introduced by VSP^+ [] and all support at least 128k input token length.
The transformer architecture and its self-attention mechanism have become indispensable in leading natural language processing (NLP) models [], demonstrating remarkable capabilities across a diverse array of applications, including language translation [], sentiment analysis [], language modeling [], the integration of differential privacy [, ], and multi-modal tasks [, , ].
Transformers’ emergent compositional abilities [, ] and proficiency in in-context learning [, , ] have led some to consider them as early indicators of Artificial General Intelligence (AGI) [].
As such, the transformer architecture continues to play a pivotal role in advancing the field of AI.

#### More about Attention Computation Acceleration.

The quadratic time complexity of attention computation with respect to the length of the input sequence [] poses significant computational challenges, especially for long sequences. Consequently, accelerating attention computation has become a crucial research area, with approaches broadly divided into two categories: (1) theoretical optimization of computational complexity [, ], and (2) experimental improvements to model performance [, , , , ].

From a theoretical standpoint, numerous works focus on approximating the attention matrix to accelerate computation. For example, AS [], AS24a [] utilize polynomial kernel approximation techniques [] to speed up both training and inference of a single attention layer, achieving almost linear time complexity, and extend this approach to multi-layer transformer [] and tensor attention [, ]. Other theoretical contributions include the conv-basis method introduced by LLS^+24a [] and a near-linear time algorithm proposed by HJK^+ [] under the assumptions of uniform softmax column norms and sparsity.

Experimental approaches involve modifying model architectures and optimizing implementations to accelerate inference. Methods such as Mamba [, ], Linearizing Transformers [, ], PolySketchFormer [, ], and various implementations of the Hopfield Model [, , , , , , ] aim to improve model performance and inference speed.
Additionally, specific techniques like
weight pruning [, ] have been developed to accelerate LLM generation.
Some other techniques are introduced for efficient adaptation, such as
LoRA [, , ] and prefix turning [, ].
System-level optimizations, such as Flash Attention [, , ] and block-wise parallel decoding [], address bottlenecks in attention mechanisms and enhance inference speed through efficient implementation strategies.
Collectively, these advancements contribute to making attention mechanisms more scalable and efficient, facilitating the deployment of large-scale language models.

#### More about Learning with Bounded Memory and I/O Complexity.

Learning with bounded memory has been studied in various fields in machine learning such as online learning [, , , ], parity learning [, , , ], convex optimization [, , ], active learning [], learning linear classifiers [], attention computation [], linear regression [, , ], linear programming [, ], semi-definite programming [], principal component analysis [], continual learning [, ], entropy estimation [, ] and others [, ].

A common memory model in computational systems is the two-level memory hierarchy. In this model, there are two layers of memory: a small but fast layer called the *cache*, and a large but slower layer called the *memory*. The I/O (input/output) complexity of an algorithm measures its efficiency based on the number of data transfer operations it performs between the cache and the memory. In domains such as big data analytics and database management, these data transfers can become significant performance bottlenecks because massive datasets cannot be entirely accommodated in the cache, and thus optimizing I/O is essential for fast data retrieval and storage, directly impacting query performance and system scalability [, ]. The early work of HK [] formulated the I/O complexity mathematically using the language of graph theory. Vit [] provides a comprehensive survey of the I/O complexity of various batched and online problems. There exists a substantial body of work on the I/O complexity of numerous problems, including sorting [], graph algorithms [, , , ], fine-grained I/O complexity [], computational trade-off in data transfers [], computing prime tables [], attention computation [], integer multiplication [, ], and matrix multiplication [, ].

## Appendix B Preliminary

In Section B.1, we define some basic notation we will use.
In Section B.2, we introduce the memory hierarchy we consider.
In Section B.3, we state important facts related to fast matrix multiplication.
In Section B.4, we define several intermediate functions which will arise in our algorithms.

### B.1 Notations

For any positive integer $n$, we define $[n]:=\{1,2,\dots,n\}$.
For two same length vector $x$ and $y$, we use $\langle x,y\rangle$ to denote the inner product between $x$ and $y$, i.e., $\langle x,y\rangle=\sum_{i=1}^{n}x_{i}y_{i}$.
We use $\circ$ to denote the Hadamard product i.e. the $(i,j)$-entry of $A\circ B$ is $A_{i,j}B_{i,j}$.
We use $x\circ y$ to denote vector that $i$-th entry is $x_{i}y_{i}$. Let ${\bf 1}_{n}$ denote the length-$n$ all ones vector.
It is not hard to see that $\langle x\circ y,{\bf 1}_{n}\rangle=\langle x,y\rangle$.
For a vector $x$, we use $x^{\top}$ to denote the transpose of $x$.
For a matrix $A$, we use $A^{\top}$ to denote the transpose of matrix $A$.
For a matrix $A$, we use $\exp(A)$ to denote the matrix that $(i,j)$-th coordinate is $\exp(A_{i,j})$.

Given a matrix $A\in\mathbb{R}^{n\times m}$, we index an individual entry as $A[i,j]$.
The $i$-th row is denoted $A[i]$ while the $j$-th column is denoted $A[*,j]$.
$A[i_{1}:i_{2},j_{1}:j_{2}]$ denotes a block of $A$ consisting of entries $(i,j)$ where $i\in[i_{1},i_{2}]$ and $j\in[j_{1},j_{2}]$.
Given a block size $B$, the block $A[(i-1)\cdot B+1:i\cdot B,(j-1)\cdot B+1:j\cdot B]$ is denoted $A^{(B)}[i,j]$.

For a vector $v\in\mathbb{R}^{n}$, we similarly denote entries $v[i]$, a contiguous block of entries as $v[i_{1}:i_{2}]$, and the $i$-th block of size $B$ as $v^{(B)}[i]$.
Let $\operatorname{diag}(v)$ denote the matrix $D\in\mathbb{R}^{n\times n}$ with $D[i,i]=v[i]$.

### B.2 Memory Hierarchy

In this study, we consider a two-level memory hierarchy composed of a small but fast layer called the *cache* and a large, slower layer referred to as the *memory*. We assume that the memory has unlimited capacity, while the cache is constrained by a finite size $M$. Moreover, all computations are performed exclusively within the cache.

### B.3 Matrix Multiplication

We define matrix multiplication notation and state some well-known facts here.

###### Definition B.1 .

Let $n_{1},n_{2},n_{3}$, denote any three positive integers. We use ${\cal T}_{\mathrm{mat}}(n_{1},n_{2},n_{3})$ to denote the time of multiplying an $n_{1}\times n_{2}$ matrix with another $n_{2}\times n_{3}$.

Then, we introduce a well-known fact.

###### Fact B.2 .

Let $n_{1},n_{2},n_{3}$, denote any three positive integers.
${\cal T}_{\mathrm{mat}}(n_{1},n_{2},n_{3})=O({\cal T}_{\mathrm{mat}}(n_{1},n_{
3},n_{2}))=O({\cal T}_{\mathrm{mat}}(n_{2},n_{1},n_{3}))=O({\cal T}_{\mathrm{
mat}}(n_{2},n_{3},n_{1}))=O({\cal T}_{\mathrm{mat}}(n_{3},n_{1},n_{2}))=O({
\cal T}_{\mathrm{mat}}(n_{3},n_{2},n_{1}))$.

###### Definition B.1 .

###### Fact B.2 .

### B.4 Definitions of Intermediate Variables

We start by some definitions about $X\in\mathbb{R}^{d\times d}$.

###### Definition B.3 (Definition 3.4 in AS24a [ 9 ] ) .

Let $A_{1},A_{2}\in\mathbb{R}^{n\times d}$ be two matrices.
Let $X\in\mathbb{R}^{d\times d}$.

Let us define function $A(X)$ to be:

$$ $\displaystyle A(X):=\underbrace{\exp(A_{1}XA_{2}^{\top})}_{n\times n}.$ $$

###### Definition B.4 (Definition 3.5 in AS24a [ 9 ] ) .

For $A(X)\in\mathbb{R}^{n\times n}$ defined in Definition B.3, we define the softmax normalizing vector $l(X)\in\mathbb{R}^{n}$ to be

$$ $\displaystyle l(X):=\underbrace{A(X)}_{n\times n}\cdot\underbrace{{\bf 1}_{n}} _{n\times 1}.$ $$

###### Definition B.5 (Definition 3.6 in AS24a [ 9 ] ) .

Suppose that $l(X)\in\mathbb{R}^{n}$ is defined as in Definition B.4.
Let $A(X)\in\mathbb{R}^{n\times n}$ be defined as in Definition B.3.
For a fixed $j_{0}\in[n]$, let us consider $f(X)_{j_{0}}$

$$ $\displaystyle f(X)_{j_{0}}:=\underbrace{l(X)_{j_{0}}^{-1}}_{\mathrm{scalar}} \underbrace{A(X)_{j_{0}}}_{n\times 1}.$ $$

Let $f(X)\in\mathbb{R}^{n\times n}$ denote the matrix where $j_{0}$-th row is $(f(X)_{j_{0}})^{\top}$.

Furthermore, the matrix form of $f(X)$ is

$$ $\displaystyle f(X)=\operatorname{diag}(l(X))A(X)$ $$

We then define $h(Y)$ related to $Y\in\mathbb{R}^{d\times d}$.

###### Definition B.6 (Definition 3.7 in AS24a [ 9 ] ) .

For $A_{3}\in\mathbb{R}^{n\times d}$ and $Y\in\mathbb{R}^{d\times d}$, we define $h(Y)\in\mathbb{R}^{n\times d}$ as:

$$ $\displaystyle h(Y):=\underbrace{A_{3}}_{n\times d}\underbrace{Y}_{d\times d}.$ $$

Let us define the forward output matrix $O$.

###### Definition B.7 .

Let $f(X),h(Y)$ be defined in Definition B.5 and B.6.
We define the output of attention as:

$$ $\displaystyle O:=\underbrace{f(X)}_{n\times n}\underbrace{h(Y)}_{n\times d}$ $$

where $O\in\mathbb{R}^{n\times d}$ is the output matrix of attention forward computation.

Now, we define $q$, which incorporates the information from upstream gradient.

###### Definition B.8 (Definition C.10 in LSS + [ 68 ] ) .

Let $\mathrm{d}O\in\mathbb{R}^{n\times d}$ be the upstream gradient, the matrix resulting from the application of the chain rule.
Define $h(Y)\in\mathbb{R}^{n\times d}$ as in Definition B.6.

We define $q(Y)\in\mathbb{R}^{n\times n}$ as

$$ $\displaystyle q(Y):=\underbrace{\mathrm{d}O}_{n\times d}\underbrace{h(Y)^{\top }}_{d\times n}$ $$

Then we use $q(Y)_{j_{0}}^{\top}$ to denote the $j_{0}$-th row of $q(Y)\in\mathbb{R}^{n\times n}$.

Finally, we define the gradient component matrix $p$.

###### Definition B.9 (Definition C.5 in AS24a [ 9 ] ) .

For every index $j_{0}\in[n]$, we define $p(X)_{j_{0}}\in\mathbb{R}^{n}$ as

$$ $\displaystyle p(X)_{j_{0}}:=(\operatorname{diag}(f(X)_{j_{0}})-f(X)_{j_{0}}f(X )_{j_{0}}^{\top})q(Y)_{j_{0}}.$ $$

We define $p(X)\in\mathbb{R}^{n\times n}$ in the sense that $p(X)_{j_{0}}^{\top}$ is the $j_{0}$-th row of $p(X)$.
Additionally, $p(X)$ has matrix form as

$$ $\displaystyle p(X)=$ $\displaystyle~{}f(X)\circ q(Y)-\operatorname{diag}((f(X)\circ q(Y))\cdot{\bf 1 }_{n})f(X)$ $\displaystyle=$ $\displaystyle~{}f(X)\circ q(Y)-\operatorname{diag}((O\circ\mathrm{d}O)\cdot{ \bf 1}_{n})f(X)$ $$

where $f(X),O$ are defined in Definition B.5 and B.7, and $q(Y),\mathrm{d}O$ are defined in Definition B.8.

###### Definition B.3 (Definition 3.4 in AS24a [ 9 ] ) .

###### Definition B.4 (Definition 3.5 in AS24a [ 9 ] ) .

###### Definition B.5 (Definition 3.6 in AS24a [ 9 ] ) .

###### Definition B.6 (Definition 3.7 in AS24a [ 9 ] ) .

###### Definition B.7 .

###### Definition B.8 (Definition C.10 in LSS + [ 68 ] ) .

###### Definition B.9 (Definition C.5 in AS24a [ 9 ] ) .

## Appendix C I/O Complexity Upper Bound for Small Cache

In this section, we prove the I/O complexity upper bound (Theorem C.12) for small cache case $M=o(d^{2})$.
Specifically, in Section C.1, we introduce an algorithm of attention gradient computation without cache to guide our algorithm design.
Section C.2 presents algorithms and analyses for attention gradient computation in the small cache setting.
Finally, Section C.3 provides the upper bound theorem for the small cache case.

### C.1 Algorithm for Attention Backward Without Cache

Using results from AS24a [], we can compute the gradient in ${\cal T}_{\mathrm{mat}}(n,d,n)+{\cal T}_{\mathrm{mat}}(n,d,d)$ time.

###### Lemma C.1 (Attention gradient computation, Lemma C.8 in AS24a [ 9 ] ) .

If it holds that

- •
Define $A_{1},A_{2},A_{3},\mathrm{d}O\in\mathbb{R}^{n\times d}$. Define $X,Y\in\mathbb{R}^{d\times d}$ to be several input fixed matrices.
- •
Let $X,Y\in\mathbb{R}^{d\times d}$ denote matrix variables (we will compute gradient with respect to $X$).
- •
Let $g=\frac{\mathrm{d}L(X)}{\mathrm{d}X}\in\mathbb{R}^{d\times d}$ (Definition 3.2).

Then, gradient $g\in\mathbb{R}^{d\times d}$ can be computed in ${\cal T}_{\mathrm{mat}}(n,d,n)+{\cal T}_{\mathrm{mat}}(n,d,d)$ time.

We first give a naive algorithm that have not utilized cache to compute the gradient (Algorithm 1).

Figure: Algorithm 1 Attention gradient computation without cache. See more details in Section B and C of AS24a [] and Section F of LSS^+ [].

###### Lemma C.2 (Correctness) .

The AttentionGradientNoCache (Algorithm 1) outputs a $d\times d$ matrix $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ defined in Definition 3.2.

###### Proof.

From Lemma C.1, we know this holds.
∎

###### Lemma C.3 (Time/space complexity) .

There exists an algorithm (see Algorithm 1) that can compute the exact gradient in Definition 3.2 in ${\cal T}_{\mathrm{mat}}(n,d,n)+{\cal T}_{\mathrm{mat}}(n,d,d)$ time and $O(n^{2}+d^{2})$ space.

###### Proof.

From Lemma C.1, we can prove the time complexity. Since the stored matrices have three sizes, namely $n\times d$, $n\times n$, $d\times d$, the space complexity is $O(n^{2}+nd+d^{2})=O(n^{2}+d^{2})$.
∎

###### Lemma C.1 (Attention gradient computation, Lemma C.8 in AS24a [ 9 ] ) .

###### Lemma C.2 (Correctness) .

###### Proof.

###### Lemma C.3 (Time/space complexity) .

###### Proof.

### C.2 Algorithms for Attention Backward in Small Cache

We now give algorithms to compute the upper bound of small cache case $M=o(d^{2})$ in attention backward computation.

First, we give the algorithm and analysis for Phase 1 (see Algorithm 2) to compute $f$ defined in Definition B.5.

###### Lemma C.4 (Correctness of Phase 1) .

The AttentionGradientCachePhase1 (Algorithm 2) outputs a $n\times n$ matrix $f$ defined in Definition B.5.

###### Proof.

The algorithm first computes $S=A_{1}X$. Then it computes $A=SA_{2}^{\top}$, $A=\exp(A)$, and $l=A\cdot{\bf 1}$. Finally, it outputs $f=\operatorname{diag}(l)^{-1}A$ which is $f$ defined in Definition B.5.
∎

###### Lemma C.5 (I/O complexity of Phase 1) .

The I/O complexity of AttentionGradientCachePhase1 (Algorithm 2) is $O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.

###### Proof.

In Phase 1 (Algorithm 2) the number of items in cache is at most $3B^{2}+B\leq 4B^{2}\leq M$. For each iteration in computing $S=A_{1}X$ and $A=SA_{2}^{\top}$, the algorithm reads $O(B^{2})$ from memory into cache. This is the dominating factor of the I/O complexity of the algorithm.
Thus, the I/O complexity of Phase 1 is $O(\frac{n^{2}d}{B^{3}}B^{2})+O(\frac{nd^{2}}{B^{3}}B^{2})=O(\frac{n^{2}d+nd^{2
}}{B})=O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.
∎

Figure: Algorithm 2 Attention gradient computation with cache phase 1. Compute $f$.

Second, we give the algorithm and analysis for Phase 2 (see Algorithm 3) to compute $q$ defined in Definition B.8.

###### Lemma C.6 (Correctness of Phase 2) .

The AttentionGradientCachePhase2 (Algorithm 3) outputs a $n\times n$ matrix $q$ defined in Definition B.8.

###### Proof.

The algorithm first computes $h=A_{3}Y$. Then, it outputs $q=\mathrm{d}Oh^{\top}$ which is exactly the same as $q$ defined in Definition B.8.
∎

###### Lemma C.7 (I/O complexity of Phase 2) .

The I/O complexity of AttentionGradientCachePhase2 (Algorithm 3) is $O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.

###### Proof.

In Phase 2 (Algorithm 3) the number of items in cache is at most $3B^{2}\leq 4B^{2}\leq M$. For each iteration in computing $h=A_{3}Y$ and $q=\mathrm{d}Oh^{\top}$, the algorithm reads $O(B^{2})$ from memory into cache. This is the dominating factor of the I/O complexity of the algorithm.
Thus, the I/O complexity of Phase 2 is $O(\frac{n^{2}d}{B^{3}}B^{2})+O(\frac{nd^{2}}{B^{3}}B^{2})=O(\frac{n^{2}d+nd^{2
}}{B})=O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.
∎

Figure: Algorithm 3 Attention gradient computation with cache phase 2. Compute $q$.

Then, we give the algorithm and analysis for Phase 3 (see Algorithm 4) to compute $p$ defined in Definition B.9.

###### Lemma C.8 (Correctness of Phase 3) .

The AttentionGradientCachePhase3 (Algorithm 4) outputs a $n\times n$ matrix $p$ defined in Definition B.9.

###### Proof.

The algorithm first computes $v=(f\circ q)\cdot{\bf 1}$. Then it outputs $p=f\circ q-\operatorname{diag}(v)f$.
∎

###### Lemma C.9 (I/O complexity of Phase 3) .

The I/O complexity of AttentionGradientCachePhase3 (Algorithm 4) is $O(\frac{n^{2}}{\sqrt{M}})$.

###### Proof.

In Phase 3 (Algorithm 4) the number of items in cache is at most $3B^{2}+B\leq 4B^{2}\leq M$. For each iteration in computing $v=(f\circ q)\cdot{\bf 1}$ and $p=f\circ q-\operatorname{diag}(v)f$. The algorithm reads $O(B^{2})$ from memory into cache. This is the dominating factor of the I/O complexity of the algorithm.
Thus, the I/O complexity of Phase 2 is $O(\frac{n^{2}}{B^{3}}B^{2})=O(\frac{n^{2}}{B})=O(\frac{n^{2}}{\sqrt{M}})$.
∎

Figure: Algorithm 4 Attention gradient computation with cache phase 3. Compute $p$.

Lastly, we give the algorithm and analysis for Phase 4 (see Algorithm 5) to compute $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$.

###### Lemma C.10 (Correctness of Phase 4) .

The AttentionGradientCachePhase4 (Algorithm 5) outputs a $d\times d$ matrix $g=\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ (Definition 3.2).

###### Proof.

The algorithm first computes $T=A_{1}^{\top}p$. Then it outputs $g=TA_{2}$.
∎

###### Lemma C.11 (I/O complexity of Phase 4) .

The I/O complexity of AttentionGradientCachePhase4 (Algorithm 5) is $O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.

###### Proof.

In Phase 4 (Algorithm 5) the number of items in cache is at most $3B^{2}\leq 4B^{2}\leq M$. For each iteration in computing $T=A_{1}^{\top}p$ and $g=TA_{2}$. The algorithm reads $O(B^{2})$ from memory into cache. This is the dominating factor of the I/O complexity of the algorithm.
Thus, the I/O complexity of Phase 2 is $O(\frac{n^{2}d}{B^{3}}B^{2})+O(\frac{nd^{2}}{B^{3}}B^{2})=O(\frac{n^{2}d+nd^{2
}}{B})=O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.
∎

Figure: Algorithm 5 Attention gradient computation with cache phase 4. Compute $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$.

###### Lemma C.4 (Correctness of Phase 1) .

###### Proof.

###### Lemma C.5 (I/O complexity of Phase 1) .

###### Proof.

###### Lemma C.6 (Correctness of Phase 2) .

###### Proof.

###### Lemma C.7 (I/O complexity of Phase 2) .

###### Proof.

###### Lemma C.8 (Correctness of Phase 3) .

###### Proof.

###### Lemma C.9 (I/O complexity of Phase 3) .

###### Proof.

###### Lemma C.10 (Correctness of Phase 4) .

###### Proof.

###### Lemma C.11 (I/O complexity of Phase 4) .

###### Proof.

### C.3 Upper Bound for Attention Backward in Small Cache M = o ⁢ ( d 2 ) 𝑀 𝑜 superscript 𝑑 2 M=o(d^{2}) italic_M = italic_o ( italic_d start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT )

When cache size is not so big, i.e. $M=o(d^{2})$, the attention backward is equivalent to matrix multiplication, thus having $O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$ bound on the I/O complexity.

We show the upper bound theorem below for the overall algorithm (see Algorithm 6) to solve the attention backward in small cache case.

###### Theorem C.12 (Small cache upper bound, formal version of Theorem 4.3 ) .

Suppose $n$ is the input length, $d$ is the head dimension, and $M$ is the cache size.
There is an algorithm (see Algorithm 6) outputs a $d\times d$ matrix $g=\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ (Definition 3.2) with I/O complexity $O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$, time complexity ${\cal T}_{\mathrm{mat}}(n,d,n)+{\cal T}_{\mathrm{mat}}(n,d,d)$, and space complexity $O(n^{2}+d^{2})$.

###### Proof.

Time/space complexity.

First, we notice that Algorithm 6 calculates the same gradients as the Algorithm 1 except that the former utilize cache to speed up the computation and specify the standard matrix multiplication computations in cache.
Thus, the overall time complexity ${\cal T}_{\mathrm{mat}}(n,d,n)+{\cal T}_{\mathrm{mat}}(n,d,d)$, and space complexity $O(n^{2}+d^{2})$ should be the same as Lemma C.3.

I/O complexity.

From Lemma C.5, C.7, C.9, and C.11, we know the overall I/O complexity is $O(\frac{n^{2}d+nd^{2}}{\sqrt{M}})+O(\frac{n^{2}}{\sqrt{M}})=O(\frac{n^{2}d+nd^
{2}}{\sqrt{M}})$.

Correctness.

From Lemma C.4, C.6, C.8, and C.10, the algorithm computes the correct $\frac{\mathrm{d}L(X)}{\mathrm{d}X}$.
∎

Figure: Algorithm 6 Attention gradient computation with small cache.

###### Theorem C.12 (Small cache upper bound, formal version of Theorem 4.3 ) .

###### Proof.

## Appendix D I/O Complexity Upper Bound for Large Cache

In this section, we establish the upper bound (Theorem D.5) for the I/O complexity in the case where the cache size is large, specifically when $M=\Omega(d^{2})$.
Section D.1 presents algorithms and analyses for attention gradient computation in the large cache setting.
Section D.2 provides the upper bound theorem for the large cache case.

Since our goal is to compute the backward pass of the attention mechanism, and the forward pass has already been performed, it is natural to assume that we have access to the softmax normalizing vector $l:=A\cdot{\bf 1}\in\mathbb{R}^{n}$ (Definition B.4) and the final attention forward output $O=\operatorname{diag}(l)^{-1}AV\in\mathbb{R}^{n\times d}$ (Definition B.7) where $A=\exp(A_{1}XA_{2}^{\top})$ (Definition B.3).

By utilizing these precomputed quantities from the forward pass, we can efficiently proceed with the backward computation while optimizing the I/O operations required.

### D.1 Algorithms for Attention Backward in Large Cache

We first give Algorithm 7 and its analysis in large cache case for computing intermediate variables $S,h$.

Figure: Algorithm 7 Attention gradient computation large cache phase 1. Compute $S,h$.

###### Lemma D.1 (Correctness of Phase 1) .

The AttentionGradientLargeCachePhase1 (Algorithm 7) outputs two $n\times d$ matrices $S=A_{1}X$ (Definition 3.1) and $h=A_{3}Y$ (Definition B.6).

###### Proof.

The algorithm first divide $A_{1},A_{3},X,Y$ into row/column blocks of size $B_{r}\times d$ or $d\times B_{c}$. Then it reads the row/column block matrices to compute the corresponding small blocks of $S,h$ by standard matrix multiplication. Thus, it computes the exact value for $S,h$.
∎

###### Lemma D.2 (I/O complexity of Phase 1) .

Suppose the cache size satisfy $nd\geq M\geq d$.
The I/O complexity of AttentionGradientLargeCachePhase1 (Algorithm 7) is $O(\frac{n^{2}d^{2}}{M}+\frac{nd^{3}}{M})$.

###### Proof.

Why such conditions for $B_{r},B_{c}$.

The cache size has three constraints, because we need matrices $A_{1,i},A_{3,i}\in\mathbb{R}^{B_{r}\times d}$, $X_{*,j},Y_{*,j}\in\mathbb{R}^{d\times B_{c}}$, and $S_{i,j},h_{i,j}\in\mathbb{R}^{B_{r}\times B_{c}}$ to fit into cache.
Thus, we have

$$ $\displaystyle B_{r}d=$ $\displaystyle~{}O(M)$ $\displaystyle B_{c}d=$ $\displaystyle~{}O(M)$ $\displaystyle B_{r}B_{c}=$ $\displaystyle~{}O(M)$ $$

Then, we need

$$ $\displaystyle B_{r}=$ $\displaystyle~{}O(M/d)$ $\displaystyle B_{c}=$ $\displaystyle~{}O(M/d)$ $$

By setting $B_{c}=\Theta(M/d)$, we have

$$ $\displaystyle B_{r}=$ $\displaystyle~{}\Theta(\min\{M/d,M/B_{c}\})$ $\displaystyle=$ $\displaystyle~{}\Theta(\min\{M/d,d\})$ $$

I/O complexity.
We know $B_{r}\leftarrow\min\{\lceil\frac{M}{4d}\rceil,d\}$ and $B_{c}\leftarrow\lceil\frac{M}{4d}\rceil$, also $T_{r}=\lceil\frac{n}{B_{r}}\rceil$ and $T_{c}=\lceil\frac{d}{B_{r}}\rceil$. Substituting $B_{r}$ into $T_{r}$, we get $T_{r}=O(\frac{nd}{M})$.
Observe that $T_{r}B_{r}=O(n)$ and $T_{c}B_{c}=O(d)$.

The I/O complexity can be computed by:

$$ $\displaystyle T_{r}(B_{r}d+T_{c}(dB_{c}))=$ $\displaystyle~{}O(nd)+T_{r}d^{2}$ $\displaystyle=$ $\displaystyle~{}O(nd)+O(\frac{nd}{M}d^{2})$ $\displaystyle=$ $\displaystyle~{}O(nd+\frac{nd^{3}}{M})$ $$

where the first step follows from $T_{r}B_{r}=O(n)$ and $T_{c}B_{c}=O(d)$, the second step follows from $T_{r}=O(\frac{nd}{M})$, and the last step follows from simple algebra.

Because $M\leq nd$, we have

$$ $\displaystyle O(nd+\frac{nd^{3}}{M})=$ $\displaystyle~{}O(\frac{ndM}{M}+\frac{nd^{3}}{M})$ $\displaystyle=$ $\displaystyle~{}O(\frac{n^{2}d^{2}}{M}+\frac{nd^{3}}{M})$ $$

Thus, the total I/O complexity is $O(\frac{n^{2}d^{2}}{M}+\frac{nd^{3}}{M})$
∎

Figure: Algorithm 8 Attention gradient computation large cache phase 2. Compute $g$.

We then give Algorithm 8 along with its analysis for computing the gradient $g$.

###### Lemma D.3 (Correctness of Phase 2) .

The AttentionGradientLargeCachePhase2 (Algorithm 8) outputs a $d\times d$ matrix $g$ (Definition 3.2).

###### Proof.

The algorithm first vertically divides the matrices $S$, $A_{2}$, $l$, $O$, $\mathrm{d}O$, $h$, and $A_{1}$ into row blocks of size $B_{r}\times d$ or $B_{c}\times d$. Following the computational graph (Fig. 2) and the no-cache algorithm (Algorithm 1), we compute the gradient $g$ exactly. It is important to note that, in algorithm design, we need to avoid reading the attention matrix $f\in\mathbb{R}^{n\times n}$ directly—even though it has been computed during the forward pass—or any matrices of size $B_{r}\times n$ or $B_{c}\times n$. Doing so would result in an $O(n^{2})$ I/O complexity, which cannot be improved through caching.
∎

###### Lemma D.4 (I/O complexity of Phase 2) .

Suppose the cache size satisfy $nd\geq M\geq d^{2}$.
The I/O complexity of AttentionGradientLargeCachePhase2 (Algorithm 8) is $O(\frac{n^{2}d^{2}}{M}+\frac{nd^{3}}{M})$.

###### Proof.

The reason for conditions of $B_{r},B_{c}$ is the same as the proof of Lemma D.2.
However, it is important to note that updating the gradient $g$ in the cache requires assuming a cache size of $M\geq d^{2}$. This is necessary because we fuse the key and query weight matrices into a single matrix $X\in\mathbb{R}^{d\times d}$. The update to the corresponding gradient $g$ in the cache is driven by the outer product representation of the matrix, as shown in Line 21 of Algorithm 8.

Next we show the I/O complexity.
Since $B_{r}\leftarrow\min\{\lceil\frac{M}{4d}\rceil,d\}$ and $B_{c}\leftarrow\lceil\frac{M}{4d}\rceil$, also $T_{r}=\lceil\frac{n}{B_{r}}\rceil$ and $T_{c}=\lceil\frac{n}{B_{r}}\rceil$, we get $T_{r}=O(\frac{nd}{M})$.
Also, we observe that $T_{r}B_{r}=O(n)$ and $T_{c}B_{c}=O(n)$.

The I/O complexity can be computed by:

$$ $\displaystyle T_{r}(B_{r}d+T_{c}B_{c}d)+d^{2}=$ $\displaystyle~{}O(nd)+T_{r}nd+d^{2}$ $\displaystyle=$ $\displaystyle~{}O(T_{r}nd)+d^{2}$ $\displaystyle=$ $\displaystyle~{}O(\frac{n^{2}d^{2}}{M})+d^{2}$ $$

where the first step follows from $T_{r}B_{r}=O(n)$ and $T_{c}B_{c}=O(n)$, the second step follows from $T_{r}\geq 1$, and the last step follows from $T_{r}=O(\frac{nd}{M})$.

Then, because $M\leq nd$, we can show

$$ $\displaystyle O(d^{2}+\frac{n^{2}d^{2}}{M})=$ $\displaystyle~{}O(\frac{d^{2}M}{M}+\frac{n^{2}d^{2}}{M})$ $\displaystyle=$ $\displaystyle~{}O(\frac{nd^{3}}{M}+\frac{n^{2}d^{2}}{M})$ $$

Thus, the total I/O complexity is $O(\frac{n^{2}d^{2}}{M}+\frac{nd^{3}}{M})$
∎

###### Lemma D.1 (Correctness of Phase 1) .

###### Proof.

###### Lemma D.2 (I/O complexity of Phase 1) .

###### Proof.

###### Lemma D.3 (Correctness of Phase 2) .

###### Proof.

###### Lemma D.4 (I/O complexity of Phase 2) .

###### Proof.

### D.2 Upper Bound for Attention Backward in Large Cache M = Ω ⁢ ( d 2 ) 𝑀 Ω superscript 𝑑 2 M=\Omega(d^{2}) italic_M = roman_Ω ( italic_d start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT )

In the large cache scenario, while it is feasible to precompute and store the $n\times n$ attention matrix, reading it will result in an unavoidable $O(n^{2})$ I/O complexity.
Inspired by FlashAttention [, , ], we present the following theorem, which provides an upper bound $O(\frac{n^{2}d^{2}+nd^{3}}{M})$ on the I/O complexity of the attention gradient algorithm in the large cache (Algorithm 9).

###### Theorem D.5 (Large cache upper bound, formal version of Theorem 4.1 ) .

Suppose $n$ is the input length, $d$ is the head dimension, and $nd\geq M\geq d^{2}$ is the cache size.
There is an algorithm (see Algorithm 9) outputs a $d\times d$ matrix $g=\frac{\mathrm{d}L(X)}{\mathrm{d}X}$ (Definition 3.2) with I/O complexity $O(\frac{n^{2}d^{2}+nd^{3}}{M})$.

###### Proof.

Correctness.
Combining Lemma D.1 and D.3, we finish the proof.

I/O complexity.
Combining Lemma D.2 and D.4, we finish the proof.
∎

Figure: Algorithm 9 Attention gradient computation with large cache.

###### Theorem D.5 (Large cache upper bound, formal version of Theorem 4.1 ) .

###### Proof.

## Appendix E Lower Bound for Attention Backward Computation

In this section, we prove the lower bound of the attention gradient computation. In Section E.1, we state some definition in graph theory that will be used to establish the framework of [] that will be used to analyze the I/O complexity. In Section E.2, we state some tools from previous works from I/O compleixty of standard matrix multiplication and attention forward computation. In Section E.3, we will establish our lower bounds of I/O complexity for attention backward passes in both large cache case and small cache case.

### E.1 Basic Definition in Graph Theory

HK [] introduces a method for analyzing I/O complexity using the concept of an $M$-partition on a graph. Before we define it, we first provide some definitions from graph theory.

###### Definition E.1 (Dominator set) .

Let $G=(V,E)$ be a directed acyclic graph and $S\subseteq V$. We define a set $D\subseteq V$ as a dominator set of $S$ if, for every path in $G$ from a input node to any node in $S$, there exists at least one node in $D$ on that path.

###### Definition E.2 (Minimum set) .

Let $G=(V,E)$ be a directed acyclic graph and $S\subseteq V$. We say that a set $M\subseteq S$ is a minimum set of $S$ if $M$ contains all nodes in $S$ that have no children in $S$.

###### Definition E.3 (Vertex subset dependence) .

Let $G=(V,E)$ be a directed acyclic graph. Let $V_{1},V_{2}\subseteq V$ be two disjoint subsets. We say that $V_{2}$ depends on $V_{1}$ if there is a directed edge from a node in $V_{1}$ to a node in $V_{2}$.

###### Definition E.4 (Cyclic dependence) .

Let $G=(V,E)$ be a directed acyclic graph. Let $V_{1},\ldots,V_{h}\subseteq V$ be $h$ disjoint subsets of $V$. We say that there is a cyclic dependence among $\{V_{1},\ldots,V_{h}\}$ if there exists a permutation $(i_{1},\ldots,i_{h})$ of $[h]$ such that $V_{i_{1}}$ depends on $V_{i_{h}}$, and for every $j\in\{2,\ldots,h\}$, $V_{i_{j}}$ depends on $V_{i_{j-1}}$.

Now, we are ready to define $M$-partitons.
In fact, the minimum number of sets in any $M$-partition provides a lower bound on the I/O complexity.

###### Definition E.5 ( $M$ -partition [ 49 ] ) .

Let $G=(V,E)$ be a directed acyclic graph. Let $V_{1},\ldots,V_{h}\subseteq V$ be $h$ disjoint subsets of $V$. We say that $\{V_{1},\ldots,V_{h}\}$ is a $M$-partition of $G$ if the following conditions are satisfied

- •
$\{V_{1},\ldots,V_{h}\}$ is a partition of $V$, i.e., $V_{1},\ldots,V_{h}$ are disjoint and $V=\bigcup_{i=1}^{h}V_{i}$.
- •
For each $V_{i}$, there exists a dominator set $D_{i}$ of $V_{i}$ such that $D_{i}$ has at most $M$ nodes.
- •
For each $V_{i}$, there exists a minimum set $M_{i}$ of $V_{i}$ such that $M_{i}$ has at most $M$ nodes.
- •
There is no cyclic dependence among $\{V_{1},\ldots,V_{h}\}$.

We use $P(G,M)$ to denote the minimum number of sets in any $M$-partition of $G$.

###### Definition E.1 (Dominator set) .

###### Definition E.2 (Minimum set) .

###### Definition E.3 (Vertex subset dependence) .

###### Definition E.4 (Cyclic dependence) .

###### Definition E.5 ( M 𝑀 M italic_M -partition [ 49 ] ) .

### E.2 Previous Tools for I/O Complexity

Now, we are ready to introduce some tools for I/O Complexity from HK [] by using an $M$-partition on a graph.

###### Lemma E.6 (Lemma 3.1 of HK [ 49 ] ) .

For any directed acyclic graph $G$ and any positive integer $M$, we have

$$ $\displaystyle Q(G,M)\geq M\cdot(P(G,2M)-1).$ $$

We omit $G$ when it is clear in the context.

We state two useful lemmas from previous works as follows.

###### Lemma E.7 (Lemma 3.3 of SY [ 102 ] ) .

Suppose that $M=\Omega(d^{2})$ and $A\in\mathbb{R}^{n_{1}\times d},B\in\mathbb{R}^{d\times n_{2}}$. Let $\mathcal{P}$ be an $M$-partition of the computational graph of any algorithm that computes $AB$ using standard matrix multiplication. Then for each $V^{\prime}\in\mathcal{P}$, $V^{\prime}$ contains at most $O(\frac{M^{2}}{d})$ product nodes $A_{i,k}B_{k,j}$, sum nodes $(AB)_{i,j}$, and all intermediate nodes in the summation trees.

In SY [], the matrices $A$ and $B$ in the above lemma are of sizes $n\times d$ and $d\times n$, respectively. We note that with slight modifications to the proofs, the result also holds when $A$ and $B$ have different sizes, specifically $n_{1}\times d$ and $d\times n_{2}$.

The next lemma gives the lower bound of I/O compleixty of standard matrix multiplication.

###### Lemma E.8 (Corollary 6.2 of HK [ 49 ] ) .

Let $A\in\mathbb{R}^{n_{1}\times d},B\in\mathbb{R}^{d\times n_{2}}$. The standard matrix multiplication algorithm computing $AB$ has I/O complexity $Q(M)=\Omega(\frac{n_{1}dn_{2}}{\sqrt{M}})$.

###### Lemma E.6 (Lemma 3.1 of HK [ 49 ] ) .

###### Lemma E.7 (Lemma 3.3 of SY [ 102 ] ) .

###### Lemma E.8 (Corollary 6.2 of HK [ 49 ] ) .

### E.3 Proof of Our Lower Bound

We establish the lower bounds of I/O complexity of attention gradient computation in both large cache case and small cache case. We first give the lower bound in the large cache case, i.e., the cache size $M=\Omega(d^{2})$.

###### Theorem E.9 (Large cache lower bound, formal version of Theorem 4.2 ) .

Suppose $n$ is the input length and $d$ is the head dimension.
Suppose the cache size $M=\Omega(d^{2})$.
Then the I/O complexity of attention gradient computation using standard matrix multiplication is $\Omega(\frac{n^{2}d^{2}+nd^{3}}{M})$.

###### Proof.

Any algorithm that computes the attention gradient needs to compute the matrix product $A_{1}XA_{2}^{\top}$ using standard matrix multiplication. Note that we compute $A_{1}XA_{2}^{\top}$ using standard matrix multiplication, so we either first compute $A_{1}X$ and then compute $(A_{1}X)A_{2}^{\top}$, or first compute $XA_{2}^{\top}$ and then compute $A_{1}(XA_{2}^{\top})$. In either case, we perform two matrix multiplications: one between an $n\times d$ matrix and a $d\times d$ matrix, and another between an $n\times d$ matrix and a $d\times n$ matrix. Without loss of generality, we assume the first case where we first compute $A_{1}X$.

Recall that the level-1 nodes are the product nodes $(A_{1})_{i,k}X_{k,j}$, the sum nodes $(A_{1}X)_{i,j}$, and all intermediate nodes in the summation trees. For every $V^{\prime}$ in an $M$-partition $\mathcal{P}$, by Lemma E.7, there are at most $O(\frac{M^{2}}{d})$ level-1 nodes in $V^{\prime}$. Since the number of sum nodes $(A_{1}X)_{i,j}$ is $nd^{2}$, the number of parts in the $M$-partition $\mathcal{P}$ is at least $\Omega(\frac{nd^{3}}{M^{2}})$. By Lemma E.6, the I/O complexity for computing $A_{1}X$ is $\Omega(\frac{n^{2}d}{M})$.

Similarly, we recall that level-2 nodes are the product nodes $(A_{1}X)_{i,k}(A_{2}^{\top})_{k,j}$, the sum nodes $((A_{1}X)A_{2}^{\top})_{i,j}$, and all intermediate nodes in the summation trees. For every $V^{\prime}$ in an $M$-partition $\mathcal{P}$, by Lemma E.7, there are at most $O(\frac{M^{2}}{d})$ level-2 nodes in $V^{\prime}$. Since the number of sum nodes $((A_{1}X)A_{2}^{\top})_{i,j}$ is $n^{2}d$, the number of parts in the $M$-partition $\mathcal{P}$ is at least $\Omega(\frac{n^{2}d^{2}}{M^{2}})$. By Lemma E.6, the I/O complexity for computing $(A_{1}X)A_{2}^{\top}$ is $\Omega(\frac{n^{2}d^{2}}{M})$.

Therefore, the I/O complexity of attention gradient computation is at least $\Omega(\frac{nd^{3}+n^{2}d^{2}}{M})$.
∎

Next, we give the lower bound in the small cache case, i.e., the cache size $M=o(d^{2})$.

###### Theorem E.10 (Small cache lower bound, formal version of Theorem 4.4 ) .

Suppose $n$ is the input length and $d$ is the head dimension.
Suppose the cache size $M=o(d^{2})$.
Then the I/O complexity of attention gradient computation using standard matrix multiplication is $\Omega(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.

###### Proof.

We show that when $M=o(d^{2})$, the attention gradient computation can be reduced to computing the matrix product $A_{1}XA_{2}^{\top}$. Note that we compute $A_{1}XA_{2}^{\top}$ using standard matrix multiplication, so we either compute $A_{1}X$ first and then compute $(A_{1}X)A_{2}^{\top}$, or we first compute $XA_{2}^{\top}$ and then $A_{1}(XA_{2}^{\top})$. However, both cases require performing one matrix multiplication between an $n\times d$ matrix and a $d\times d$ matrix, and one matrix multiplication between an $n\times d$ matrix and a $d\times n$ matrix. Hence, without loss of generality, we assume that $A_{1}X$ is computed first. By Lemma E.8, the I/O complexity of computing $A_{1}X$ is $\Omega(\frac{nd^{2}}{\sqrt{M}})$, and the I/O complexity of computing $(A_{1}X)A_{2}^{\top}$ is $\Omega(\frac{n^{2}d}{\sqrt{M}})$. Hence, the total I/O complexity of computing $A_{1}XA_{2}^{\top}$ is $\Omega(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.

Suppose that there is an algorithm $\mathcal{A}$ for attention gradient computation which has I/O complexity $o(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$. We construct an algorithm $\mathcal{B}$ that computes the matrix product $A_{1}XA_{2}^{\top}$ with I/O complexity $o(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$. Since $M<o(d^{2})$, we have $\frac{n^{2}d+nd^{2}}{\sqrt{M}}>\omega(n^{2}+nd)>\omega(n^{2})$, so algorithm $\mathcal{A}$ is able to transfer the all entries of matrix product $(A_{1}X)A_{2}^{\top}$ from cache to memory. In the language of the red-blue pebble game, algorithm $\mathcal{B}$ works as follows: whenever algorithm $\mathcal{A}$ delete a blue pebble from a node in $(A_{1}X)A_{2}^{\top}$, do not delete it; whenever algorithm $\mathcal{A}$ place a red pebble on a node in $(A_{1}X)A_{2}^{\top}$, also place a blue pebble on it. Since the I/O complexity of algorithm $\mathcal{A}$ is $o(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$ and we need an additional $n^{2}$ I/O operations to transfer the entries of the matrix product $(A_{1}X)A_{2}^{\top}$ from cache to memory. Since $n^{2}<o(\frac{n^{2}d}{\sqrt{M}})$, the overall I/O complexity of $\mathcal{B}$ is still $o(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$. However, this contradicts the fact that the I/O complexity of computing $A_{1}XA_{2}^{\top}$ is $\Omega(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$. Therefore, the I/O complexity of attention gradient computation using standard matrix multiplication is $\Omega(\frac{n^{2}d+nd^{2}}{\sqrt{M}})$.
∎

###### Theorem E.9 (Large cache lower bound, formal version of Theorem 4.2 ) .

###### Proof.

###### Theorem E.10 (Small cache lower bound, formal version of Theorem 4.4 ) .

###### Proof.

## Appendix F Sparse Attention Computation

In this section, we provide the lower bounds of sparse attention computation for both forward and backward passes. In Section F.1, we state previous tools of sparse matrix multiplication. In Section F.2, we provide the proofs of the lower bounds of sparse attention.

### F.1 Previous Tools For I/O complexity of Sparse Matrix Multiplication

We assume that sparse matrices are stored by listing only their non-zero entries along with their coordinates. Sparse semi-ring matrix multiplication restricts operations to addition and multiplication of these entries, which means that each output entry $(AB)_{i,j}$ can only be computed as the sum of products given by $\sum_{k}A_{i,k}B_{k,j}$.

###### Lemma F.1 (Theorem 2 of [ 85 ] ) .

Let $A\in\mathbb{R}^{n_{1}\times d}$ and $B\in\mathbb{R}^{d\times n_{2}}$ be two matrices such that $R_{1}:=\operatorname{nnz}(A)+\operatorname{nnz}(B)$ and $R_{2}:=\operatorname{nnz}(AB)$. The sparse semi-ring matrix multiplication that computes $AB$ has I/O complexity $\Omega(\min\{\frac{R_{1}^{2}}{M},\frac{R_{1}\sqrt{R_{2}}}{\sqrt{M}}\})$.

Note that in this statement, the I/O complexity also separates into the large cache case and the small cache case, but the dividing point may not be $d^{2}$. It depends on whether all the necessary values for computing each output entry can be stored in the cache during the computation.

###### Lemma F.1 (Theorem 2 of [ 85 ] ) .

### F.2 Our Lower Bounds for Sparse Attention Computation

We first prove a useful lemma which state the lower bound of I/O complexity of computing the attention matrix.

###### Lemma F.2 .

Let $A_{1}\in\mathbb{R}^{n\times d},X\in\mathbb{R}^{d\times d},A_{2}\in\mathbb{R}^{
d\times n}$ be three matrices. Let $Z_{A}:=\min\{\operatorname{nnz}(A_{1}),\operatorname{nnz}(A_{2})\},Z_{X}:=
\operatorname{nnz}(X),Z_{AX}=\min\{\operatorname{nnz}(A_{1}X),\operatorname{
nnz}(XA_{2}^{\top})\},Z_{AXA}:=\operatorname{nnz}(A_{1}XA_{2}^{\top})$. Then the sparse semi-ring matrix multiplication that computes $A_{1}XA_{2}^{\top}$ has I/O complexity $\Omega(\min\{\frac{Z_{A}^{2}+Z_{A}Z_{X}}{M},\frac{Z_{A}\sqrt{Z_{AXA}}+\sqrt{Z_
{A}Z_{X}Z_{AX}}}{\sqrt{M}}\})$.

###### Proof.

We first consider the case where all the necessary values for computing each output entry can be stored in the cache during the computation. Suppose that $A_{1}X$ is computed first, by Lemma F.1, computing $A_{1}X$ has I/O compleixty

$$ $\displaystyle\Omega(\frac{(\operatorname{nnz}(A_{1})+\operatorname{nnz}(X))^{2 }}{M})=$ $\displaystyle~{}\Omega(\frac{\operatorname{nnz}(A_{1})^{2}+2\operatorname{nnz} (A_{1})\operatorname{nnz}(X)+\operatorname{nnz}(X)^{2}}{M})$ $\displaystyle\geq$ $\displaystyle~{}\Omega(\frac{Z_{A}^{2}+2Z_{A}Z_{X}+Z_{X}^{2}}{M})$ $\displaystyle\geq$ $\displaystyle~{}\Omega(\frac{Z_{A}^{2}+2Z_{A}Z_{X}}{M})$ $$

where the first step follows by the basic algebra, the second step uses the definition of $Z_{A},Z_{X}$, and the last step follows from the basic algebra.
Then we compute the product $(A_{1}X)A_{2}^{\top}$, by Lemma F.1, computing $A_{1}X$ has I/O compleixty

$$ $\displaystyle\Omega(\frac{(\operatorname{nnz}(A_{1}X)+\operatorname{nnz}(A_{2} ))^{2}}{M})=$ $\displaystyle~{}\Omega(\frac{\operatorname{nnz}(A_{1}X)^{2}+2\operatorname{nnz }(A_{1}X)\operatorname{nnz}(A_{2})+\operatorname{nnz}(A_{2})^{2}}{M})$ $\displaystyle\geq$ $\displaystyle~{}\Omega(\frac{\operatorname{nnz}(A_{2})^{2}}{M})$ $\displaystyle=$ $\displaystyle~{}\Omega(\frac{Z_{A}^{2}}{M})$ $$

where the first and second steps follow by the basic algebra, and the last step uses the definition of $Z_{A}$. Therefore, computing $A_{1}XA_{2}^{\top}$ in this way has I/O complexity $\Omega(\frac{2Z_{1}^{2}+2Z_{1}Z_{2}}{M})=\Omega(\frac{Z_{1}^{2}+Z_{1}Z_{2}}{M}).$ Similary, suppose that $XA_{2}^{\top}$ is computed first. Then we can also get the I/O complexity $\Omega(\frac{Z_{1}^{2}+Z_{1}Z_{2}}{M})$.

Next, we consider the case where some elementary products of matrix multiplication needs to be written in the memory during the computation. Suppose that $A_{1}X$ is computed first, and then $(A_{1}X)A_{2}^{\top}$ is computed. By Lemma F.1, computing $(A_{1}X)$ has I/O compleixty

$$ $\displaystyle\Omega(\frac{(\operatorname{nnz}(A_{1})+\operatorname{nnz}(X)) \sqrt{\operatorname{nnz}(A_{1}X)})}{\sqrt{M}})\geq$ $\displaystyle~{}\Omega(\frac{2\sqrt{\operatorname{nnz}(A_{1})\operatorname{nnz }(X)}\sqrt{\operatorname{nnz}(A_{1}X)}}{\sqrt{M}})$ $\displaystyle\geq$ $\displaystyle~{}\Omega(\frac{2\sqrt{Z_{A}Z_{X}Z_{AX}}}{\sqrt{M}})$ $$

where the first step uses Cauchy-Schwarz inequality, the second step uses the definition of $Z_{A}$, $Z_{X}$ and $Z_{AXA}$.

By Lemma F.1, computing $(A_{1}X)A_{2}^{\top}$ has I/O compleixty

$$ $\displaystyle\Omega(\frac{(\operatorname{nnz}(A_{1}X)+\operatorname{nnz}(A_{2} ))\sqrt{\operatorname{nnz}(A_{1}XA_{2}^{\top})}}{\sqrt{M}})\geq$ $\displaystyle~{}\Omega(\frac{\operatorname{nnz}(A_{2})\sqrt{\operatorname{nnz} (A_{1}XA_{2}^{\top})}}{\sqrt{M}})$ $\displaystyle\geq$ $\displaystyle~{}\Omega(\frac{Z_{A}\sqrt{Z_{AXA}}}{\sqrt{M}}).$ $$

where the first step follows by the basic algebra, the second step uses the definition of $Z_{A}$ and $Z_{AXA}$. Therefore, computing $A_{1}XA_{2}^{\top}$ in this way has I/O complexity $\Omega(\frac{Z_{A}\sqrt{Z_{AXA}}+\sqrt{Z_{A}Z_{X}Z_{AX}}}{\sqrt{M}})$. Similary, suppose that $XA_{2}^{\top}$ is computed first. Then we can also get the I/O complexity $\Omega(\frac{Z_{A}\sqrt{Z_{AXA}}+\sqrt{Z_{A}Z_{X}Z_{AX}}}{\sqrt{M}})$.

Therefore, the sparse semi-ring matrix multiplication that computes $A_{1}XA_{2}^{\top}$ has I/O complexity $\Omega(\min\{\frac{Z_{A}^{2}+Z_{A}Z_{X}}{\sqrt{M}},\frac{Z_{A}\sqrt{Z_{AXA}}+
\sqrt{Z_{A}Z_{X}Z_{AX}}}{\sqrt{\sqrt{M}}}\})$.
∎

Next, we can apply Lemma F.2 to get the lower bound of sparse attention forward and backward passes.

###### Theorem F.3 (Lower bound for sparse attention forward) .

Suppose $n$ is the input length, $d$ is the head dimension, and $M$ is the cache size. Let $Z_{A}:=\min\{\operatorname{nnz}(A_{1}),\operatorname{nnz}(A_{2})\},Z_{X}:=
\operatorname{nnz}(X),Z_{AX}=\min\{\operatorname{nnz}(A_{1}X),\operatorname{
nnz}(XA_{2}^{\top})\},Z_{AXA}:=\operatorname{nnz}(A_{1}XA_{2}^{\top})$. Then any algorithm for attention forward computation using sparse semi-ring matrix multiplication has I/O complexity $\Omega(\min\{\frac{Z_{A}^{2}+Z_{A}Z_{X}}{M},\frac{Z_{A}\sqrt{Z_{AXA}}+\sqrt{Z_
{A}Z_{X}Z_{AX}}}{\sqrt{M}}\})$.

###### Proof.

Any algorithm for attention forward computation needs to compute the matrix product $A_{1}XA_{2}^{\top}$ to obtain the attention matrix. Thus by applying Lemma F.2, we complete the proof.
∎

###### Theorem F.4 (Lower bound for sparse attention backward) .

Suppose $n$ is the input length, $d$ is the head dimension, and $M$ is the cache size. Let $Z_{A}:=\min\{\operatorname{nnz}(A_{1}),\operatorname{nnz}(A_{2})\},Z_{X}:=
\operatorname{nnz}(X),Z_{AX}=\min\{\operatorname{nnz}(A_{1}X),\operatorname{
nnz}(XA_{2}^{\top})\},Z_{AXA}:=\operatorname{nnz}(A_{1}XA_{2}^{\top})$. Then any algorithm for attention backward computation using sparse semi-ring matrix multiplication has I/O complexity $\Omega(\min\{\frac{Z_{A}^{2}+Z_{A}Z_{X}}{M},\frac{Z_{A}\sqrt{Z_{AXA}}+\sqrt{Z_
{A}Z_{X}Z_{AX}}}{\sqrt{M}}\})$.

###### Proof.

Any algorithm for attention backward computation needs to compute the matrix product $A_{1}XA_{2}^{\top}$ to obtain the attention matrix. Thus by applying Lemma F.2, we complete the proof.
∎

###### Lemma F.2 .

###### Proof.

###### Theorem F.3 (Lower bound for sparse attention forward) .

###### Proof.

###### Theorem F.4 (Lower bound for sparse attention backward) .

###### Proof.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: None
soundness: 3
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.
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.
weaknesses: Some practical implementation results could enhance the paper.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
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?
soundness: 3
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.
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.
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.

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Under what hardware and model will the attention computation fall in different cache cases?
soundness: 3
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.
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.
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.

Meta-review: 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.

Decision: Reject