面向硬件的FP4 FlashAttention-4

中文摘要 AI 辅助

Blackwell架构的4位浮点（FP4）张量核心并不会自动提升注意力机制的速度，因为当矩阵乘积缩小时，softmax转换和片上依赖会成为主要开销。针对非因果推理，我们提出了Direct-P方法；针对因果推理路径，我们将前向量化直接传递到反向传播中。Direct-P将注意力分数直接映射到FP4概率，在NVIDIA GB200上实现了bfloat16（BF16）前向吞吐量的最高2.13倍。因果路径从保存的量化查询和键中重构概率，并使用8位浮点（FP8）梯度操作数，使单GPU 80亿参数的完整更新速度最高提升1.14倍。匹配的分布式训练保留FP8概率和值，所有测试的MXFP4概率/值训练轨迹均发生发散。

英文摘要

Blackwell's 4-bit floating-point (FP4) tensor cores do not automatically make attention faster because softmax conversion and on-chip dependencies dominate once its matrix products shrink. We address this with \emph{Direct-P} for noncausal inference and a causal path that passes the forward quantization directly into backward. Direct-P maps scores directly to FP4 probabilities and reaches up to 2.13$\times$ the bfloat16 (BF16) forward throughput on an NVIDIA GB200. The causal path reconstructs probabilities from saved quantized queries and keys and uses 8-bit floating-point (FP8) gradient operands, accelerating a complete single-GPU 8-billion-parameter update by up to 1.14$\times$. Matched distributed training retains FP8 probabilities and values; every tested MXFP4 probability/value training trajectory diverges.

发表机构

• HAO AI Lab（HAO人工智能实验室）

机构由 AI 辅助整理，请以论文原文为准。