Meta's MXFP8 FlashAttention-4 kernel, tackling TMEM scale-factor constraints, hits up to 1.6x over BF16

Blackwell's block-scaled MMA runs MXFP8, MXFP6, MXFP4 and NVFP4 natively at two to four times BF16 MMA throughput ^[5]. The attention module built on it gains up to 1.6x forward and 1.52x backward ^[2]. Softmax still runs in FP32, and only then is P converted and quantized, with the scales computed in the same step ^[12]. Scale factors have to be fetched from global memory into shared memory by TMA and copied into TMEM before the UMMA fires, and the post names TMEM contention as the primary challenge ^[17].

TMEM is fixed at 512 columns, and existing Blackwell FlashAttention kernels already use all of it for MMA operands and accumulators ^[7]. Input scale factors cannot be parked on accumulator TMEM ^[7]. So the scales borrow space from buffers that are not busy yet. FA4 forward ping-pongs across two [128, 128] Q tiles while looping over K/V tiles along N ^[8]. Prologue S(i) scale factors go in O(i), which has not started; the O(i) scales go in a region of S(i) picked to avoid P(i), which works because a barrier already guarantees S(i) has been consumed; the S(i) scales overlap S(1-i), and that overlap needs a new barrier ^[9]. The room comes from the precision change: P(i) in FP8 occupies 32 columns where S(i) in F32 occupies 128 ^[10], which leaves 96 columns of the S(i) region outside P(i) ^[11].

Scales must be computed along the MMA K dimension, which is the embedding dimension D for Q and K but the sequence dimension N for V ^[13]. dS is quantized online in transpose-invariant square blocks using Blackwell's redux.sync.max.abs.f32 warp-wide reduction ^[14]. Both are easy to get wrong in your own version. The implementation follows the Quack GEMM kernels in CuTe DSL and the CUTLASS C++ block-scaled examples ^[21], so a team that has never allocated TMEM by hand starts there.

The quantization is fused into the producers, with RMSNorm+Quantize and GEMM+Quantize kernels emitting FP8 and dual scale factor layouts in a single pass ^[15]. The benchmark subject is a cross-attention module for Ads training with fused producer and output epilogues ^[18]. Jagged inputs never get gathered: FP8 data stays at unpadded positions and only the much smaller scale factors are scattered, padded and permuted to 128-aligned addresses for TMA ^[16]. Shape sensitivity is visible inside the post itself, since the forward figure falls about 11 percent going from the LLM shapes to Meta's internal shapes and the backward figure about 21 percent ^[19].