研究提出FlashAttention-4新路径：FP4下注意力前向吞吐最高达BF16的2.13倍
论文作者称，直接以FP4映射注意力分数可在GB200上实现最高2.13倍于BF16的前向吞吐，并给出单GPU 80亿参数更新加速1.14倍的因果训练方案。
AI解读：Blackwell的FP4张量核心虽然能加速矩阵乘法，但并不会自动让注意力计算变快，因为softmax转换和芯片上的依赖会拖慢整体速度。这篇论文提出了两条路径来绕过瓶颈：非因果推理用Direct-P把注意力分数直接映射成FP4概率，省去中间转换，在NVIDIA GB200上前向吞吐最高达到BF16的2.13倍；因果训练则把前向量化结果直接传给反向传播，并用FP8梯度操作数来重建概率，使单个GPU上80亿参数的一次完整更新最高提速1.14倍。简单说，这项研究针对的是FP4精度下注意力计算效率低的问题——如果采用直接映射，就能让推理更快；如果做训练，需要同时处理正向和反向的精度匹配。值得注意的限制是，论文称在匹配的分布式训练中保留FP8的概率和值可行，但所有测试的MXFP4概率/值训练轨迹都发散，这意味着FP4直接用于分布式训练还不成熟。不过这些都是作者在arXiv预印本中的宣称，尚未经过同行评审，效果数据也来自测试条件，实际应用时需结合具体任务验证。
arXiv预印本论文《Hardware-Aware FP4 FlashAttention-4》提出，针对Blackwell FP4张量核心的注意力计算优化方案，可使前向吞吐最高达BF16的2.13倍。论文作者Robert Hu称，FP4张量核心并不会自动加速注意力，因为softmax转换和片上依赖在矩阵乘积缩小时成为主导。
Direct-P与非因果推理加速
论文提出“Direct-P”方法用于非因果推理，将注意力分数直接映射为FP4概率，从而跳过转换步骤。在NVIDIA GB200上，该方法实现的前向吞吐最高达到BF16（bfloat16）的2.13倍。
- 该方法宣称在FP4精度下有效，但论文未提供具体数据集或测试条件的细节。
因果训练路径与分布式限制
针对因果训练，论文提出将前向量化直接传入反向传播，从保存的量化查询和键中重建概率，并使用FP8（8位浮点）梯度操作数。作者称在单GPU上完成一次80亿参数的完整更新可加速最高1.14倍。
- 论文称在匹配的分布式训练中保留FP8概率和数值可行；相反，所有测试的MXFP4概率/值训练轨迹（training trajectory）均出现发散，表明该精度组合在分布式训练中不可用。