FlashAttention压到4位：矩阵乘快一倍多，训练却只快14%

注意力慢在哪儿：先学会 FlashAttention 的语言

这篇论文在 GB200 上把注意力矩阵乘的操作数全部压到 4 位浮点（FP4），随即出现一道反常的落差：FP4 张量核让矩阵乘吞吐翻了一倍，前向内核确实也快了近两倍，S8192/H64 形状下耗时从 1.611 毫秒压到 0.758 毫秒；可同一套量化端到端一算，完整训练只快了约 1.14 倍。多出来的提速去了哪里？为什么会有这道落差？先回到 FlashAttention 怎样处理非矩阵乘瓶颈。这一节接下来的内容，都是这道题目的背景。

先补 FlashAttention 的语言。一段 8192 个 token 的文本走进注意力，要算的其实只有两件事：query 和 key 的相似度矩阵，以及拿相似度当权重去加权 value。可这两步中间夹着一张 N 乘 N 的中间矩阵，N 是序列长度，算下来是 6700 多万个数。把这张矩阵写进显存再读出来，比矩阵乘本身贵得多，因为显存带宽一直是 GPU 上最稀缺的资源。所以 FlashAttention 系列存在了近四年，核心只做一件事：别把中间矩阵物化出来。这套「躲中间矩阵」的功夫，是解释落差的第一层背景：瓶颈从来不在算术本身，而在算术周围。

FlashAttention-2 的答案叫 tiling（分块）。K 和 V 被切成小块，一块一块算注意力分数，结果在片上逐步累积，S 和 P 这两张中间矩阵从头到尾不落盘，内存占用从 O(N²) 压到 O(N)。这中间有个巧劲叫 online softmax：处理每个分块时只随身维护两个运行统计量，一个是当前遇到的最大值 m，一个是未归一化的加权和 ℓ，每吃进一块就把之前的输出重缩放一次，最后得到的结果和完整算一遍 softmax 分毫不差。图里这条数据流把这段话画了出来：K、V 各分两块，逐块计算、逐块重缩放，最花心思的正是矩阵乘旁边那道 softmax。

反向传播是这套设计里更重的部分。为了省内存，backward 不保存完整的概率矩阵 P，而是靠每行的 LSE 把 P 重算出来，一口气要算五个矩阵积，片上要塞下前向两倍的值。这条伏笔后面还会回来。

图注：FlashAttention 前向数据流，K 与 V 各分成两块，逐块计算注意力分数并重缩放输出，全程不把中间矩阵 S、P 写回显存。图中省略了 softmax 减去行最大值的步骤，那是 online softmax 维持数值稳定的关键一环。FA-2

打个比方：标准注意力像把一整张「每个字和每个字的关系表」打印出来再读回去，分块算法像一边读一边在手边攒一张小表，只随身带一行统计量，读完全文答案自然就齐了。

这个家族还有个规律：硬件每换一代，内核就多一套把 softmax 藏起来的功夫，这是落差背后的第二层背景。H100 上 FlashAttention-2 只跑到约 35% 的利用率，因为指数运算太慢了，FP16 矩阵乘大约每秒 989 TFLOPs，特殊函数单元只有约 3.9 TFLOPs，差了约 256 倍。FlashAttention-3 在 Hopper 上引入 warp specialization，把 softmax 夹在两个 warpgroup 之间错峰执行，一个在做 softmax 时另一个在跑矩阵乘，前向从约 570 提到 620 到 640 TFLOPs。图里同色代表同一轮迭代，两个 warpgroup 的角色随迭代交替。

图注：两个 warpgroup 的 pingpong 调度示意，一个在做 softmax 时另一个在跑矩阵乘，同一颜色代表同一轮迭代，角色随迭代交替，softmax 被藏进对方矩阵乘的时间片里。FA-3

到现在为止，谁最慢，谁就决定注意力的快慢，而矩阵乘从来不是最慢的那个。FA-3 的论文里早就写过一句伏笔：FP8 in Hopper and FP4 in Blackwell。到 Blackwell 这一代，NVIDIA 直接给了更狠的礼物：4 位浮点张量核，矩阵乘吞吐理论上再翻一倍。问题是，注意力里最慢的环节一个都没动，开头那道落差在这里显出第一层：内核提速是真的，可瓶颈根本没被碰到。要看清这篇论文在跟什么较劲，得先认识 4 位浮点本身。

