Approaching I/O-optimality for Approximate Attention

arXiv:2605.23751v2 Announce Type: replace Abstract: We revisit the I/O complexity of attention in large language models. Given query-key-value matrices $Q,K,V\in\mathbb{R}^{n\times d}$, and a machine with fast memory size $M$, the goal is to compute the "attention matrix" $A=\text{softmax}(Q K ^{\top}/\sqrt{d}) V$ with the minimal number of data transfers between fast and slow memory. Existing methods in the literature, most notably FlashAttention and its variants, incur an I/O cost that depends quadratically on the context length $n$, while a trivial lower bound only requires $\Omega(nd)$ I/O's to read the inputs and write the output. In this work we show that the dependence on $n$ can be reduced to almost-linear, by developing I/O-efficient algorithms inspired by the recent approximate attention method of Alman and Song [NeurIPS'23]. We also provide (algorithm-specific) lower bounds in different parameter regimes, which nearly match the corresponding upper bounds, indicating that the proposed algorithms are indeed close to I/O-optimal.