Low-Precision FlashAttention-4: End-to-End Block-Scaled Attention on Blackwell

From FlashAttention-4 to Low-Precision

The FlashAttention series has long been a cornerstone of efficiency optimization for large model training and inference. By computing attention in tiles with online softmax, it avoids materializing the full attention matrix in GPU memory, dramatically reducing memory footprint and boosting throughput. The latest low-precision FlashAttention-4 (referred to here as FA4 MX8) takes this further by migrating both the forward and backward passes of attention entirely to MXFP8 — a microscaling low-precision format — purpose-built for NVIDIA's latest Blackwell architecture.

According to the paper, this work extends FlashAttention-4 with MXFP8 forward and backward implementations, achieving 2.85 PF/s forward and 2.0 PF/s backward throughput on typical LLM shapes. On the authors' internal benchmark shapes, FA4 MX8 forward also reaches 2.54 PF/s. These numbers directly demonstrate the potential of low-precision block-scaled computation on Blackwell hardware.

Why MXFP8

MXFP8 belongs to the Microscaling floating-point format family. Its core idea is to share a single scaling factor across a small block of values, preserving the representable range of 8-bit floats while mitigating the limited dynamic range inherent to low precision. Compared to traditional FP8 with a single global or tensor-level scale, block-scaled quantization adapts more granularly to the varying numerical distributions within a tensor — which is especially critical for attention, where values can span an enormous range.

The QK^T dot product, softmax exponentiation, and weighted sum with V in attention computation can all exhibit significant numerical variation. Using low precision without block-level scaling risks introducing substantial precision loss around the softmax operation. FA4 MX8 adopts MXFP8 as its end-to-end data format precisely to strike a balance between compute efficiency and numerical stability.

In practice, MXFP8 comes in two variants: E4M3 (4-bit exponent, 3-bit mantissa) and E5M2 (5-bit exponent, 2-bit mantissa). The former has a narrower dynamic range but higher precision, making it suitable for forward activations; the latter offers a wider range and is better suited for gradient representation. The MX (Microscaling) format is defined by the OCP Microscaling Formats specification, an industry standard jointly promoted by AMD, Intel, NVIDIA, Microsoft, and others. Its block-shared scaling design typically groups 32 elements per block, storing a single FP8 exponent as the shared scale — with negligible overhead. Compared to the tensor-level FP8 scaling used in Transformer Engine, this is one to two orders of magnitude finer in granularity, enabling more accurate capture of local numerical distributions and allowing quantization without significant precision degradation.

What End-to-End Block Scaling Means

Many prior low-precision optimizations only cover forward inference, or apply low precision to a subset of operators. This work emphasizes end-to-end coverage — both forward and backward passes run entirely in MXFP8 block-scaled arithmetic. This means not only inference but also gradient computation during training can benefit from the throughput gains of low precision.

Backward passes are generally more sensitive to numerical precision, since accumulated gradient errors can compound over long training sequences. Achieving a stable 2 PF/s in the backward pass demonstrates that the block-scaling strategy is viable along the gradient path as well — a practically meaningful result for teams looking to train large models at lower cost.

Hardware Adaptation for Blackwell

NVIDIA's Blackwell architecture natively supports tensor core operations in MXFP8 and other microscaling formats, which is the hardware foundation that enables FA4 MX8 to fully realize its performance potential. Tightly coupling algorithm design with hardware characteristics is a prevailing trend in high-performance attention kernel optimization — only with software-hardware co-design can theoretical peak throughput be effectively converted into real-world gains.

The reported PF/s (PetaFLOPs per second) figures show that FA4 MX8 pushes attention kernels into the petaFLOP-scale measured performance range — a level that was previously out of reach for attention implementations relying primarily on BF16 or FP16.

NVIDIA's Blackwell architecture (exemplified by the GB200/B200) features fifth-generation Tensor Cores with native MXFP8 matrix multiplication support, offering roughly twice the theoretical peak compute of the previous Hopper generation (H100) at FP8 precision. While Hopper also supports FP8, its scaling operates at the tensor granularity — requiring a global scaling factor to be precomputed before each matrix multiply, which introduces additional synchronization overhead in streaming computations like attention. Blackwell's native hardware support for MX block-level scaling allows scaling factors to be fused directly into matrix multiply execution, eliminating this bottleneck. This is the fundamental reason why FA4 MX8's performance numbers are reproducible only on Blackwell and cannot be directly ported to Hopper.

During backpropagation, the attention mechanism must either recompute or cache intermediate forward results (such as the softmax output matrix P) to compute gradients. The FlashAttention family addresses this through a recomputation strategy — re-executing portions of the forward pass during backward to trade compute for memory, avoiding storage of the full attention matrix. Extending this to MXFP8 is challenging because gradient value ranges are often harder to predict than activations, especially for the dQ, dK, and dV computations, which involve matrix multiplications over softmax gradients with higher numerical stability requirements. Successfully maintaining low precision along this path without gradient explosion or vanishing is precisely where block-level scaling holds a key advantage over tensor-level scaling.

Implications for Large Model Engineering

The continued evolution of low-precision attention is reshaping the cost structure of large model training and inference. Higher compute density means the same hardware can handle larger batch sizes, longer contexts, or reduced training time within the same budget. In long-context scenarios where attention is the dominant compute and memory bottleneck, optimizations like FA4 MX8 yield proportionally greater marginal gains.

It's worth noting that the practical benefits of low-precision approaches are highly dependent on specific model shapes. The paper reports two sets of numbers — one for general LLM shapes and one for internal shapes — with a visible gap (2.85 PF/s vs. 2.54 PF/s forward). This serves as a reminder that real-world deployment requires empirical validation against your own workload, rather than assuming peak numbers will generalize directly.

Summary

FA4 MX8 represents the latest step in low-precision attention kernel development: it applies MXFP8 block-scaled quantization throughout both forward and backward passes, fully leverages Blackwell's hardware capabilities, and achieves petaFLOP-scale forward and backward throughput on LLM shapes. For engineering teams pursuing maximum training and inference efficiency, this is a direction worth watching closely — though shape-specific tuning and precision validation remain essential before production deployment.