FlashAttention 论文精读：一个IO 感知的注意力算法，如何改变了大模型的训练速度

本文来自微信公众号： 歪睿老哥 ，作者：歪睿老哥，原文标题：《FlashAttention 论文精读：一个 IO 感知的注意力算法，如何改变了大模型的训练速度》

故事是这样的。

2022夏天，斯坦福的一帮人发了篇论文，标题叫《FlashAttention:Fast and Memory-Efficient Exact Attention with IO-Awareness》。

这名字听着挺无聊对吧，一堆缩写，一个技术名词，跟每天几百篇论文里随便哪篇长得都一样。

但我看了之后，觉得这件事挺有意思的。

因为它解决了一个几乎所有做Transformer的人都踩过坑的问题：大模型推理慢，到底慢在哪里。

很多人第一反应是"算不动"，觉得算力不够，换更大的GPU就行。

但这篇文章的作者Tri Dao一帮人说了句不一样的话：

不是算不动，是数据搬不动。

他们管这个叫IO-Awareness——在写注意力算法的时候，不要只盯着FLOP数看，要把GPU显存层级之间的读写次数也算进去。

听起来像废话对吧，但就是这个"废话"，让Transformer的训练速度提升了3倍，内存消耗降低了20倍，还让Transformer第一次在64K长度的序列上跑出了好结果。

64K。

之前所有方法，要么跑不了，要么跑出来跟随机猜差不多。

这篇文章，我想把这个论文的核心内容，用人话讲清楚。

1.背景：Transformer的瓶颈到底在哪里

问题：注意力机制的平方复杂度

先看最基本的注意力公式：

O=softmax(QK^T/sqrt(d))V

三个矩阵相乘，再加一个softmax。看起来很简单对吧。

但问题在于，Q和K的矩阵乘法会生成一个N×N的注意力矩阵，其中N是序列长度。

这意味着时间和空间复杂度都是O(N²)。

序列长度翻倍，计算量翻四倍。

这个平方复杂度是Transformer的天然缺陷，从2017年那篇"Attention Is All You Need"出来就没变过。

现有方案的局限

过去几年，一堆人想解决这个问题。

方案主要分两类：

第一类，近似注意力。用稀疏化、低秩分解、核函数近似等等手段，把注意力矩阵从N×N压缩到接近N×1。

这类方法理论上把计算复杂度降到了线性或近线性，但实际跑起来，墙钟时间（wall-clock time）并没有明显加速。

为什么？

因为很多方案只关注减少FLOP，忽略了内存访问的开销。

第二类，稀疏注意力。让每个token只关注有限的其他token，直接剪掉大量注意力连接。

这种方法确实减少了计算量，但稀疏模式本身也有内存访问的overhead，而且效果往往不如dense attention。

作者的判断

这篇文章的作者认为，现有方案没效果的根本原因不是算法不行，而是没有考虑GPU内存层级的IO特性。

现代GPU的内存层级是这样的：

DRAM（系统内存）：容量最大（几十GB到几百GB），速度最慢（大约12.8 GB/s）

HBM（高带宽内存）：GPU显存，容量中等（40-80 GB），速度中等（大约1.5-2.0 TB/s）

SRAM（片上缓存）：容量最小（A100每个SM大约192 KB），速度最快（大约19 TB/s）

从HBM到SRAM，带宽差了一个数量级。

但现代GPU的计算速度已经超过内存速度了。操作越来越被内存访问（IO）而不是计算本身瓶颈住。

所以，关键问题不是FLOP多不多，而是有多少数据在HBM和SRAM之间来回搬。

这就是"IO-Awareness"的核心思想。

2.标准注意力实现的问题

标准算法：三步走

标准的注意力实现，通常是三步：

第一步：计算QK^T

把Q和K从HBM读到SRAM，在芯片上算QK^T，结果写回HBM。

这一步产生了一个N×N的注意力分数矩阵S。

第二步：Softmax

把S从HBM读出来，逐行做softmax，得到P矩阵。P再写回HBM。

第三步：PV

把P和V从HBM读出来，在芯片上算PV，结果写回HBM。

这三步看起来很自然对吧，但每步都在做一件事：把中间结果从HBM写出去，再从HBM读进来。

HBM访问次数的分析

让我们算一下HBM的访问总量。

前向传播：

第一步：读Q、K，写S→O(N²d+N²)次HBM访问

第二步：读S，写P→O(N²)次HBM访问

第三步：读P、V，写O→O(N²d+N²)次HBM访问

前向传播总共：O(N²d+N²)次HBM访问，是序列长度的平方级。

反向传播：

反向传播需要用到前向计算的S和P矩阵来计算梯度，所以同样的，要读S、P写回dQ、dK、dV。

反向传播也大约O(N²d+N²)次HBM访问。

核心矛盾

整个前向+反向传播，HBM访问总量大约是O(N²d)次。

但输入Q、K、V本身的总大小只有O(Nd)，输出O也只有O(Nd)。

数据量是O(Nd)的东西，为什么要做O(N²d)的HBM访问？

多出来的O(N²)次访问，全是用在那个大得离谱的N×N注意力矩阵上。

这个矩阵太大，放不下SRAM，只能在HBM和SRAM之间反复搬。

这就是标准实现的根本问题。

3.FlashAttention算法：核心思路

两个关键技术

FlashAttention的思路很简单：用两个经典技术，避免把N×N注意力矩阵写到HBM上。

这两个技术是：

1.Tiling（分块计算）

2.Recomputation（重计算）

Tiling：分块做Softmax

标准softmax需要对整行做归一化，看起来必须把整行读进来才能算。

但math上有个技巧：softmax可以分块计算。