FP4 只有约 15 个可表示的数，凭什么进注意力？

4 位浮点（FP4）把 1 个符号位、2 个指数位、1 个尾数位打包成一个数。能表示的非负幅值只有 0、二分之一、1、二分之三、2、3、4、6 这 8 个，加上负侧一共约 15 个值，范围大概在 ±6。这么粗的格子，直接套在任何数值分布上都会砸出大坑，所以 4 位格式的生存法则只有一个：靠缩放因子 scale 把动态范围「接」回来，重建公式是 x 等于量化值乘上 s。像只有 4 档灰度的照片配一个曝光旋钮，旋钮把整幅画面的亮度范围扳到能用的区间。

旋钮隔多远拧一次、拧多细，就是三种 4 位格式的分水岭，也是后文所有分歧的源头。

MXFP4 的 E8M0 是个纯指数，一块 32 个元素共享一个 2 的幂乘子，硬件成本低但粗，按元素算大约 4.25 位。NVFP4 把块砍到 16，共享 scale 换成 FP8 E4M3 小数，还能再叠一层 per-tensor 的 FP32 二级缩放，匹配局部动态范围的机会多了一倍。Blackwell 的第五代张量核原生支持 NVFP4，自动处理元素分组、动态缩放和 4 位矩阵运算。算上每组 1 个 FP8 scale 和每张量 1 个 FP32 标量，NVFP4 每值摊下来约 4.5 位。NVFP4 博客

这里有个容易被忽略的伏笔：OCP 的 MX 规范把 MXFP4 定义为块 32 加 E8M0 共享指数，但规范论文从来没做过 MXFP4 的权重加激活加梯度全 4 位训练，MXFP4 只以权重格式出现过。MX 规范 它到底行不行，要等这篇论文把它推到注意力的概率和值上，答案才会揭晓。

读这篇论文前，还得先拆一个命名陷阱。标题里有两个 4，含义完全不同。FlashAttention-4 的 4 是版本号，指 Dao-AILab 的第四代 BF16 Blackwell 内核。FP4 的 4 是位宽。本文继承的是 HAO AI Lab 的执行调度，不是 Dao-AILab 的官方后继，后面还会遇到 Attn-QAT 的 FA4-FP4 kernel，那又是 FA-4 的 NVFP4 扩展。三个「4」各说各话，看到名字别急着对号入座。FA-4

格式懂了，问题回到硬件上：Blackwell 的 FP4 张量核把矩阵乘吞吐翻了一倍多，注意力为什么没有自动跟着快？

矩阵乘翻倍了，注意力没跟着快：Direct-P 的绕行方案

先看硬件账。B200 的 BF16 矩阵乘吞吐是 Hopper 的 2 倍，每 SM 每时钟 8192 次运算对 4096 次，论文口径的实测峰值约 1613 TFLOP/s，达到理论值的约 71%。可同一张卡上，指数单元还是每 SM 每时钟 16 次，shared memory 带宽纹丝不动。论文管这叫异步硬件缩放（asymmetric hardware scaling）：矩阵乘单元激进翻倍，其余单元原地踏步。注意力内部有两个矩阵乘 S 等于 QK 转置除以根号 d、O 等于 PV，中间隔着 softmax。矩阵乘翻倍之后，中间那串「取最大值、重缩放、算概率、打包、发布」的动作成了新瓶颈。锚论文把事情说得很直白：Blackwell 的 4 位张量核不会自动让注意力变快，因为矩阵乘积缩小之后，softmax 转换和片上依赖开始主导一切。FP4 FlashAttention-4

把这条流水线排出来看一眼，就知道代价有多大。图里是 FA-4 为藏住 softmax 编排的前向流水线：高、低两个 128 token 的 Q tile 轮流进场，两个 softmax warpgroup 加一个校正 warpgroup 穿插调度，分数和概率经 TMEM 流转，概率分阶段触发 PV 矩阵乘。这张图的复杂程度本身就是证据，翻倍的是矩阵乘，不是它周围的世界。

图注：FlashAttention-4 前向流水线，高、低两个 128-token 的 Q tile 轮流进入，两个 softmax warpgroup 与一个 correction warpgroup 穿插调度，S 与 P 经 TMEM 流转，P 分阶段触发 PV 矩阵乘，所有阶段重叠执行。FA-4

这篇论文的解法叫 Direct-P，思路是把「算概率」从算术问题改成分类问题。传统做法先算精确的 2 的 x 次方，再把它量化成 4 位，论文直说这没必要：计算精确指数再丢弃几乎全部精度，是不必要的。Direct-P 对就绪的 FP32 分数片段直接套一个仿射映射，取 max(0, Ax + B)，fast 档拟合的系数是 A 等于 1.50、B 等于 1.20，把分数分类到 E2M1 的码本格子里，只有少数位置动用原生指数指令。Q 和 K 用 NVFP4，P 和 V 用 MXFP4，每个 N32 概率块配一个 E8M0 的 2 的幂 scale。整条链路是下面这张图的样子。

省掉的精度要靠三处细节补回来。第一，分母归一化必须和 PV 消费的是同一批舍入码、同一个块 scale，否则会出现「E2M1 分子除以无关近似 FP32 指数和」这种错位。第二，有个专门对付极端层的 sampled guard，只在 logit 特别极端的层启动，比如 Wan 模型晚期某几层的 BF16 logits 大于 500 甚至 1000，先采样锚点、用编译期 margin 保证分母不下溢。第三，两个操作点供选择。fast 档 12 个 K/V 阶段，GB200 上原生指数指令用 0 次。accurate 档 13 个阶段，约 25% 的原生指数，用一点指数运算换回精度。

免训练这条路不是从零开始的。SageAttention3 已经证明过消费级 Blackwell 上能做 4 位量化注意力，但它瞄准的是 RTX5090，P 的处理机制也不同。本文补的是数据中心这一格：GB200 上的免训练 FP4 注意力，前向内核平均快约 2 倍，具体数字下一章拆开看。SageAttention3

推理侧绕过去了，训练侧还有一道更硬的关卡。反向传播要过 softmax 梯度，前向存的 4 位量化结果能直接喂给 backward 吗？

训练端：MXFP4 全军覆没

论文的训练路径叫因果训练，核心是「前向量化、直传反向」。前向只保存 NVFP4 Q/K 的 payload 字节、块 scale，以及每行一个的 LSE，也就是 log-sum-exp 统计量。反向直接拿这些量化操作数重算 QK 转置，用 LSE 重构概率，不另建一条 BF16 的精确分数路径。重构之后走五个分块矩阵积：dP 等于 dO 乘 V 转置，dV 等于 P 转置乘 dO，dS 等于 P 逐元素乘上 dP 减 r 的差，r 是 O 与 dO 的内积，最后 dQ 等于 dS 乘 K、dK 等于 dS 转置乘 Q。概率重构和 dP 必须先于 dQ 和 dK，这是一条串行的依赖链。

梯度操作数的精度选择是这条链的生死线。前向发布的梯度里，dO 曾用 E4M3 表示，结果约 97% 的数被舍入成 0，整个梯度形同虚设。换成 E5M2 之后，零的比例才降到约 14%。孤立地看，从量化操作数重建 backward 核心能省约 29% 的延迟，但一旦加上训练安全的统计发布，省下的时间又被吃掉，整体回到 0.986 倍。

真正炸开的是 MXFP4。论文做了四臂诊断，把投影精度和概率值精度两个变量交叉，结果两个 MXFP4 P/V 臂在约 0.1B token 处就和 FP8 控制组分离，随后 loss 不降反升、梯度范数爆开。匹配 B4 配置的复现更触目惊心：前 300 步还跟着 FP8 走，然后 loss 从第 325 步的 7.01 一路涨到第 350 步的 16.25。论文的结论写得很干脆：匹配的分布式训练保留 FP8 概率和值，所有测试过的 MXFP4 概率/值训练轨迹全部发散。FP4 FlashAttention-4

