硬件感知的 FP4 FlashAttention-4
来自 arXiv · 期号 2026-09-04
摘要
Blackwell 的 FP4 张量核心不会自动让注意力计算变快，因为一旦矩阵乘积变小，softmax 转换和片上依赖就成了瓶颈。 本文提出 Direct-P 方法用于非因果推理：直接把注意力分数映射成 FP4 概率，在 NVIDIA GB200 上达到 BF16 前向吞吐量的 2.13 倍。因果训练路径则从保存的量化查询和键中重建概率，并使用 FP8 梯度操作数，将单 GPU 上 8B 参数的一次完整更新加速至 1.14 倍。 分布式训练中，匹配的路径保留了 FP8 概率和值；所有测试的 MXFP4 概率/值训练轨迹均发散。