# The I/O Complexity of Attention, or How Optimal is Flash Attention?

arXiv: https://arxiv.org/abs/2402.07443 · HTML: https://mlsanity.com/paper/2402.07443/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 35 references that are in the map.
- First posted on arXiv 12 Feb 2024.

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Scatterbrain: Unifying Sparse and Low-rank Attention Approximation](https://mlsanity.com/paper/2110.15343/) (2021)
- [Fast Attention Requires Bounded Entries](https://mlsanity.com/paper/2302.13214/) (2023)
- [HyperAttention: Long-context Attention in Near-Linear Time](https://mlsanity.com/paper/2310.05869/) (2023)
- [New Bounds for Matrix Multiplication: from Alpha to Omega](https://mlsanity.com/paper/2307.07970/) (2023)
- [Fast Learning Requires Good Memory: A Time-Space Lower Bound for Parity Learning](https://mlsanity.com/paper/1602.05161/) (2016)
- [Memory-Sample Tradeoffs for Linear Regression with Small Error](https://mlsanity.com/paper/1904.08544/) (2019)
- [Memory Bounds for the Experts Problem](https://mlsanity.com/paper/2204.09837/) (2022)
- [Memory Bounds for Continual Learning](https://mlsanity.com/paper/2204.10830/) (2022)
- [Online Prediction in Sub-linear Space](https://mlsanity.com/paper/2207.07974/) (2022)
- [Near Optimal Memory-Regret Tradeoff for Online Learning](https://mlsanity.com/paper/2303.01673/) (2023)
- [Space lower bounds for linear prediction in the streaming model](https://mlsanity.com/paper/1902.03498/) (2019)
- [Efficient Convex Optimization Requires Superlinear Memory](https://mlsanity.com/paper/2203.15260/) (2022)
- [Improved Rectangular Matrix Multiplication using Powers of the Coppersmith-Winograd Tensor](https://mlsanity.com/paper/1708.05622/) (2017)
- [On Communication Complexity of Classification Problems](https://mlsanity.com/paper/1711.05893/) (2017)
- [Towards a combinatorial characterization of bounded memory learning](https://mlsanity.com/paper/2002.03123/) (2020)
- [Open Problem: The Oracle Complexity of Convex Optimization with Limited Memory](https://mlsanity.com/paper/1907.00762/) (2019)
- [Memory-Query Tradeoffs for Randomized Convex Optimization](https://mlsanity.com/paper/2306.12534/) (2023)
- [The Input/Output Complexity of Triangle Enumeration](https://mlsanity.com/paper/1312.0723/) (2013)
- [The Input/Output Complexity of Sparse Matrix Multiplication](https://mlsanity.com/paper/1403.3551/) (2014)
- [The I/O complexity of Strassen's matrix multiplication with recomputation](https://mlsanity.com/paper/1605.02224/) (2016)

## Built on by (most cited first, 6 of 6)

- [A Tighter Complexity Analysis of SparseGPT](https://mlsanity.com/paper/2408.12151/) (2024)
- [Debiasing Watermarks for Large Language Models via Maximal Coupling](https://mlsanity.com/paper/2411.11203/) (2024)
- [Compression Barriers for Autoregressive Transformers](https://mlsanity.com/paper/2502.15955/) (2025)
- [Model Agnostic Hybrid Sharding For Heterogeneous Distributed Inference](https://mlsanity.com/paper/2407.19775/) (2024)
- [kNN Attention Demystified: A Theoretical Exploration for Scalable Transformers](https://mlsanity.com/paper/2411.04013/) (2024)
- [FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention](https://mlsanity.com/paper/2609.24089/) (2026)

## Abstract (arXiv)

Self-attention is at the heart of the popular Transformer architecture, yet suffers from quadratic time and memory complexity. The breakthrough FlashAttention algorithm revealed I/O complexity as the true bottleneck in scaling Transformers. Given two levels of memory hierarchy, a fast cache (e.g. GPU on-chip SRAM) and a slow memory (e.g. GPU high-bandwidth memory), the I/O complexity measures the number of accesses to memory. FlashAttention computes attention using fracN2d2M I/O operations where N is the dimension of the attention matrix, d the head-dimension and M the cache size. However, is this I/O complexity optimal? The known lower bound only rules out an I/O complexity of o(Nd) when M=Theta(Nd), since the output that needs to be written to slow memory is Omega(Nd). This leads to the main question of our work: Is FlashAttention I/O optimal for all values of M? We resolve the above question in its full generality by showing an I/O complexity lower bound that matches the upper bound provided by FlashAttention for any values of M ≥ d2 within any constant factors. Further, we give a better algorithm with lower I/O complexity for M < d2, and show that it is optimal as well. Moreover, our lower bounds do not rely on using combinatorial matrix multiplication for computing the attention matrix. We show even if one uses fast matrix multiplication, the above I/O complexity bounds cannot be improved. We do so by introducing a new communication complexity protocol for matrix compression, and connecting communication complexity to I/O complexity. To the best of our knowledge, this is the first work to establish a connection between communication complexity and I/O complexity, and we believe this connection could be of independent interest and will find many more applications in proving I/O complexity lower bounds in the future.