Meta engineers extended FlashAttention-4 with end-to-end MXFP8 block-scaled support for both forward and backward passes on NVIDIA Blackwell GPUs, addressing TMEM allocation, online quantization of softmax outputs and gradients, and jagged/variable-length sequence handling for production ads recommendation training. The implementation reaches up to 2.85 PF/s forward and 2 PF/s backward on LLM-shaped workloads, delivering 1.6x and 1.52x throughput gains over BF16 respectively. Fusing quantization into RMSNorm and GEMM epilogues yields additional 4.4-4.7x speedups by eliminating separate quantization passes. The kernel is being open-sourced in Meta’s Ads Model Kernel Library.
