大模型训练(lian)推理神作，又更新了！
主流大模型都在用(yong)的(de)FlashAttention，刚(gang)刚(gang)升级第三代。
时隔(ge)一年，FlashAttention-3已经全方位升级。
训练(lian)速度(du)提升1.5-2倍，FP16下(xia)计算吞吐(tu)量高(gao)达740TFLOPs/s，达理论最大吞吐(tu)量75%，更充分利用(yong)计算资源，此前只能做到35%。
FP8下(xia)速度(du)接近1.2PFLOPs/s！
同时误差也进一步减小，FP8下(xia)的(de)误差比(bi)标准Attention减少2.6倍。
而且(qie)这(zhe)一次，不再是一作Tri Dao单(dan)打独斗，FlashAttention-3直接和(he)英(ying)伟(wei)达、Meta、谷歌等合作，针对(dui)最强(qiang)芯片H100专门做优(you)化。
英(ying)伟(wei)达CUTLASS团(tuan)队和(he)cuDNN团(tuan)队，都直接为(wei)该研(yan)究提供支持(chi)。
同时和(he)前作一样，FlashAttention-3也将开源，PyTorch和(he)Hugging Face中都集成(cheng)。
作者之一Vijay Thakkar激(ji)动(dong)表示：
曾经在FA2发布时，我就说(shuo)过这(zhe)句话。今天，我想(xiang)再说(shuo)一次：
看(kan)到CUTLASS和(he)CuTe被用(yong)来开让Tensor Core大显身手的(de)新算法(fa)，真的(de)泰裤辣。
曾经在FA2发布时，我就说(shuo)过这(zhe)句话。今天，我想(xiang)再说(shuo)一次：
看(kan)到CUTLASS和(he)CuTe被用(yong)来开让Tensor Core大显身手的(de)新算法(fa)，真的(de)泰裤辣。
前Stable Diffusion老板Emad也非常关(guan)注这(zhe)一进展，他推测使用(yong)FlashAttention-3，能将4090的(de)FP8计算吞吐(tu)量推升到700+TFLOPs。
充分利用(yong)Hopper架构特点
自初代发布以来，FlashAttention已经使大模型速度(du)提高(gao)了4-8倍，但(dan)还有一个遗憾：尚未充分利用(yong)现代 GPU。
针对(dui)英(ying)伟(wei)达H100倍后的(de)Hopper架构新特性，三代进行了专门优(you)化。
整(zheng)个系列的(de)核(he)心(xin)思路，是IO感知优(you)化和(he)分块处理。
作者认为(wei)，传(chuan)统的(de)注意力机制效率低的(de)原因，在处理长序(xu)列时，会出现内存(cun)访问操作频繁，以及算法(fa)复杂(za)度(du)指数级暴增这(zhe)两大问题。
FlashAttention通(tong)过IO感知优(you)化将数据从较大但(dan)缓慢的(de)高(gao)带宽内存(cun)（HBM）加载(zai)到较小但(dan)更快的(de)片上内存(cun)（SRAM），在SRAM中执行计算，减少了内存(cun)读写操作的(de)次数。
分块处理则是将输(shu)入(ru)序(xu)列分成(cheng)若干小块，每次只处理一个小块的(de)数据。这(zhe)种方法(fa)使得每次处理的(de)数据量减少，从而降低了内存(cun)使用(yong)和(he)计算复杂(za)度(du)。
这(zhe)样一来，两个关(guan)键问题就得到了解决，这(zhe)两大核(he)心(xin)思想(xiang)也在本次的(de)FlashAttention-3中得到了继承。
但(dan)是，第一代的(de)FlashAttention也遗留下(xia)了并行性不够强(qiang)、工作分区划分不合理，以及非矩阵(zhen)乘法(fa)较多(duo)（GPU计算单(dan)元处理矩阵(zhen)乘法(fa)比(bi)非矩阵(zhen)速度(du)更快）的(de)问题。
针对(dui)这(zhe)一问题，第二代FlashAttention通(tong)过重写softmax，减少了重新缩放(fang)操作、边界检查和(he)因果(guo)屏蔽(bi)操作的(de)次数，使得大部分计算集中在矩阵(zhen)乘法(fa)上。
另外，FlashAttention-2引入(ru)了序(xu)列长度(du)维度(du)上的(de)并行化，并针对(dui)工作在线程块之间的(de)分配进行了优(you)化，GPU利用(yong)效率更高(gao)了。
可以说(shuo)前两代当中，作者一直坚持(chi)着充分利用(yong)硬件(jian)特点这(zhe)一思路，但(dan)站在今天的(de)视角来看(kan)，对(dui)硬件(jian)的(de)挖掘仍然(ran)不够充分。
到了这(zhe)次的(de)FlashAttention-3，由于是直接和(he)英(ying)伟(wei)达官(guan)方合作，对(dui)英(ying)伟(wei)达Hopper架构特点的(de)理解更加透彻，软(ruan)硬件(jian)之间的(de)协同进一步增强(qiang)了。
FlashAttention-3的(de)技术(shu)报告(gao)显示，为(wei)了充分匹配Hopper架构，团(tuan)队主要做了三方面的(de)技术(shu)升级。
首先，Hopper架构的(de)一个重要特点是Tensor Core的(de)异步性，FlashAttention-3针对(dui)性地提出了一种异步方式。
具体(ti)来说(shuo)，FlashAttention-3引入(ru)了一种“生产者（Producer）-消费(fei)者（Consumer）”的(de)编程模型，将注意力的(de)计算划分为(wei)两个角色。
- “生产者”负责将数据从HBM异步加载(zai)到片上共享内存(cun)（SMEM）。这(zhe)个过程主要利用(yong)了Hopper GPU的(de)张量内存(cun)加速器（TMA），可以在不阻塞CUDA核(he)心(xin)的(de)情况下(xia)进行数据传(chuan)输(shu)。
- 消费(fei)者直接从共享内存(cun)读取数据，并使用(yong)Tensor Core执行矩阵(zhen)乘法(fa)等计算密集型任务。由于共享内存(cun)的(de)访问延迟远低于全局(ju)内存(cun)，消费(fei)者可以快速获取所(suo)需数据，提升计算效率。
为(wei)了实现角色的(de)划分，作者引入(ru)了warp专门化技术(shu)，用(yong)不同的(de)warp分别匹配生产者和(he)消费(fei)者，让两者可以并行执行。
这(zhe)其中利用(yong)了Hopper架构的(de)动(dong)态warp寄存(cun)器分配特性，通(tong)过setmaxnreg指令优(you)化了寄存(cun)器资源的(de)利用(yong)。
为(wei)了进一步提高(gao)GPU的(de)利用(yong)率，作者又提出了一种“乒乓调度(du)”策略，让一个warp组执行矩阵(zhen)乘法(fa)时，另一个warp组执行softmax，从而实现计算的(de)重叠。
具体(ti)讲，FlashAttention-3使用(yong)CUDA的(de)同步原语控制不同warp组之间的(de)执行顺序(xu)，让不同warp组分别执行两种运算，然(ran)后像乒乓球一样交替运行。
第二大技术(shu)特点，是warp组内部GEMMs和(he)softmax的(de)重叠，核(he)心(xin)奥义是重新安排计算的(de)执行顺序(xu)以提高(gao)GPU利用(yong)率。
与乒乓调度(du)不同，这(zhe)里的(de)计算重排处理的(de)是warp组内部的(de)重叠，而乒乓调度(du)更关(guan)注组间协调。
实现方式上，FlashAttention-3提出了一种两阶段GEMM-softmax流水线方案，以打破不同操作之间的(de)数据依赖。
- 第一阶段，当前迭代（iteration）的(de)softmax操作与下(xia)一个迭代的(de)Q·K^T矩阵(zhen)乘法(fa)重叠执行。
- 第二阶段，当前迭代的(de)P·V矩阵(zhen)乘法(fa)与下(xia)一个迭代的(de)softmax操作重叠执行。
通(tong)过引入(ru)额外的(de)寄存(cun)器和(he)共享内存(cun)缓冲区，FlashAttention-3实现了跨迭代的(de)数据传(chuan)递和(he)重用(yong)。
在每个迭代中，Q·K^T的(de)结果(guo)首先存(cun)储在名为(wei)S_cur的(de)缓冲区中，用(yong)于当前迭代的(de)softmax计算，同时异步执行下(xia)一个迭代的(de)Q·K^T矩阵(zhen)乘法(fa)，结果(guo)存(cun)储在名为(wei)S_next的(de)缓冲区中。
在执行当前迭代的(de)P·V矩阵(zhen)乘法(fa)时,异步执行下(xia)一个迭代的(de)softmax操作，并更新S_cur和(he)S_next缓冲区。
第三项更新，是用(yong)更低的(de)FP8精度(du)替代FP16。
实际上，降低数值精度(du)是一种常见(jian)的(de)优(you)化策略，可以显著提高(gao)GPU的(de)计算吞吐(tu)量和(he)能效，Hopper GPU也引入(ru)了FP8精度(du)的(de)Tensor Core支持(chi)。
但(dan)是，直接将注意力计算从FP16转(zhuan)换为(wei)FP8可能会引入(ru)较大的(de)精度(du)损失。
另外，FP8 Tensor Core对(dui)输(shu)入(ru)数据的(de)布局(ju)也有特定的(de)要求（K维度(du)连续(xu)），不幸的(de)是，注意力计算中的(de)输(shu)入(ru)数据存(cun)储格式（头维度(du)连续(xu)）并不符合这(zhe)样的(de)要求。
所(suo)以FlashAttention-3首先引入(ru)了一系列内存(cun)布局(ju)转(zhuan)换技术(shu)，动(dong)态转(zhuan)置V矩阵(zhen)的(de)块，改变其连续(xu)方式，从而适配FP8 Tensor Core的(de)布局(ju)要求。
在此基础之上，为(wei)了获得更高(gao)的(de)计算精度(du)，FlashAttention-3又采用(yong)了分块量化和(he)非相干处理技术(shu)。
传(chuan)统的(de)量化方法(fa)通(tong)常对(dui)整(zheng)个矩阵(zhen)使用(yong)一个统一的(de)缩放(fang)因子（per-tensor quantization），无法(fa)很(hen)好地适应不同区域的(de)数值范围。
FlashAttention-3则采用(yong)了分块量化（block-wise quantization）的(de)策略，为(wei)每个块单(dan)独设置缩放(fang)因子，更好地捕捉局(ju)部的(de)数值分布。
非相干处理（incoherent processing）技术(shu)则是通(tong)过随机正交矩阵(zhen)对(dui)输(shu)入(ru)数据进行旋转(zhuan)，破坏(huai)不同块之间的(de)相干性，减少量化误差的(de)传(chuan)播。
这(zhe)两项技术(shu)的(de)结合使得FlashAttention-3在FP8精度(du)下(xia)取得了更高(gao)的(de)计算精度(du)，显著优(you)于传(chuan)统的(de)量化方法(fa)。
结果(guo)，与基于传(chuan)统量化方法(fa)的(de)FP8实现相比(bi)，FlashAttention-3的(de)使得精度(du)提高(gao)了2.6倍。
比(bi)标准Attention快16倍
以上就是FlashAttention-3在充分研(yan)究Hopper架构特点后做出的(de)三大更新，针对(dui)更新后的(de)表现，作者主要进行了3方面测试。
- 注意力基准测试
- 消融实验
- FP8注意力准确(que)性测试
首先来看(kan)注意力基准测试。
通(tong)过改变序(xu)列长度(du)（512、1k、……16k），并设置批大小以确(que)保总token数为(wei)16k。研(yan)究人员将隐藏维度(du)设置为(wei)2048，头维度(du)设置为(wei)64、128或(huo)258，计算前向传(chuan)播、后向传(chuan)播。
对(dui)比(bi)标准Attention、FlashAttention-2、Triton、cuDNN和(he)FlashAttention-3，在H100 80GB SXM5上FP16的(de)运行时间。
FlashAttention-3的(de)前向传(chuan)播比(bi)FlashAttention-2快1.5-2倍，后向传(chuan)播快1.5-1.75倍。
与标准Attention相比(bi)，FlashAttention-3的(de)速度(du)快了3-16倍。
对(dui)于中长序(xu)列（1k以上），FlashAttention-3甚至超过了专门为(wei)H100优(you)化的(de)cuDNN。
在消融实验中，通(tong)过对(dui)非因果(guo)FP16 FlashAttention-3进行了2阶段WGMMA-softmax流水线和(he)warp特殊化的(de)消融研(yan)究，参数固定为(wei){batch, seqlen, nheads, hdim} = {4, 8448, 16, 128}。
结果(guo)证实，FlashAttention-3改进带来了显著加速，从570提升到661。
另外，因为(wei)对(dui)FlashAttention的(de)数值误差感兴(xing)趣，研(yan)究团(tuan)队还将FlashAttention-2、FlashAttention-3和(he)标准Attention进行了比(bi)较。
为(wei)了模拟LLMs中的(de)异常特征和(he)激(ji)活(huo)，研(yan)究团(tuan)队生成(cheng)了Q、K、V的(de)条目(mu)，分布为(wei)：N(0,1)+N(0,100)·Bernoulli(0.001)
也就是说(shuo)，每个条目(mu)都服从均值为(wei)0、标准差为(wei)1的(de)正态分布，但(dan)对(dui)于0.1%的(de)条目(mu)，增加了一个独立的(de)项，其标准差为(wei)10。然(ran)后测量均方根误差（RMSE）。
结果(guo)显示，在FP16中，由于中间结果(guo)（softmax）保留在FP32中，FlashAttention-2和(he)FlashAttention-3的(de)RMSE比(bi)标准Attention减少1.7倍。
FP8的(de)标准Attention使用(yong)每个张量的(de)缩放(fang)，matmul累加器在FP32中，中间softmax结果(guo)保留在FP16中。由于块量化和(he)非相干处理，FP8中的(de)FlashAttention-3比(bi)这(zhe)个基线更准确(que)2.6倍。
最后，论文还表示目(mu)前工作专注于Hopper架构，后续(xu)将推广到其他硬件(jian)。
除了英(ying)伟(wei)达为(wei)研(yan)究提供了技术(shu)支持(chi)外，Meta、Together AI和(he)普林(lin)斯顿(dun)大学为(wei)研(yan)究提供了计算支持(chi)。
本文来源：量子位，原文标题：《H100利用(yong)率飙升至75%！英(ying)伟(wei)达亲自下(xia)场FlashAttention三代升级，比(bi)标准注意力快16倍》
*免责声明(ming)：文章内容仅供参考，不构成(cheng)投资建(jian)议
*风险提示：股市有风险，入(ru)市需谨慎