FlashAttention通过减少数据搬运让精确注意力更快

长序列Transformer速度慢且内存消耗大，因为标准自注意力会物化大型中间矩阵，并在GPU高带宽内存与计算单元之间反复搬运这些矩阵。

诱人的修复方案是近似注意力：改变模型的数学计算以降低复杂度。FlashAttention走了相反的路线。它保持精确注意力，但在SRAM中处理数据块，并重新计算运行中的softmax统计量，因此完整的注意力矩阵无需被存储。

论文报告称，与MLPerf 1.1训练纪录相比，在BERT-large上实现了15%的端到端实际耗时加速；在序列长度为1K的GPT-2上实现了3倍加速；在Long Range Arena任务上实现了2.4倍加速。

优化内存流量，而非数学计算

核心约束是HBM与SRAM之间的IO，而不仅仅是浮点运算。

分块让每个数据块在数据已位于芯片上时完成有用的工作。

保持算法精确避免了使许多近似方法缺乏吸引力的质量退化问题。

同一原语在不改变Transformer抽象的情况下解锁了更长的上下文和新的序列任务。

当硬件数据搬运成为瓶颈时，更好的算法可以是那种在更少次穿越内存层级的情况下计算出相同答案的算法。

FlashAttention成为高效Transformer训练和推理的标准构建模块。其更大的启示超越了注意力本身：现代加速器算法往往先围绕内存层级进行设计，然后再发明新的模型数学计算，从而取胜。