换成保留 FP8 P/V 的分布式训练，100B token 的 schedule 就能跑完，最后同一步的校验 loss，BF16 是 2.3048，保留 FP8 P/V 的路线是 2.3948（差距 0.09），中位数吞吐比约 1.112 倍。练得动，还更快，只是校验 loss 略高于 BF16，值不值，只看你给概率和值留多少精度。

为什么同一个「4 位注意力能不能训」的问题，答案差这么多？回到格式本身。MXFP4 是 32 个元素共享一个 2 的幂 scale，动态范围不愁，但单个元素只有 4 位精度，概率这种需要精细相对精度的数值吃不住。再加上前面那个伏笔，MX 规范从没做过 MXFP4 的全 4 位训练，本文等于把 MXFP4 推到注意力数值上，才踩到这颗雷。

同一问题，别的路线给出了正面答案。Attn-QAT 用量化感知训练稳定 4 位注意力，提炼出两条原则：backward 重算的 P 必须和 forward 一样做 fake-quantize，softmax 梯度的恒等式需要一条高精度辅助输出 O' 来维持。Attn-QAT Full-Stack FP4 走混合精度：QK 转置量化到 NVFP4，softmax 敏感的 PV、P 转置 dO、dO V 转置留在 BF16，3B 模型 64B token 预训练 loss 差距只有 0.838%，而且明确自限，不迁移到 MXFP4。Full-Stack FP4

三份工作放在一起，答案很清晰：格式与缩放粒度的选择，直接决定训练是稳是炸，不是位宽越低就一定越快。训练有红线，那论文宣称的 2.13 倍、1.14 倍，到底在什么条件下成立？

2.13× 是真是假：把数字放回测量边界

先给结论：2.13× 是真的，但它是最优形状下的前向内核峰值，不是典型值，更不是端到端收益。精确口径是，S8192/H64 这个形状下，FP4 前向内核耗时 0.758 ms，对照 BF16 的 1.611 ms，正好 2.125 倍。同一套九形状网格的几何均值是 2.023 倍。headline 形状 B1/S4096/H24/D128 只到 1.783 倍，本地 NV/FP8 控制组更只有 1.118 倍。这些数字全部是 GB200 上、单一形状、kernel 级的吞吐。质量数字则是另一组口径。ViT 在 4096 长序列上 fast 档 top-1 88.5%，和 BF16 持平。Wan2.1 视频模型 fast 档 1.75 倍（1.3B）、2.09 倍（14B）。MAE 重建 PSNR 差约 0.02 dB。注意这些都是固定输入的推理评估，不是微调或预训练。FP4 FlashAttention-4

速度和质量怎么权衡，第一张图看得很清楚。headline 形状下，越靠左上越好，最快的 full-FP4 点都在参照点左上角，代价是接受更多算子误差。

图注：headline 形状（B1/S4096/H24/D128）的速度-误差平面，横轴为算子误差、纵轴为速度，越靠左上越好。最快的 full-FP4 点比 NV/FP8、FP8 PV 参照点接受了更多算子误差。FP4 FlashAttention-4

第二张图则是形状依赖的证据。2.125 倍只是 S8192/H64 这一个点，换个形状就掉下来，同样网格里 S4096/H64 只有 1.828 倍。

图注：跨多种注意力形状（startup、transition、saturated，含 H64）的配对散点，每点标出相对 BF16 的加速比与相对 L2 误差，约 2.125 倍的最高点出现在 S8192/H64 这一个形状上。FP4 FlashAttention-4

训练端的数字更现实。单卡 8B 模型的完整参数更新，batch 4 时约 1.14 倍，FP8 口径 1.137 倍、MXFP4 口径 1.141 倍，MFU 从 41.12% 升到 46.74%。可这是峰值，batch 1 只有约 1.09 倍，而且这张表里明明白白写着：只测性能，不构成训练质量结论。

把三条路线摆在一起看，锚论文的位置也很清晰。它站在「免训练 + 数据中心 Blackwell」这一格，和另外两条路线各管一段。

