FlashAttention-4 Brings End-to-End MXFP8 Attention to Blackwell

Introduction

Using FP8 efficiently is not simply a matter of changing the input and output types of an attention kernel. Block-scaled matrix multiplication also requires scale factors to be generated, moved, synchronized, and aligned with the tensor-core K dimension. PyTorch has now described an extension of FlashAttention-4 that brings MXFP8 to both forward and backward attention on NVIDIA Blackwell GPUs, while integrating quantization with surrounding operators and variable-length data handling.

What the implementation changes

• MXFP8 covers both passes. On LLM-oriented shapes, the reported peak is 2.85 PFLOPS for forward and 2 PFLOPS for backward. On the team’s internal shapes, forward reaches 2.54 PFLOPS and backward 1.58 PFLOPS. The stated improvements over BF16 are up to 1.6x and 1.52x respectively.
• Scale factors are fitted into a crowded TMEM. Existing FlashAttention-4 kernels already make extensive use of Blackwell’s 512-column TMEM. The new design overlaps scale-factor storage with regions whose accumulators have already been consumed, while adding barriers where asynchronous MMA could otherwise create write conflicts.
• Quantization is fused with producers. Fused RMSNorm-plus-quantization and GEMM-plus-quantization kernels can emit FP8 data together with scale factors in layouts required by later operations. This avoids treating conversion as a separate, exposed stage in the pipeline.
• The softmax output is quantized online. The probability matrix P is computed during attention and must be converted for the subsequent P·V operation. The implementation reuses row maxima already produced by softmax and applies instruction-level optimizations to reduce the cost of MXFP8 conversion.
• Jagged inputs avoid large gathers. In the variable-length module, FP8 activations remain at their unpadded positions. Only the smaller scale-factor data is scattered, padded, and rearranged to addresses that are friendly to TMA transfers.

Why it matters

Blackwell introduces block-scaled MMA instructions for formats including MXFP8, MXFP6, MXFP4, and NVFP4. Their potential throughput advantage over BF16 is substantial, but the benefit can disappear if conversion and scale management are handled outside the main pipeline. This work therefore illustrates a broader principle: low-precision attention must be designed across memory layout, synchronization, quantization, and neighboring kernels.

The optimization also changes where the bottleneck appears. Faster MMA exposes the cost of softmax and other special-function operations that were previously hidden behind slower BF16 matrix multiplication. The authors adjust KV unrolling and barrier placement to reduce bubbles at tile boundaries and keep the asynchronous pipeline moving.

The implementation is used in an internal Meta GEM training workload and has been open-sourced in the associated repository. Its benchmark figures are tied to particular shapes and kernel configurations, so they should not be read as universal speedups. The more general contribution is an end-to-end route for keeping FP8 data in the attention path, including forward, backward, fused producers, and jagged-sequence handling.

Source: PyTorch Blog