Hardware-Aware FP4 FlashAttention-4

TL;DR - Hardware-Aware FP4 FlashAttention-4 restructures attention around NVIDIA Blackwell tensor cores to avoid softmax conversion and on-chip dependency bottlenecks. It improves inference and training throughput, but experiments show that FP4 probability/value representations can destabilize distributed training.

• Direct-P maps attention scores directly to FP4 probabilities, delivering up to 2.13× BF16 forward throughput on an NVIDIA GB200 for noncausal inference.
• The causal training path reuses forward quantization in backward, reconstructing probabilities from quantized queries and keys and using FP8 gradient operands.
• This approach speeds a complete single-GPU 8B-parameter update by up to 1.14×.
• Distributed training remained stable with FP8 probabilities and values, while every tested MXFP4 probability/value trajectory diverged.