Reading Barna Saha and Christopher Ye, The I/O Complexity of Attention, or How Optimal is Flash Attention? (2024). Original paper
We often measure an algorithm’s cost in FLOPs. On a GPU, however, less arithmetic does not necessarily mean less time: arithmetic units may be waiting for data to travel from device memory to on-chip storage. FlashAttention’s central contribution is to reduce that movement. This paper asks a further question: how much room remains?
The answer depends on the memory regime. In the paper’s large-cache model, FlashAttention already achieves the optimal asymptotic I/O complexity. With a small cache, a different tiling strategy does better. Optimality here concerns the order of data movement, rather than the fastest possible implementation.
Attention as a mathematical problem
For a single attention head, let
Here $N$ is the sequence length and $d$ the head dimension. The score between query $i$ and key $j$ is
Actual models usually divide this score by $\sqrt d$. That scaling does not change the I/O orders discussed here, so we omit it.
Softmax turns each row of scores into weights:
The output is a weighted average:
In matrix form,
Computing both products with classical matrix multiplication takes roughly $4N^2d$ FLOPs, counting multiplication and addition separately and ignoring lower-order terms. Yet the intermediate score matrix has $N^2$ entries, while the output has only $Nd$. Must we store that enormous intermediate matrix?
Two memory levels: a warehouse and a workbench
The paper uses an idealized model. Slow memory is large; fast memory holds only $M$ numbers. All computation happens in fast memory, and each number loaded or stored counts as one I/O operation.
On a GPU, we can think of the slow level as device memory and the fast level as on-chip storage. Real hardware has registers, shared memory, and multiple cache levels; the model abstracts one boundary within that hierarchy.
The units matter. $M$ counts elements, rather than bytes. If fast memory holds $C$ bytes and each element occupies $w$ bytes, then
Moving $L$ elements means transferring $wL$ bytes. The paper seeks the smallest possible asymptotic order of $L$, rather than directly predicting seconds.
Softmax without storing the score matrix
Fix one query and define
Then its output is $o=u/z$. We can process keys and values block by block, adding each block’s contributions to $z$ and $u$ without retaining every score.
For numerical stability, implementations also maintain a maximum $m$ and use
If the new block’s score indices form a set $J$, update the maximum to
The existing accumulators must be rescaled:
The final output remains $o=u/\ell$. These equations explain why score computation can be fused with accumulation and the score blocks discarded. They are background for understanding FlashAttention, rather than new contributions of this lower-bound paper. Exact attention means evaluating the same mathematical function; changing floating-point execution order can still change rounding.
Large caches: where FlashAttention’s I/O comes from
Let $B$ be the number of queries in a fast-memory block. Each query needs $d$ elements, alongside output accumulators and temporary storage. Ignoring constants,
One intuitive schedule keeps a query block and its output state in fast memory, then scans all of $K,V$. Each scan costs $\Theta(Nd)$, and there are about $N/B$ query blocks. The leading term is therefore
As a teaching example, take $N=4096,d=64,M=65536$ elements.
These are values of asymptotic expressions, rather than exact load and store counts. Actual block sizes must reserve room for accumulators, temporary scores, and other state. The example illustrates the scaling: in this regime, doubling fast memory can roughly halve the leading movement term.
The paper proves a matching lower bound for the classical matrix multiplication model when $M\ge d^2$. This is more than a convenient tiling strategy: alternative schedules cannot improve its asymptotic order.
Small caches: when storing an intermediate matrix helps
When $M<d^2$, tiling around whole vectors of dimension $d$ is too restrictive. Instead, use square blocks with side length
so that two input blocks and an output block fit in fast memory. A local product uses $\Theta(b^2)$ elements to support $\Theta(b^3)$ scalar multiply-adds, giving reuse efficiency $\Theta(b)=\Theta(\sqrt M)$.
The two attention products require $\Theta(N^2d)$ multiply-adds in total, so this strategy achieves
It writes the exponentiated score matrix to slow memory, consuming $\Theta(N^2)$ space. But the small-cache condition implies
Writing $N^2$ elements therefore does not increase the asymptotic leading term. Never storing an intermediate matrix is not a universally optimal rule.
Compare the two strategies:
For $M
Putting the regimes together
Theorem 3.1 gives the following leading term for classical matrix multiplication:
For ordinary attention dimensions, including compulsory input and output traffic, it is useful to understand this as
The added $Nd$ emphasizes the cost of reading inputs and writing outputs. With an extremely large cache, the leading-term expression cannot be extrapolated to traffic smaller than the output itself. Here we assume the common regime $d\le N$; extreme rectangular shapes require explicit checks of block dimensions and compulsory I/O.
| Fast-memory regime | Suitable strategy | Leading movement term |
|---|---|---|
| $M<d^2$ | Square matrix multiplication tiles; store the intermediate matrix | $N^2d/\sqrt M$ |
| $M\ge d^2$ | Fused streaming accumulation; avoid the full score matrix | $N^2d^2/M$ |
Proving a lower bound for a fixed computation graph
The red-blue pebble game represents dependencies as a DAG. Inputs and intermediate values are vertices, and dependencies are edges. Red pebbles represent values in fast memory, with at most $M$ available; blue pebbles represent values in slow memory.
There is no need to enumerate every schedule. Divide an execution into phases of $M$ I/O operations, then ask how many scores one phase can complete.
In the large-cache regime, the paper proves that a phase can complete only
entries of $QK^\top$. Intuitively, each query or key vector requires $d$ elements. About $M/d$ queries and $M/d$ keys can form only about $(M/d)^2$ pairs.
A rigorous proof must also handle partially computed dot products. Counting complete vectors alone is insufficient. The paper analyzes summation trees and partition boundaries to account for these partial computations.
There are $N^2$ scores, so at least $\Omega(N^2d^2/M^2)$ phases are needed. Multiplying by the $\Theta(M)$ movement budget per phase gives
What if we use a Strassen-like algorithm?
A proof for a fixed DAG does not automatically rule out another matrix multiplication algorithm: changing the algorithm changes the graph. The paper’s second part addresses this distinction.
The authors introduce a communication problem called matrix entry compression. Alice knows the input matrices and sends Bob a short message. Bob must recover $B$ entries of their product. If those entries contain substantial independent information, the message cannot be arbitrarily short.
The key communication lower bound has the form
measured in finite-field elements. An I/O phase has an information budget of only $O(M)$ elements, so it can produce at most
product entries. This yields the lower bound in Theorem 4.8:
When $M\ge d^2$, this matches FlashAttention’s leading term. The proof uses a finite field $\mathbb F_q$ with $q>N$ and constructions based on linear independence. The corresponding binary-input result loses logarithmic factors. It cannot be applied without qualification as an information-theoretic theorem for arbitrary floating-point representations.
The scope also matters: the extension allows different matrix multiplication methods while still computing entries of $QK^\top$. It is not an unrestricted lower bound for every conceivable attention algorithm that might bypass those explicit entries, use approximation, or exploit special inputs.
For small caches, the authors connect attention to the I/O complexity of rectangular matrix multiplication. Once fast matrix multiplication is allowed, the paper does not supply one tight closed-form expression covering every case.
From a movement lower bound to a performance upper bound
The following is an application of the paper to performance modeling. If at least $L_{\min}$ elements must cross a memory boundary, each occupies $w$ bytes, and that boundary has a bandwidth ceiling $\beta$, then
If the chosen algorithm class also requires at least $F_{\min}$ FLOPs, and the compute ceiling at the relevant precision is $P$, then
Taking the maximum allows ideal overlap between computation and transfer. This is an optimistic time lower bound, rather than a prediction of actual runtime. An asymptotic $\Omega$ expression alone also supplies no reliable numerical constant.
For a Bound Estimator on a Mac, first specify the memory boundary, effective capacity, precision, and algorithm class. Unified memory does not eliminate movement between on-chip storage and DRAM.
The paper studies full $N\times N$ attention. Single-token decoding has query length one, so its cost cannot be obtained by copying an $N^2$ formula. It needs rectangular analysis and a separate account of where the existing KV state resides.
What remains after reading the paper?
FlashAttention lets a score complete its entire purpose in fast memory: it is generated, participates in normalization and weighted accumulation, then disappears. The lower-bound paper shows that, in the specified model’s large-cache regime, this strategy reaches the asymptotic limit of data movement.
There is still distance between optimal movement and minimum runtime. Parallelism, instruction throughput, synchronization, register occupancy, and constant factors remain open to improvement. Theory is most useful when it tells us where an improvement in asymptotic order is still possible, and where the next gains must come from implementation efficiency.
Where to read in the original
- Pages 4–6: two-level memory, red-blue pebbling, and the partition lemma.
- Pages 7–13, Theorems 3.1, 3.2, and 3.5: large- and small-cache results for classical matrix multiplication.
- Pages 14–20, Theorems 4.8 and 4.14: communication complexity and broader matrix multiplication algorithms.
- Page 21, Theorem 4.17 and the conclusion: rectangular matrix multiplication in the small-cache regime and open questions.
This is an original exposition. Online softmax, numerical examples, and hardware performance conversions provide intuition; the original paper determines each theorem’s precise scope.