Meta Brings Low-Precision FlashAttention-4 to Blackwell

Meta has open-sourced an extension of FlashAttention-4 that introduces end-to-end MXFP8 support for both forward and backward passes. Optimized for Nvidia Blackwell's block-scaled tensor cores, this implementation achieves 2.85 petaflops in the forward pass and 2 petaflops in the backward pass on large language model shapes. On Meta's internal recommendation system shapes, the low-precision kernel reaches 2.54 petaflops forward and 1.58 petaflops backward, representing speedups of up to 1.6x and 1.52x over traditional BF16 formats.

To achieve these speeds, the developers resolved several hardware bottlenecks on Blackwell. They designed a memory allocation strategy to fit scale factors into the architecture's tightly constrained 512-column tensor memory. Additionally, they introduced a transpose-invariant square quantization scheme using 32-by-32 blocks, allowing both forward and backward passes to reuse a single quantized representation. For variable-length data common in recommendation models, a zero-gather jagged module keeps the primary data compact while only padding and aligning the much smaller scale factors.

The team also eliminated quantization overhead by fusing MXFP8 conversion directly into preceding operations. On Blackwell B200 shapes, fusing quantization into the matrix multiplication epilogue boosted throughput from 0.20 petaflops to between 0.91 and 0.98 petaflops, a 4.4x to 4.7x improvement. Similarly, fusing RMSNorm with MXFP8 quantization increased effective memory bandwidth from 0.75–0.87 terabytes per second to 2.6–4.0 terabytes per second.

For machine learning practitioners, these developments make low-precision training highly practical. The code has been integrated into Meta's internal GEM training workloads and released publicly in the Ads Model Kernel Library. By shifting the bottleneck away from memory bandwidth and quantization overhead, the release allows developers to leverage Blackwell's native microscaling formats for production-scale training without sacrificing numerical stability.