Blackwell上的低精度注意力：FlashAttention-4迈向端到端MXFP8

导语

低精度计算的难点，往往不在于把数据类型从BF16改成FP8，而在于让量化、缩放因子搬运和矩阵乘法能够融入完整的计算流水线。PyTorch团队近日介绍了面向Blackwell GPU的低精度FlashAttention-4实现，将MXFP8支持扩展到注意力前向与反向，并进一步接入带有融合算子和变长序列处理的训练模块。

核心进展

• 前向与反向均支持MXFP8。 在LLM形状上，前向性能最高达到2.85 PFLOPS，反向达到2 PFLOPS；在团队内部形状上，前向达到2.54 PFLOPS、反向达到1.58 PFLOPS。相比BF16，内部测试中的前向最高提升1.6倍，反向最高提升1.52倍。
• 让缩放因子适应Blackwell的TMEM。 Blackwell的块缩放矩阵乘法需要额外管理scale factor，但FlashAttention-4原有流程已经基本占满512列TMEM。实现通过复用不同阶段的TMEM区域，并增加必要的屏障同步，在不显著扩大资源占用的情况下放入缩放因子。
• 把量化放进生产者算子。 团队实现了融合RMSNorm与量化、融合GEMM与量化的内核，使上游算子可以直接生成FP8输出和不同布局的缩放因子，减少独立格式转换带来的开销。
• 优化在线计算的P量化。 注意力中的softmax结果P原本需要从FP32转为适合后续PV乘法的低精度格式。实现利用softmax阶段已经计算出的行最大值，并结合Blackwell上的指令级优化，降低MXFP8转换成本。
• 面向变长输入采用零Gather设计。 在jagged模块中，FP8激活尽量保留在未填充的位置，仅对规模更小的缩放因子进行散列、填充和重排，使数据更适合通过TMA搬运到张量核心。

为什么重要

Blackwell的块缩放MMA原生支持MXFP8、MXFP6、MXFP4和NVFP4等微缩放格式，理论吞吐可达到BF16的数倍。但如果量化发生在主计算之外，格式转换、scale factor搬运和同步就可能抵消硬件收益。该工作说明，低精度注意力需要从数据布局、内存分配、异步执行到上下游算子进行整体设计，而不是只替换一个数据类型。

另一个值得注意的变化是性能瓶颈的转移。随着矩阵乘法加速，softmax等仍受特殊函数单元和同步流程限制的阶段更容易暴露。团队因此调整KV展开方式，并重新安排屏障等待位置，以减少tile边界处的流水线空泡。

从应用角度看，这套实现已被接入Meta内部的GEM训练场景，并开放了相关代码。不过，文中数据来自特定硬件、算子实现和测试形状，不能直接等同于所有模型或部署环境的普遍加速。它更重要的价值在于展示了一条路径：将FP8从单个矩阵乘法的优化，推进到覆盖注意力前后处理、反向计算和变长数据组织的端到端训练组件。

来源：PyTorch Blog