低精度FlashAttention-4：Blackwell上的端到端块级缩放注意力

从FlashAttention-4到低精度版本

FlashAttention系列一直是大模型训练与推理效率优化的核心技术之一。它通过分块计算与在线softmax，避免了将完整注意力矩阵物化到显存中，从而大幅降低内存占用并提升吞吐。最新的低精度FlashAttention-4（下称FA4 MX8）在此基础上更进一步，把注意力的前向与反向传播全面迁移到MXFP8这一微缩放低精度格式上，专为NVIDIA最新的Blackwell架构量身打造。

据原文披露，这项工作在FlashAttention-4的基础上扩展了MXFP8的前向（forward）与反向（backward）实现，在典型的LLM形状下达到了前向2.85 PF/s、反向2 PF/s的性能表现；而在其内部测试形状上，FA4 MX8前向也可以达到2.54 PF/s的算力水平。这些数字直接体现了低精度块级缩放在Blackwell硬件上的潜力。

为什么是MXFP8

MXFP8属于微缩放（Microscaling）浮点格式家族，其核心思想是在一小块数据（block）内共享一个缩放因子，从而在保持8位浮点表达范围的同时，缓解低精度带来的数值动态范围不足问题。相比传统的FP8单一全局或张量级缩放，块级缩放（block-scaled）能够更细粒度地适应张量内部数值分布的差异，这对于注意力这种数值跨度极大的运算尤其关键。

注意力计算中的QK^T点积、softmax指数运算以及与V的加权求和，都可能出现较大的数值波动。若直接使用低精度而不做块级缩放，容易在softmax前后引入显著的精度损失。FA4 MX8选择MXFP8作为端到端的数据格式，正是为了在算力提升与数值稳定性之间取得平衡。

MXFP8在具体实现上通常有两种变体：E4M3（4位指数、3位尾数）和E5M2（5位指数、2位尾数）。前者动态范围较窄但精度更高，适合前向激活值；后者动态范围更宽，更适合梯度表示。微缩放格式（MX格式，Microscaling Format）由OCP Microscaling Formats规范定义，是AMD、英特尔、NVIDIA、微软等公司联合推动的行业标准。其"块级共享缩放因子"的设计通常以32个元素为一个缩放块，每块仅需存储一个FP8指数作为共享缩放因子，额外开销极小。这与之前Transformer Engine中使用的"张量级"FP8缩放相比，粒度细了一到两个数量级，能更准确地捕捉局部数值分布，从而在不引入显著精度损失的前提下完成量化。

端到端块级缩放的意义

以往许多低精度优化只覆盖前向推理，或仅在部分算子上使用低精度。而这项工作强调的是端到端（End-to-End）——前向与反向传播全程采用MXFP8块级缩放。这意味着不仅推理阶段能受益，训练阶段的梯度计算同样可以享受低精度带来的算力红利。

反向传播通常对数值精度更为敏感，因为梯度的累积误差可能在长序列训练中被放大。能够在反向阶段稳定达到2 PF/s，说明块级缩放策略在梯度路径上同样具备可用性。这对于希望用更低成本训练大模型的团队而言，具有相当的现实价值。

针对Blackwell的硬件适配

NVIDIA Blackwell架构原生支持MXFP8等微缩放格式的张量核心运算，这是FA4 MX8能够充分发挥性能的硬件基础。将算法设计与硬件特性紧密绑定，是当前高性能注意力内核优化的普遍趋势——软硬件协同才能把理论峰值算力尽可能地转化为实际吞吐。

从公布的PF/s（PetaFLOPs per second，每秒千万亿次浮点运算）数据看，FA4 MX8已经把注意力内核推进到了PetaFLOP级别的实测区间，这在过去主要依赖BF16/FP16的注意力实现中是难以企及的。

NVIDIA Blackwell架构（以GB200/B200为代表）的第五代张量核心（Tensor Core）原生支持MXFP8矩阵乘法，理论峰值算力相比上一代Hopper（H100）的FP8计算提升约一倍。Hopper虽然也支持FP8，但其FP8缩放仍以张量为粒度，需要在每次矩阵乘前预先计算全局缩放因子，这在注意力这种流式计算中会引入额外的同步开销。Blackwell通过硬件级别对MX块级缩放的原生支持，使得缩放因子可以与矩阵乘运算融合执行，消除了这一瓶颈。这也是FA4 MX8性能数字只能在Blackwell上复现、无法直接迁移到Hopper的根本原因。

在反向传播中，注意力机制需要重新计算或缓存前向的中间结果（如softmax输出矩阵P）来计算梯度。FlashAttention系列通过"重计算"策略（recomputation）在反向时重新执行部分前向计算，以时间换空间，避免存储完整注意力矩阵。将这一过程扩展到MXFP8时，挑战在于梯度路径中的数值范围往往比激活值更难预测，尤其是dQ、dK、dV的计算涉及对softmax梯度的矩阵乘，数值稳定性要求更高。能够在此路径上稳定维持低精度而不出现梯度爆炸或消失，正是块级缩放相比张量级缩放的关键优势所在。

对大模型工程的影响

低精度注意力的持续演进，正在重塑大模型训练与推理的成本结构。更高的算力密度意味着相同硬件可以承载更大的批量、更长的上下文，或者在同等预算下缩短训练周期。对于长上下文场景，注意力往往是计算与内存的主要瓶颈，FA4 MX8这类优化的边际收益会更加明显。

提一嘴，低精度方案的实际收益高度依赖具体模型形状（shape）。原文同时给出了LLM通用形状与内部形状两组数据，两者存在差异（2.85 PF/s vs 2.54 PF/s前向），提醒使用者在落地时仍需针对自身工作负载做实测验证，而非直接套用峰值数字。

小结

FA4 MX8代表了注意力内核低精度化的最新一步：它把MXFP8块级缩放贯穿前向与反向，充分利用Blackwell的硬件能力，在LLM形状上实现了PetaFLOP级的前向与反向吞吐。对于追求极致训练/推理效率的工程团队，这是一个值得密切关注的方向，但落地前的形状适配与精度验证仍不可或缺。