什么是 FlashAttention-2？一种改进 GPU 工作划分的注意力计算加速方法

FlashAttention-2 是一种针对 Transformer 注意力层（attention layer）的 GPU 计算优化方案，由 Tri Dao 在论文《FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning》中提出，2023 年 7 月 17 日提交至 arXiv。它要解决的问题是：注意力层的运行时间与显存占用随序列长度呈二次增长，是模型扩展到长序列时的主要瓶颈；而前代 FlashAttention 虽然已经做到线性显存占用与相对优化基线 2 至 4 倍的加速，却仍只能达到理论最大 FLOPs/s 的 25% 至 40%，远不及优化后的矩阵乘（GEMM）运算。FlashAttention-2 通过改进 GPU 上线程块（thread block）与 warp 之间的工作划分，缩小这一效率差距。

为什么重要

论文指出，把 Transformer 扩展到更长序列长度是过去数年的重要问题，它有望改善语言建模与高分辨率图像理解的表现，并解锁代码、音频、视频生成等新应用。但注意力层是长序列扩展的主要瓶颈：其运行时间和显存都随序列长度二次增长。

FlashAttention 的做法是利用 GPU 显存层次的不对称性，在不做近似的前提下带来显著的显存节省（从二次降为线性）和运行时间加速（相对优化基线 2 至 4 倍）。然而论文同时观察到，FlashAttention 仍然远不如优化后的矩阵乘运算快，只达到理论最大 FLOPs/s 的 25% 至 40%。论文把这种低效归因于 GPU 上不同线程块与 warp 之间的工作划分不够理想，导致占用率偏低，或产生不必要的共享内存读写。换言之，问题不在于算法是否需要近似，而在于并行任务在硬件上的切分方式没有充分利用算力。FlashAttention-2 的价值就在于针对这一层「工程划分」问题做改进，使注意力计算效率接近矩阵乘运算。

工作机制

根据论文摘要，FlashAttention-2 在 FlashAttention 的基础上提出更好的工作划分，具体包含三项改动：

1. 调整算法以减少非矩阵乘 FLOPs。 论文对算法做了微调，降低不经过矩阵乘单元完成的那部分浮点运算数量，使更多计算落在 GPU 高效的矩阵乘路径上。
2. 把注意力计算并行化到更多线程块。 即使对于单个注意力头（head），也把计算拆分到不同线程块上并行执行，以提高 GPU 的占用率（occupancy）。
3. 在线程块内部重新分配 warp 间的工作。 在每个线程块内，把任务在 warp 之间重新划分，以减少经由共享内存（shared memory）的通信开销。

论文报告，这三项改动合计带来相对 FlashAttention 约 2 倍的加速，在 A100 上达到理论最大 FLOPs/s 的 50% 至 73%，接近 GEMM 运算的效率。论文还做了端到端验证：在训练 GPT 风格模型时，FlashAttention-2 在单块 A100 GPU 上达到最高 225 TFLOPs/s 的训练速度，模型 FLOPs 利用率为 72%。需要说明的是，以上数字均为该论文自述的实验结果，针对其设定的硬件与模型配置，并非跨平台、跨模型的通用保证。

典型例子

论文中给出的验证场景是端到端训练 GPT 风格（GPT-style）模型：把 FlashAttention-2 用于完整训练流程，而不是单独测一个注意力算子，测得单块 A100 GPU 上最高 225 TFLOPs/s 的训练速度与 72% 的模型 FLOPs 利用率。这一例子说明论文关注的不只是算子层面的微基准，也包括它在真实训练负载中的表现。

论文的另一组对照对象是前代 FlashAttention 与优化后的矩阵乘基线：前者被用来衡量加速倍数（约 2 倍），后者被用来衡量效率差距（从 25% 至 40% 提升到 50% 至 73% 的理论最大 FLOPs/s）。此外，论文的动机部分提到长序列 Transformer 在语言建模、高分辨率图像理解以及代码、音频、视频生成方面的潜在应用，但这些属于问题背景的陈述，论文并未在这些任务上逐一给出 FlashAttention-2 的实测结果。

边界与常见误解

首先，FlashAttention-2 不是一种新的注意力数学定义，也不引入近似。论文明确把「无近似」作为 FlashAttention 系列的性质来对比，FlashAttention-2 的改动集中在工作划分与并行方式上，而不是改变注意力计算的结果。把它理解成「一种近似注意力」是常见误解。

其次，论文报告的加速与利用率数字有明确的前提：约 2 倍加速是相对 FlashAttention 而言，50% 至 73% 的理论最大 FLOPs/s 与 225 TFLOPs/s、72% 模型 FLOPs 利用率是在 A100 GPU 上测得。这些数字不能被外推为在其他硬件、其他序列长度或其他模型规模上的表现。论文没有声称它适用于所有 GPU 架构，也没有给出跨硬件的对比数据。

第三，FlashAttention-2 解决的是 GPU 上线程块与 warp 的工作划分问题，其收益依赖于具体的硬件执行模型与内核实现。论文把 FlashAttention 的低效归因于占用率偏低与共享内存读写浪费，这意味着如果实现层面不遵循相应的划分策略，理论上就得不到论文所述的效率。对于不使用 GPU 或计算模式差异较大的场景，论文未做讨论。

最后，需要注意信息来源：上述关于动机、机制与实验数字的表述均来自 Tri Dao 的这篇论文，属于作者的主张与其实验设置下的观测，不应被当作整个行业已经验证的共识。

参考资料

• arXiv:2307.08691 FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning