FlashAttention 是什么？让 Transformer 又快又省显存的注意力算法

FlashAttention 是一种快速、省显存且结果精确的注意力计算算法，由斯坦福大学的 Tri Dao 等人在 2022 年提出，合作者包括纽约州立大学布法罗分校的研究人员，论文标题是《FlashAttention：具有 IO 感知的快速且内存高效的精确注意力》。它已被集成到 PyTorch 2.0 中，成为如今几乎所有大模型训练和推理的标配技术。

注意力为什么慢

Transformer 中的自注意力，需要计算每个 token 和其他所有 token 之间的注意力分数，形成一个巨大的矩阵：如果序列有 N 个 token，这个矩阵就有 N 乘 N 个元素，计算量和显存占用都随序列长度平方增长。很多加速方法试图通过近似来减少计算量，但往往会降低模型效果，实际上也不一定更快。

Tri Dao 等人的关键洞察是：注意力计算的真正瓶颈，往往不在于计算本身，而在于数据在 GPU 不同存储层级之间的搬运。GPU 有两种主要的存储：一种是容量大但读写较慢的高带宽显存（HBM），另一种是容量很小但速度极快的片上存储（SRAM）。传统的注意力实现会把那个巨大的 N 乘 N 矩阵写到 HBM 里，再读出来做 softmax，再写回去，反复地搬运大量数据，时间大部分花在了等待数据上。

FlashAttention 怎么做

FlashAttention 是一种“IO 感知”的算法，它的核心是减少对慢速显存的读写。分块计算：把查询、键、值切成小块，每次只把一小块数据载入快速的片上存储，在里面完成这一块的注意力计算，而不是一次性算出整个大矩阵。为了在分块的情况下正确地计算 softmax，它使用了一种在线 softmax 的技巧，边算边更新归一化所需的统计量。这样，那个巨大的 N 乘 N 矩阵从来不需要完整地写进显存。重计算：在训练的反向传播时，不保存庞大的注意力矩阵，而是在需要时根据保存下来的少量统计量重新计算一遍。多做一点计算，却省下了大量的显存读写，总体反而更快。

关键在于，FlashAttention 的结果是精确的，和标准注意力的结果完全一致，不是近似，因此不会影响模型效果。

影响

论文报告，FlashAttention 让 BERT、GPT-2 等模型的训练速度提升了数倍，显存占用从随序列长度平方增长降为线性增长，从而能够训练更长上下文的模型。此后 Tri Dao 又推出了 FlashAttention-2 和针对新一代 GPU 优化的 FlashAttention-3，进一步提升了效率。如今大模型能够支持几十万甚至上百万 token 的上下文，FlashAttention 功不可没。