逼近 I/O 最优性

arXiv:2605.23751v2 公告类型：替换 摘要：我们重新审视大型语言模型中注意力机制的 I/O 复杂度。给定查询矩阵、键矩阵和值矩阵 $Q,K,V\in\mathbb{R}^{n\times d}$，以及一台快速内存容量为 $M$ 的机器，目标是在快速内存与慢速内存之间进行最少次数的数据传输，以计算“注意力矩阵”$A=\text{softmax}(Q K ^{\top}/\sqrt{d}) V$。现有文献中的方法，尤其是 FlashAttention 及其变体，其 I/O 开销依赖于上下文长度 $n$ 的平方；而一个平凡的下界只要求进行 $\Omega(nd)$ 次 I/O，以读取输入并写入输出。在本文中，我们通过开发受 Alman and Song [NeurIPS'23] 最新近似注意力方法启发的 I/O 高效算法，证明了对 $n$ 的依赖可以降至近似线性。我们还在不同参数范围内给出了（针对具体算法的）下界，这些下界与相应的上界几乎匹配，表明所提出的算法确实接近 I/O 最优。