具体来说，如果我把一个向量x拆成两段x¹和x²，那么整个向量x的softmax结果，可以用x¹和x²各自的softmax统计量（最大值m和归一化因子ℓ）来逐步合并。

公式大概是：

这样，每次处理一个块，只需要记录两个小值（m和ℓ），就能把结果正确合并起来。

所以FlashAttention的做法是：

把K、V分成多个块

每次只把一个块加载到SRAM

对Q的每个块，和K的这个块算QK^T

在SRAM里算softmax，更新m和ℓ

逐步累积输出，最后写回HBM

关键：整个过程中，N×N注意力矩阵从来没有完整地出现在HBM上。

Recomputation：反向传播时不再读矩阵

那反向传播怎么办？

反向传播需要用到前向的S和P矩阵。标准做法是把前向的S和P存在HBM上，反向时直接读。

但FlashAttention说：不存了，反向时重新算。

它只存前向的输出O和softmax的统计量m、ℓ，这两个东西很小，O(Nd)的大小。

反向传播时，从HBM读Q、K、V，重新在SRAM里算S和P，然后再算梯度。

虽然多算了一些FLOP，但因为避免了从HBM读N×N矩阵的开销，实际运行时间反而更快。

这在学术上叫selective gradient checkpointing——梯度检查点的一种选择。

4.IO复杂度分析

理论保证

文章给出了一个严格分析。

标准注意力的HBM访问次数是Θ(N²d+N²)。

FlashAttention的HBM访问次数是Θ(N²d/M)，其中M是SRAM的大小。

为什么是N²d/M？

因为SRAM能放下大小为Θ(M)的K、V块，每次能处理Θ(M/d)个K行。

对于N行的Q，需要N/(M/d)=Nd/M次扫描。每次扫描加载O(Nd)数据，所以总共O(N²d/M)次HBM访问。

实际差距

拿A100来算：

d=64（head维度）

M≈100KB（每个SM的SRAM）

标准注意力：O(N²×64)次HBM访问

FlashAttention：O(N²×64/100000)≈O(N²×0.00064)次HBM访问

HBM访问量减少了大约100倍。

虽然实际不可能完全达到理论极限（因为SRAM利用率、块大小选择等因素），但文章实验显示前向传播减少了约8倍，反向传播减少了约7倍，合计约9倍的HBM访问量降低。

下界证明

文章还证明了一个有意思的结论：

对于任何精确注意力算法，在所有可能的SRAM大小范围内，不可能渐近地优于O(N²d/M)的HBM访问下界。

换句话说，FlashAttention在这个意义上是最优的。

5.Block-Sparse FlashAttention

扩展：稀疏注意力

FlashAttention不只是精确注意力，还可以扩展到稀疏注意力。

思路很简单：如果注意力矩阵是块稀疏的（比如某些块全是零），那么在Tiling循环中直接跳过这些块就行。

算法跟FlashAttention几乎一样，只是加了一个if判断：如果当前块M_ij=0，跳过计算。

文章证明了Block-Sparse FlashAttention的HBM访问次数是Θ(N²d·s/M)，其中s是非零块的比例。

s越小，加速越多。

实验显示，在LRA benchmark上，Block-Sparse FlashAttention相对于标准FlashAttention有2.8倍的加速，同时精度相当。

6.实验结果

训练速度

BERT-large：

在MLPerf 1.1上，FlashAttention比Nvidia记录快了15%（从20.0分钟降到17.4分钟）。

GPT-2 small：

比HuggingFace实现快3.5倍（从9.5天降到2.7天）

比Megatron-LM快2.0倍（从4.7天降到2.7天）

GPT-2 medium：

比HuggingFace实现快3.0倍（从21.0天降到6.9天）

比Megatron-LM快1.7倍（从11.5天降到6.9天）

Long-Range Arena：

平均加速2.4倍

模型质量提升

FlashAttention不只是更快，还能训练出更好的模型。

GPT-2长上下文：

用FlashAttention训练GPT-2 small，上下文长度从1K提升到4K，仍然比Megatron的1K版本快30%，且perplexity低了0.7。

长文档分类：

在MIMIC-III（医疗文本分类）和ECtHR（法律判决分类）上，增加序列长度带来显著提升：

MIMIC-III：16K序列比512序列提升4.3分

ECtHR：8K序列比512序列提升8.5分

PathFinder挑战：

Path-X（16K序列）：FlashAttention的Transformer达到61.4%准确率，是第一个在这个任务上超过随机猜测的Transformer

Path-256（64K序列）：Block-Sparse FlashAttention达到63.1%准确率

基准测试

在不同序列长度下：

序列长度128-512：FlashAttention比PyTorch标准实现快2-3倍

序列长度1024-2048：FlashAttention比所有近似注意力方法都快

内存占用：FlashAttention比PyTorch标准实现低20倍，比Linformer低2倍

不同硬件上的表现

A100：2-4倍加速

RTX 3090：2.5-4.5倍加速（HBM带宽更低，加速效果更明显）

T4：加速较少（SRAM更小，块大小需要更小）

7.总结

这篇文章的核心贡献，可以用一句话概括：

写注意力算法的时候，要把GPU内存层级的读写开销也算进去。

这个"IO-Awareness"的思想听起来简单，但在深度学习这个领域里，很少有人真正认真对待过。

大家习惯了看FLOP数，看理论复杂度，看benchmark上的accuracy。

但FLOP不等于wall-clock time，不等于内存使用量，不等于实际训练出来的模型质量。

FlashAttention用两个经典技术——Tiling和Recomputation——把注意力机制的HBM访问量从O(N²)降到了O(N²/M)，在保持精确计算的同时实现了3倍的加速和20倍的内存节省。

而且它不只是更快，还让Transformer第一次真正具备了建模64K长度上下文的能力。

这就是IO-Awareness的力量。