PyTorch发布FlashAttention-4 MXFP8实现

TL;DR

We extend FlashAttention-4 [1] with MXFP8 forward and backward, reaching 2.85 PF/s forward and 2 PF/s backward on LLM shapes. On our internal shapes, FA4 MX8 reaches 2.54 PF/s forward and 1.58 PF/s backward, delivering up to 1.6× and 1.52× gains over BF16. We fuse quantization into the surrounding producers and develop an end-to-end zero-gather jagged module in which most activations and compute remain in FP8. This module is used internally at Meta for GEM training [2]. To our knowledge, this is one of the first SoTA implementations of MXFP8 FA4 forward and backward being used in production training workloads. We have open sourced the code in https://github.com/facebookresearch/ads_model_kernel_library/tree/main/lp_fa4

1. Introduction

Blackwell’s tensor cores introduce block-scaled MMA instructions (tcgen05.mma.block_scale) that operate natively on microscaling formats — MXFP8, MXFP6, MXFP4 & NVFP4 — delivering 2-4x the throughput of BF16 MMA [4,5]. However, exploiting this in real training workloads requires more than just swapping the attention kernel’s data type. Scale factors must be managed in the already-saturated TMEM, quantization along the GEMM K-dimension must be handled for every operand (including online-computed intermediates like P and dS), and the overhead of converting between precisions must be overlapped so MMA can continue to execute at full speed.

In this work, we extend the FA4 attention kernel with end-to-end MXFP8 support for both forward and backward passes, and integrate it into a cross-attention module for Ads training with fused producer and output epilogues. The key contributions are: (1) TMEM allocation strategies that fit scale factors into fully-utilized 512-column TMEM with minimal new barriers, (2) online transpose invariant square block scale quantization of dS using Blackwell’s redux.sync.max.abs.f32 warp-wide reduction, (3) fused RMSNorm+Quantize and GEMM+Quantize kernels that eliminate quantization overhead by producing FP8 output with dual scale factor layouts in a single pass, and (4) a zero-gather jagged module where FP8 data stays at unpadded positions and only the much smaller scale factors are scattered, padded, and permuted to tensor-core friendly 128-aligned addresses for TMA.

2. Implementation Details

2.1 Attention Forward

Attention forward consists of the following primary operations:

S = Q @ K.T

P = Softmax(S)

O = P @ V

To enable blockscaled MMA, we follow existing CuTe DSL examples from Quack GEMM kernels and CUTLASS C++ examples [5,6]. We use TMA loads to fetch scale factors from GMEM to SMEM, and copy SFs from SMEM to TMEM before triggering the UMMA. The primary challenge here is TMEM contention, which is explained in later.

Currently, in the softmax warp, softmax computation happens in FP32, and then the results are converted to BF16 before the PV multiplication. We convert P to MXFP8 while also computing the scales. We deep dive into the PTX optimizations done to achieve this efficiently later.

One subtle thing to note is that for P.V blockscaled MMA to work, the scales need to be computed along the MMA K-dim. For Q and K, this is the embedding dimension (D) of attention, but for V, the scales and quantization need to be computed along the sequence dimension (N).

2.1.1 TMEM Allocation and barrier synchronization

Blackwell architecture has a fixed TMEM size of 512 column, which is completely utilized for MMA operands and accumulators in existing Blackwell FA kernels. This makes it challenging to add block scaled MMA, since scales also need to be in TMEM.

FA4 forward uses a ping-pong computation between two Q tiles. We load two Q tiles, Q0 & Q1 of size [128, 128], and loop over K/V tiles (N dimension). The order of GEMMs is:

GEMM Prologue S0 = Q0 @ K0 S1 = Q1 @ K0 Mainloop (for n in 0 .. N-1) O0 = P0_n * V_n S0 = Q0 * K_{n+1} O1 = P1_n * V_n S1 = Q1 * K_{n+1} Epilogue O0 = P0_N * V_N O1 = P1_N * V_N

This is the TMEM Alloc:As you can see, TMEM is fully utilized with no spare room for SFs. Note that we cannot overlap input SFs with the accumulator TMEM. To solve this problem, we overlap the SF’s in the following way:

• For the prologue S(i) GEMMs, we can use O(i) for S(i) SFs, as O(i) hasn’t started yet
• SFs for O(i) can live overlapped with S(i) – this is the same strategy used by regular FA which overlaps P(i) with S(i). Thus, there already exists a barrier which ensures that S(i) TMEM is consumed, before we copy O(i) SFs. We just need to choose a region in S(i) distinct from P(i). Note that for FP8, P(i) is 32 columns, vs S(i) being F32 is 128 columns.
• SFs for S(i) live overlap with S(1-i). This requires an additional barrier between the MMA and Softmax warps, since MMA is executed asynchronously, it is possible that there is a write-write conflict between S(1-i) accumulator and S(i) SFs. Hence, we add a barrier in MMA warp which waits before copying S(i) SFs, and arrives when S(1-i)’s accumulator is read by the softmax warp. Note that this barrier doesn’t usually add any extra cost, since there is an O(i) GEMM to overlap the TMEM->reg reads that need to precede the barrier. Typically, the GEMM takes longer than TMEM->reg read.

Concretely, the GEMM execution order with SF placement is shown below:

2.1.2 Improved unroll-KV

As we move to lower precision, MMA throughput increases (2x for MXFP8, 4x for MXFP4 vs BF16), the softmax warp’s SFU-bound computation stays the same. This shifts the bottleneck: softmax bubbles that were hidden behind slow BF16 MMA now become exposed.

In the persistent kernel, the tile boundary is particularly problematic. Without unroll-KV, the last two GEMMs of a tile are both PV, followed by both QK of the next tile. Softmax for the next tile’s Q0 can’t start until QK0 completes — but QK0 can’t start until all of tile(n)’s PV GEMMs finish. This creates a bubble at every tile boundary. Unroll-KV (inspired by GDPA) [3] interleaves the last PV of the current tile with the first QK of the next tile:

Now softmax for Q0 triggers one GEMM earlier, hiding the tile boundary latency behind the PV pipeline.

However, enabling this for BF16 caused a regression: the correction warp (which processes both stages sequentially) was delayed because PV1 was pushed later, which cascaded through the softmax_corr_empty barrier into the next softmax. We fixed this by moving the barrier wait to after the row-sum computation (which has no dependency on the barrier), decoupling the correction pipeline from the critical path.

2.1.3 Optimized Online MXFP8 conversion

Since we want to do both QK and PV GEMMs using block-scaled attention, this necessitates an online conversion of the post-softmax output (P) from FP32 to MXFP8, instead of BF16. Applying blockscaling is a 3-step process. For a block x consisting of 32 elements:

Here, a is the amax for the block, and sigma is the scaling factor. In order to do this optimally in Blackwell, we make use of 3-instruction max, and fmul2 instructions (exposed in CuteDSL nvvm) for step (1) and (3) – which happen for every element. For computing the scale factor (Step 2) – we use an optimized PTX sequence which avoids log2 and division by extracting FP32 exponent and mantissa bits.

Additionally, we notice that computation of softmax already involves calculating row maxes for the 128 elements before the exponentiation. Since exp is a monotonically increasing function, we can re-use the max computed for softmax, thus preventing additional max operations while adding 1 exp operation per 32 elements. We find this approach to yield better performance.