闻乐 发自 凹非寺
量子位 | 公众号 QbitAI
所有使用英伟达Blackwell B200 GPU的用户，可能正为未被释放的算力买单。
普林斯顿大学联合Meta、Together AI及英伟达等机构研究发现：由于软硬件协同不足，B200在注意力计算中平均浪费约60%的理论算力。
全新优化算法FlashAttention-4应运而生，专为Blackwell架构设计，将注意力计算利用率从行业普遍的20%–30%提升至71%。
该算法由Tri Dao领衔开发，参与方包括Meta、Together AI及英伟达。
Blackwell B200有力使不出
英伟达Blackwell B200张量核心算力达2.25 PFLOPS，为Hopper H100的2倍，理论上可大幅提升注意力计算效率。
但实际性能受限于硬件设计的结构性失衡：
- MUFU指数运算单元吞吐量未升级，仍与Hopper持平；
- 共享内存带宽亦未同步提升。
导致注意力计算中，原本耗时较长的矩阵乘法已非瓶颈，反而是共享内存读写和指数运算耗时增加25%–60%，Tensor Core长期闲置。
结果：超六成算力因核心单元与配套单元脱节而被浪费。
FlashAttention-4三招破解瓶颈
第一招：重构指数运算与内存访问
通过多项式近似实现软件模拟指数函数，将原属MUFU的计算任务迁移至高吞吐FMA单元，在保证精度前提下显著提速；
引入条件性softmax rescaling策略，仅在必要时执行缩放，跳过冗余计算；
采用2-CTA MMA模式，双计算单元协同完成矩阵乘法，各自仅加载半量数据，使共享内存读写量减半，并减少原子操作，缓解带宽压力。
第二招：全异步流水线调度
深度适配Blackwell的全异步MMA与新增张量内存（TMEM），重构前向与反向传播流水线，实现softmax计算与矩阵乘法完全重叠。
当张量核心处理一个数据块时，其余硬件资源同步执行另一块的softmax计算，消除空闲周期。
第三招：面向演进的弹性设计
针对下一代B300/GB300 GPU指数单元吞吐量已翻倍至32 ops/clock/SM的趋势，FlashAttention-4保留软件模拟方案的动态切换能力，确保持续适配硬件迭代。
告别C++，编译速度提升30倍
FlashAttention-4全面转向Python生态，基于CuTe-DSL框架开发，实现零C++代码编写。
编译效率大幅跃升：
- 前向传播内核：55秒 → 2.5秒（提速22倍）；
- 反向传播内核：45秒 → 1.4秒（提速32倍）；
- 整体最高提速达30倍。
B200实测显示，前向传播峰值达1613 TFLOPS/s，对应71%理论利用率。
相较主流方案优势显著：
- 比cuDNN 9.13快1.1–1.3倍；
- 比Triton快2.1–2.7倍；
- 在长序列、因果掩码等大模型核心场景中性能增益更明显。
One More Thing
论文指出，cuDNN 9.13起已反向集成FlashAttention-4多项核心技术。
