《FlashAttention-2：更好并行与工作划分的更快注意力》

问题：注意力计算没有喂饱 GPU

注意力（attention）让 Transformer 模型中的每个位置去参考序列里其他位置的信息，再据此更新自己的表示。它出现在每一层里，因此注意力这一步算得快不快，直接决定了这类模型训练和推理的速度。Tri Dao 于 2023 年 7 月 17 日把 FlashAttention-2 的论文提交到 arXiv ^[1]。

衡量硬件速度常用 FLOPs/s，即每秒浮点运算次数。每块 GPU 都有一个理论峰值，也就是只做运算、不做别的事情时能达到的上限。这条工作线上更早的 FlashAttention 只达到了这个上限的 25–40% ^[1]。论文认为，问题不是算得太多，而是算得不够整齐：线程块（thread block）与线程束（warp）之间的工作划分不理想 ^[1]。线程块是硬件调度的单位，一个线程块里的各个线程束同步执行；只要有的块早早做完、有的束在等别人，运算单元就会空转。

FlashAttention-2 改了什么

论文的改动有三处 ^[1]。第一，减少非矩阵乘（non-matmul）的浮点运算——指数、取最大值、求和这类操作不属于矩阵乘法，用不上专门跑矩阵乘法的单元。第二，即使只有一个注意力头，也把它的计算切分到多个线程块上并行做，让单个头的工作铺开到更多硬件上，提高占用率（occupancy），也就是真正在忙的那部分硬件的比例。第三，在每个线程块内部把工作分给各个线程束，让线程束之间少通过共享内存（shared memory）交换数据；共享内存是片上的小块存储，它的读写本身也会成为瓶颈。

结果

三处改动合起来，相比 FlashAttention 大约有 2 倍加速，在 A100 显卡上达到理论最大 FLOPs/s 的 50–73% ^[1]。论文还按前向和反向分别给了数字：前向最多达到理论最大吞吐的 73%，反向最多 63% ^[1]。反向传播要算梯度，通常比前向更难保持效率。

对使用者更直接的数字是端到端训练：用它训练 GPT 风格的模型时，每块 A100 上的训练速度最高可达 225 TFLOPs/s ^[1]。这个数字覆盖的是完整的训练步骤，而不只是单独一个注意力算子，所以它描述的是整个模型，而不是注意力操作本身。

为什么重要

注意力在每一层都要调用一次，把它算快一点，等于把每一步训练、每一个生成 token 的时间都缩短一点。这些百分比也是一种有用的度量：它说明这套实现离硬件上限还有多远，而不只是说明它比上一版快多少。从 25–40% 提到 50–73%，意味着同一块显卡在同样的工作上能提供大约两倍的有效算力 ^[1]。