Meta's FlashAttention-4 Achieves 2.85 PF/s with MXFP8 on Blackwell GPUs
Breakthrough in low-precision attention for LLM training with end-to-end FP8 support and optimized memory management
1Breakthrough in low-precision training
Meta's FlashAttention-4 achieves 2.85 PF/s forward and 2 PF/s backward performance on LLM shapes using MXFP8 precision. The implementation shows 1.6× gains over BF16 in forward passes and 1.52× in backward passes. This marks significant progress in low-precision training efficiency for large language models, particularly on NVIDIA's Blackwell architecture with its native support for microscaling formats.
2End-to-end FP8 pipeline
The system maintains most activations and computations in FP8 throughout the entire training process. Key innovations include fused quantization in surrounding producers and a zero-gather jagged module that keeps FP8 data at unpadded positions. This approach minimizes precision conversions and memory overhead while maintaining numerical stability during training.
3TMEM allocation strategies
The team developed novel techniques to fit scale factors into Blackwell's fully-utilized 512-column TMEM. They implemented ping-pong computation between Q tiles and overlapped scale factor storage with existing accumulator regions. Additional barriers were carefully placed to prevent write-write conflicts while maintaining high throughput, with GEMM operations typically taking longer than TMEM-register reads.
4Optimized online conversion
The implementation features an efficient MXFP8 conversion process that reuses softmax computations to minimize overhead. By leveraging Blackwell's specialized instructions like redux.sync.max.abs.f32 and fmul2, the team achieved high-performance online quantization. The system also benefits from constant scaling for softmax outputs, which remains numerically robust due to the natural [0,1] constraint of softmax.
5Jagged sequence handling
The solution addresses variable sequence lengths through split addressing, where only small scale-factor tensors are padded to 128-aligned positions while FP8 data remains compact. Lightweight scatter and layout-transformation kernels rearrange scale factors without changing values, enabling efficient handling of jagged inputs common in advertising models.
6Backward pass innovations
The backward pass implementation handles five GEMM operations with careful TMEM management. Square [32,32] block quantization makes representations transpose-invariant, allowing reuse across operations. FP16 dQ reduction was identified as critical for throughput, with static scaling providing optimal performance while maintaining acceptable numerical accuracy in production workloads.
7Quantization overhead reduction
The team fused MXFP8 quantization into surrounding operations to eliminate overhead, particularly for smaller sequence lengths. This includes developing fused RMSNorm+Quantize and GEMM+Quantize kernels that produce FP8 output with dual scale factor layouts in a single pass, significantly reducing conversion bottlenecks.
8Production deployment
The system is currently used internally at Meta for GEM training, representing one of the first state-of-the-art implementations of MXFP8 FlashAttention in production training workloads. The open-sourced code provides a reference implementation for the research community to build upon in the ads_model_kernel_library repository.
9Blackwell architecture advantages
The implementation leverages Blackwell's tensor cores with block-scaled MMA instructions that natively support MXFP formats. These deliver 2-4× the throughput of BF16 MMA operations. The work demonstrates how to fully exploit these hardware capabilities in real training workloads beyond simple data type swaps.
10Numerical stability techniques
Despite the aggressive use of low precision, the system maintains training stability through several innovations. These include square block scale quantization, careful scale factor management, and numerical analysis of precision requirements at each operation. The team validated the approach's robustness across various model architectures and training scenarios.
11Performance trade-offs
The transition to lower precision exposed new bottlenecks, particularly in softmax computation. As MMA throughput increased with MXFP8, previously hidden softmax bubbles became visible. The team addressed this through techniques like unroll-KV that interleave operations across tile boundaries, hiding latency behind the PV pipeline.
12Future directions
The work opens several research directions, including extension to even lower precisions like MXFP6 and MXFP4, broader application across model architectures, and further optimization of the quantization pipeline. The successful production deployment demonstrates the viability of aggressive low-precision approaches for modern LLM training at scale.