Hardware-Aware FP4 FlashAttention-4

📅 2026-09-03
📈 Citations: 0
Influential: 0
📄 PDF
🤖 AI Summary
本文解决了FP4张量核心未自动加速注意力机制的问题,通过Direct-P方法直接映射分数到概率,并采用8位浮点数梯度操作加速训练。
📝 Abstract
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.
Problem

Research questions and friction points this paper is trying to address.

FP4
FlashAttention-4
Hardware-Aware
softmax conversion
on-chip dependencies
Innovation

Methods, ideas, or system contributions that make the work stand out.

Direct-P
FP4
softmax conversion
on-chip dependencies
causal path
🔎 Similar Papers
No similar papers found.