PyTorch FlashAttention-4 adds MXFP8 block-scaled forward and backward for Blackwell

PyTorch published details on extending FlashAttention-4 with MXFP8 support for both forward and backward passes on Blackwell tensor cores. The implementation reaches 2.85 PF/s forward and 2 PF/s backward on LLM shapes, and on internal shapes it delivers 2.54 PF/s forward and 1.58 PF/s backward, up to 1.6x and 1.52x gains over BF16. The code is open sourced in the ads_model_kernel_library repository under the lp_fa4 directory.

TMEM allocation and scale factor placement

Blackwell's block-scaled MMA instructions operate natively on MXFP8, MXFP6, MXFP4, and NVFP4, delivering 2 to 4 times the throughput of BF16 MMA. The challenge is fitting scale factors into the fixed 512-column TMEM, which is already fully utilized for MMA operands and accumulators in existing Blackwell FlashAttention kernels. The forward pass uses a ping-pong scheme with two Q tiles, and the team overlaps scale factors with accumulator regions that are not yet active. For the prologue S GEMMs, scale factors live in the O accumulator region since O has not started. Scale factors for O overlap with S, using an existing barrier that ensures S is consumed before O scale factors are copied. Scale factors for S overlap with the other S accumulator, requiring one additional barrier between the MMA and softmax warps to avoid write-write conflicts. This barrier typically adds no cost because an O GEMM overlaps the TMEM-to-register reads that precede it.

The backward pass has five GEMMs, and dK and dV are persistent accumulators that occupy their TMEM regions for the entire kernel lifetime. Scale factor placement is restricted to the S and dP regions, with dP used for the first two GEMMs and S for the next three. Only one additional barrier is needed before the S GEMM.

Quantization along both axes and jagged tensors

Block-scaled MMA requires quantization along the GEMM K-dimension. In the forward pass, Q and K quantize along the embedding dimension, while V quantizes along the sequence dimension. The backward pass has GEMMs that need quantization along both axes, such as dP and dV. The solution uses square 32 by 32 blocks for Q, K, and dO, making the E4M3 payload transpose-invariant so both GEMMs reuse one quantized representation. The smaller E8M0 scales are still laid out and loaded separately for each GEMM's tensor core scale layout. The same square scheme handles dS, producing one representation that both dK and dQ can consume.

Jagged sequence lengths in Ads models do not naturally align to the 512-byte swizzled scale-factor atom that block-scaled MMA expects. Padding full data tensors would add memory traffic and storage overhead. Instead, the implementation uses split addressing: only the much smaller scale-factor tensors are padded to 128-aligned positions, while FP8 data remains compact. Lightweight scatter and layout-transformation kernels move scale factors from compact jagged layouts into padded, swizzled layouts, allowing each downstream GEMM or attention kernel to consume scales directly in its preferred layout.

Softmax conversion and unroll-KV

The softmax warp computes in FP32 and converts P to MXFP8 online. The conversion uses a three-step process: finding the amax for a 32-element block, computing the scale factor, and applying it. The team uses a three-instruction max and fmul2 instructions for the per-element steps, and an optimized PTX sequence for the scale factor that avoids log2 and division by extracting FP32 exponent and mantissa bits. Since softmax already computes row maxes before exponentiation, and exp is monotonically increasing, the max computed for softmax is reused, preventing additional max operations while adding one exp operation per 32 elements. Recent experiments indicate that constant scaling for P can further improve performance, remaining numerically robust because softmax constrains outputs to the 0 to 1 interval.

Lower precision increases MMA throughput, shifting the bottleneck to the softmax warp's SFU-bound computation. Tile boundaries are particularly problematic without unroll-KV, creating a bubble at every tile boundary. Unroll-KV interleaves the last PV GEMM of the current tile with the first QK GEMM of the next tile, letting softmax for Q0 trigger one GEMM earlier. For BF16 this caused a regression because the correction warp was delayed, but moving the barrier wait to after the row-sum computation decoupled the correction pipeline from the critical path.