Self-attention is at the heart of the popular Transformer architecture, yet suffers from quadratic time and memory complexity.
In a recent significant development, FlashAttention shows that the I/O complexity of attention is the true bottleneck in scaling Transformers.
Given two levels of memory hierarchy, a fast cache (e.g. GPU on-chip SRAM)  where computation happens and a slow memory (e.g. GPU high-bandwidth memory)  where the data resides, the I/O complexity measures the number of accesses to the slow memory.
FlashAttention is an I/O-aware algorithm for self-attention that requiresI/O operations whereis the dimension of the attention matrix,is the head-dimension andis the size of cache.However, is this I/O complexity optimal?The known lower bound only rules out an I/O complexity ofwhen, since the output of the attention mechanism that needs to be written in the slow memory is.
The main question that remained open after FlashAttention is whether this I/O complexity is optimal for any value 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 ofwithin any constant factors.
Further, we give a better algorithm with lower I/O complexity for, 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 future.