读这类论文的第一课是问测量边界：数字在哪一层测的、哪个形状、哪个操作点、哪块硬件。2.13 倍是特定形状的内核最高值，RTX5090 上的 1038 TOPS 和 GB200 上的峰值 2998 TFLOP/s 各说各话，跨平台不能加总。论文还顺手测了 B300：指数吞吐翻倍到每 SM 每时钟 32 次，短形状上比 GB200 好 5.6% 到 7.7%——但这反过来说明更快的算术并不会让整个 kernel 等比例变快，论文据此把真正上限更多归因于 TMEM 列所有权与操作数就绪，而非指数运算本身。那这套设计离生产还有多远，论文自己画了哪些边界？

离生产还有多远：局限与那句被点破的判断

局限一条条摆出来。硬件依赖是第一条，整个设计是 Blackwell 专属的，NVFP4 格式、TMEM、异步指令缺一不可。TMEM 每 SM 只有 256 KB，两个 score bank 加两个输出 bank 占满 512 列，可分配的重叠窗口被锁死。诊断数字很说明问题。同 98,304 条 tensor 指令下，真实概率路径的 tensor-pipe 活动率只有 18.8%，固定概率时有 26.6%。2669 个未发布样本里，1715 个在等 long-scoreboard 依赖。矩阵乘在闲着等数据，这就是「张量核快不等于注意力快」的硬件侧证据。FP4 FlashAttention-4

其次，为什么 FP8 P/V 可行而 MXFP4 不行，论文只给了诊断，没有给出完备解释。质量评估都是固定输入下的推理评估，因果测量止于算子或整卡更新级别，100B 分布式训练是单轨迹，没有 run-to-run 的变异性估计。状态上更要拎清：论文 2026 年 9 月 3 日提交，仅 v1，发布三天，未经同行评审，复现仓库 github.com/MrHuff/fp4-fa4 只确认了可访问。Direct-P 是一条面向数据中心 Blackwell 的免训练研究路线，能否接入现有部署通道仍待完整复现与工程验证。

生态坐标倒是不含糊。截至 2025 年 6 月的官方博客口径，NVFP4 已经是 NVIDIA 主推的 4 位推理格式，TensorRT-LLM 和 vLLM 都在接，Hugging Face 上已经有预量化 checkpoint 可下载。NVFP4 博客 锚论文的 Direct-P 正是往这条通道里填的一块拼图，只是这块拼图能否合上，要等复现与工程验证来回答。

论文的结语更像一声提醒，原文大意是：更快的矩阵算术本身是不够的，低精度注意力需要更大的可分配重叠窗口，就像它需要更快的张量核一样。硬件把矩阵乘翻倍，只是把瓶颈搬了个家。想真把 FP4 的速度装进兜里，还得自己啃下 softmax 转换和格式选择这两块骨头。作者自己的另一项工作 UE5M3 块缩放保留了 E2M1 的 payload、给了更多 scale range，已经实现了稳定的 FP4 语言模型预训练，也许通往完全 FP4 训练的路，就在那张地图的下一页。这张「FP4 注意力」的地图值得收藏，也值得等复现检验。

想深入可以读原文：Hardware-Aware FP4 FlashAttention-4 — arXiv:2609.04105

参考文献

1. Hardware-Aware FP4 FlashAttention-4（锚论文，Robert Hu，arXiv:2609.04105，v1 2026-09-03）— https://arxiv.org/abs/2609.04105
2. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning（Tri Dao，arXiv:2307.08691）— https://arxiv.org/abs/2307.08691
3. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision（arXiv:2407.08608）— https://arxiv.org/abs/2407.08608
4. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling（Dao-AILab，arXiv:2603.05451）— https://arxiv.org/abs/2603.05451
5. SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training（arXiv:2505.11594，v3）— https://arxiv.org/abs/2505.11594
6. Attn-QAT: 4-Bit Attention With Quantization-Aware Training（Hao AI Lab，arXiv:2603.00040，v3）— https://arxiv.org/abs/2603.00040
7. Introducing NVFP4 for Efficient and Accurate Low-Precision Inference（NVIDIA 官方博客，2025-06-24）— https://developer.nvidia.com/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference/
8. Microscaling Data Formats for Deep Learning（OCP MX 规范，arXiv:2310.10537）— https://arxiv.org/abs/2310.10537
9. Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention（arXiv:2607.04422，v2）— https://arxiv.org/abs/2607.04422