所有用英伟达 Blackwell B200 的人，都在花冤枉钱？？

普林斯顿大学等联合团队指出，这款 GPU 居然因为软硬件适配问题白白浪费了 60% 的计算资源。

算力浪费了，咋办呢——FlashAttention-4给出了答案。

这款专为 Blackwell 架构 GPU 量身打造的注意力算法，一举将利用率从行业普遍的 20%-30% 推至71%。

FlashAttention-4 由 Tri Dao 领衔、携手 Meta、Together AI 等团队共同研发。

嗯，英伟达自己也参与其中了……

Blackwell B200 有力使不出

英伟达 Blackwell B200 作为新一代数据中心 GPU，其 tensor core 张量核心算力达到 2.25 PFLOPS，是上一代 Hopper H100 的 2 倍。

理论上能让注意力计算的速度实现跨越式提升。

但理想很丰满……

这款 GPU 发生了严重的偏科。

核心算力猛增的同时，关键的配套计算单元却原地踏步。

其中，负责指数运算的 MUFU 单元吞吐量与 Hopper 架构完全一致，没有任何提升；

共享内存的带宽也保持原样，并未跟随张量核心同步升级。

这一硬件设计的不对称性，直接导致了性能瓶颈的反转。

在大模型核心的注意力计算负载中，原本的性能瓶颈矩阵乘法，如今耗时远低于辅助环节，共享内存的读写操作和指数运算的耗时，反而比矩阵乘法多出 25%-60%。

算力翻倍的 Tensor Core 长期处于等待状态，大量计算资源就这么被闲置了。

于是，大量开发者花费重金部署的 B200 GPU，因核心算力与配套单元的脱节，超六成资源被白白浪费。

算力翻倍？

No！明明是有力使不出……

FlashAttention-4 三招破解瓶颈

针对 Blackwell GPU 的偏科问题，FlashAttention-4 量身打造了三大优化策略。

第一招，多管齐下化解指数运算与内存读写难题。

团队一方面通过软件模拟指数函数，借助多项式近似的方法，让高速的 FMA 计算单元参与到原本由 MUFU 单元负责的指数运算中，大幅提升指数计算的吞吐量；

同时通过混合硬件计算与软件模拟的方式，在提速的同时保证计算精度。

另一方面推出条件性 softmax rescaling策略，仅在必要时执行 softmax 的缩放操作，直接跳过大量无用的计算步骤，减少非矩阵乘法的运算量。

此外，团队充分利用 Blackwell 架构的2-CTA MMA模式，让两个计算单元搭档完成矩阵运算，各自仅加载一半的运算数据。

这就将共享内存的读写量直接砍半，同时还减少了后续的原子操作，从根源上缓解共享内存的带宽压力。

第二招，重构计算流水线，实现算力的并行最大化。

FlashAttention-4 深度适配 Blackwell 架构的全异步 MMA 操作和新增的张量内存 TMEM，重新设计了注意力计算的前向和反向流水线。

让 softmax 计算与矩阵乘法这两个核心环节实现完全的计算重叠。

当硬件的张量核心在处理一个矩阵块时，另一部分硬件资源可同时对另一个数据块执行 softmax 计算，避免硬件算力的空闲。

第三招，兼顾硬件迭代，为下一代 GPU 预留优化空间。

研发团队同时考虑到 Blackwell 架构的硬件升级趋势，目前 B300/GB300 GPU 的指数运算单元吞吐量已翻倍至32 ops/clock/SM。

针对这一变化，团队明确表示，FlashAttention-4 当前的软件模拟指数运算方案，在下一代硬件上会根据实际性能表现重新权衡，确保算法能持续适配硬件的迭代升级。

告别 C++，编译速度狂飙 30 倍

除了算法层的深度优化，FlashAttention-4 在开发层面也带来了变化。

与此前基于 C++ 模板开发的 FlashAttention-3 不同，FlashAttention-4 的全部代码基于 Python 的领域专用版本 CuTe-DSL 框架编写，实现了零 C++ 代码开发。

这一设计带来的是编译的效率跃升。

前向传播内核的编译时间从 FlashAttention-3 的 55 秒缩短至 2.5 秒，提速 22 倍；

反向传播的编译时间从 45 秒降至 1.4 秒，提速 32 倍，整体编译速度最高狂飙 30 倍。

在 B200 GPU 上的实测数据显示，其前向传播算力最高达到1613 TFLOPS/s，一举实现 71% 的理论峰值利用率。

对比主流的计算框架，FlashAttention-4 的优势也比较明显。

比英伟达官方的 cuDNN 9.13 快 1.1-1.3 倍，比常用的 Triton 框架快 2.1-2.7 倍。

且在长序列、因果掩码等大模型训练推理的核心场景中，性能优势更为突出。

One More Thing

论文还指出，cuDNN 从 9.13 版本开始就已经开始反向吸收了 FA4 的核心技术。

看来，英伟达自己也忍不住抄作业了（doge）。

论文地址：https://arxiv.org/abs/2603.05451

参考链接：https://x.com/alex_prompter/status/2033885345935462853?s=20

一键三连「点赞」「转发」「小心心」

欢迎在评论区留下你的想法！

—  完  —

点亮星标

科技前沿进展每日见