论文

硬件感知 FP4 FlashAttention-4

Hardware-Aware FP4 FlashAttention-4

摘要

Blackwell 的 4 位浮点 (FP4) 张量核心不会自动加快注意力速度,因为一旦矩阵乘积缩小,softmax 转换和片上依赖性就会占主导地位。我们用 \emph{Direct-P} 来解决这个问题,用于非因果推理,以及将前向量化直接传递到后向量化的因果路径。 Direct-P 将分数直接映射到 FP4 概率,并在 NVIDIA GB200 上达到 bfloat16 (BF16) 前向吞吐量的 2.13$\times$。因果路径根据保存的量化查询和键重建概率,并使用 8 位浮点 (FP8) 梯度操作数,将完整的单 GPU 80 亿参数更新加速高达 1.14$\times$。匹配分布式训练保留FP8概率和值;每个测试的 MXFP4 概率/值训练轨迹都有所不同。