FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao Department of Computer Science, Stanford University Daniel Y. Fu Department of Computer Science, Stanford University Stefano Ermon Department of Computer Science, Stanford University Atri Rudra Department of Computer Science and Engineering, University at Buffalo, SUNY Christopher Ré Department of Computer Science, Stanford University
Abstract
Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware—accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3 $\times$ speedup on GPT-2 (seq. length 1K),
中文速览
传统Transformer的自注意力机制在处理长序列时面临计算量和显存占用随序列长度平方增长的瓶颈，而以往的近似注意力方法虽然减少了浮点运算量，却因忽视GPU内存读写开销（IO）而未能带来实际加速。FlashAttention的核心思路是让注意力计算变得"IO感知"——通过分块（tiling）技术将输入切成小块逐步完成softmax计算，并在反向传播时通过重新计算代替存储庞大的中间矩阵，从而大幅减少慢速高带宽内存（HBM）与快速片上SRAM之间的数据搬运次数。实验结果表明，FlashAttention在保持精确计算结果的同时，训练速度比现有基准最高提升7.6倍，显存占用仅随序列长度线性增长，还首次让Transformer在序列长度16K乃至64K的极难任务上取得了有意义的准确率。这项工作表明，优化硬件内存访问模式而非单纯削减算术运算量，才是真正加速深度学习算子的关键所在。
原文 arXiv:2205.14135;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2205.14135v2