AMD把softmax搬进矩阵核：MXFP4注意力在MI355X上达4.6 PFLOP/s，视频生成端到端最高提速
核心概要
AMD为CDNA 4架构推出免训练的Fused-Softmax Attention，把softmax中占主导的向量运算折叠进注意力的矩阵乘法，让低精度矩阵吞吐真正被用上。在MI355X上，MXFP8、MXFP6、MXFP4内核分别达到约3.0、4.0、4.6 PFLOP/s，而新增误差仍低于量化GEMM操作数本身带来的误差下限。在Wan2.2-T2V-A14B上，MXFP6/4混合调度在保持画面质量的同时带来端到端最高加速，内核已随AMD开源AITER库发布。
深度剖析
注意力里的softmax不再主要占用向量单元，其大部分工作被移入两次矩阵乘法与一次格式转换，使低精度矩阵吞吐可以直接转化为注意力吞吐。 此前FlashAttention系列把矩阵乘法与softmax的向量工作流水重叠，但GPU代际间矩阵吞吐增长远快于向量吞吐，量化操作数只会进一步拉大这一差距，softmax因此成为瓶颈。 在AMD Instinct MI355X上，MXFP8、MXFP6、MXFP4内核分别达到约3.0、4.0、4.6 PFLOP/s，相对BF16注意力取得加速，且完整流水线在量化内核上最多只给操作数量化已有的误差增加几个百分点。
做法是把行最大值先冻结并提前减进打分矩阵乘法，把指数近似成一次仿射变换后并入同一矩阵乘法，行和也改由矩阵乘法完成。 冻结的最大值只在后续分块溢出时才触发回滚，回滚时按真实最大值与推测最大值之差重新平移并缩放累加结果，因此稳态循环里不再需要逐块重缩放。 指数用Schraudolph仿射近似，经两级转换得到量化概率，块最大值同时用于推导E8M0缩放因子和溢出门控；在Wan2.2-T2V-A14B的真实注意力调用上，回滚很少发生，且关掉门控后各内核输出为NaN，说明冻结最大值单独使用会溢出。
在视频扩散模型上，纯MXFP4会带来可见的画面漂移，而MXFP6/4混合调度把精度损失压回接近BF16的水平。 此前的低比特注意力工作多保持softmax不变，或只替换指数运算，行最大值与行和仍留在向量单元上。 在Wan2.2-T2V-A14B的248条VBench提示上，MXFP8、MXFP6与混合调度的VBench总分与BF16基线相当，单独使用MXFP4则明显落后；在HunyuanVideo与MiniMax-H3上，混合调度同样恢复了参考画面的外观。
启示与展望
评估集中在AMD Instinct MI355X与CDNA 4架构，精度与吞吐对比使用固定形状的非因果前向注意力，视频质量评估覆盖Wan2.2-T2V-A14B、HunyuanVideo与MiniMax-H3三个扩散模型。
冻结最大值的滞后上限是按扩散负载的长序列长度选定的启发式，累加范围还取决于序列长度与数值幅度，无法仅由BF16指数范围确定；MXFP4使用的快速最大值归约只在有限范围内保证门控正确，超出该范围时真实溢出量会丢失。