Hardware-Aware FP4 FlashAttention-4
cs.LG
Submitted: 2026-09-03
Updated: 2026-09-03
Code: https://github.com/hao-ai-lab/flash-attention-fp4
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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.
Terminology
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 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.
Sources
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
- Masked Autoencoders Are Scalable Vision Learners
- UE5M3 FP4 Block Scaling for Stable Language Model Pretraining
- Microsoft COCO: Common Objects in Context
- FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
- Wan: Open and Advanced Large-Scale Video Generative Models
- FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling
- SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training
- Attn-QAT: 4-Bit Attention With Quantization-Aware Training
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks