FlashAttention

概述

FlashAttention 是斯坦福大学 Tri Dao 等人于 2022 年提出的 IO 感知精确注意力算法。其出发点是:标准注意力实现的瓶颈不在浮点运算,而在 GPU 高带宽显存(HBM)与片上 SRAM 之间反复读写 N×N 注意力矩阵的内存访问;序列越长,这一开销越占主导。

原理

FlashAttention 采用分块(Tiling)策略,把 Q、K、V 切成小块装入 SRAM,在片上完成局部注意力计算,借助在线 softmax 技巧增量合并各块结果,自始至终不在显存中物化完整的注意力矩阵;反向传播则用重计算代替存储中间矩阵。这使显存占用从序列长度的平方降为线性,同时因大幅减少 HBM 访问而获得数倍实际加速,且计算结果与标准注意力完全一致。

演进与影响

FlashAttention-2(2023)重构了并行化,速度约提升一倍;FlashAttention-3(2024)针对 NVIDIA Hopper 架构利用异步与 FP8 特性进一步提速。该算法已被 PyTorch、Hugging Face、vLLM、Megatron-LM 等广泛集成,是长上下文大模型训练与推理的事实标配。