Hardware-Aware FP4 FlashAttention-4
Robert Hu shows that Blackwell FP4 tensor cores do not speed up attention on their own, because softmax conversion and on-chip dependencies dominate once the matrix products shrink, and gives separate forward and causal paths that recover the speedup.
Ask this paper
The bottleneck is not the matmul: shrinking the matrix products with FP4 exposes softmax conversion and on-chip dependencies as the cost, which is why naive FP4 attention does not get faster.
Direct-P for noncausal inference maps scores directly to FP4 probabilities and reaches up to 2.13x bfloat16 forward throughput on an NVIDIA GB200.
The causal path passes forward quantization into backward, reconstructing probabilities from saved quantized queries and keys and using FP8 gradient operands, which accelerates a full single-GPU 8B update by up to 1.14x.
A hard limit reported honestly: matched distributed training retains FP8 probabilities and values, because every tested MXFP4 probability and value training trajectory diverges.
The practical reading: FP4 is usable in the forward pass and for gradient operands, and not yet usable for probabilities and values in distributed training.
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.