🤖 machine learning

Hardware-Aware FP4 FlashAttention-4

This paper introduces Hardware-Aware FP4 FlashAttention-4, a method that overcomes the limitations of Blackwell's FP4 tensor cores by employing Direct-P for noncausal inference and a causal path with FP8 gradients, achieving up to 2.13× faster forward throughput and 1.14× faster single-GPU training updates while avoiding the divergence issues observed in MXFP4 training.

Original authors: Robert Hu

Published 2026-09-04✓ Author reviewed ⓘ

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Robert Hu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of artificial intelligence, the most powerful models rely on a mechanism called "attention" to understand context. Imagine a reader scanning a long document; attention allows the model to focus on the most relevant words in a sentence while ignoring the rest, connecting ideas across vast distances. To do this, the computer performs a massive series of calculations, comparing every word against every other word to determine their relationships. For years, these calculations have been performed with high-precision numbers, similar to using a ruler with tiny, precise markings to measure a building. This ensures the model learns correctly, but it is slow and consumes enormous amounts of energy. As models grow larger, the need for speed becomes critical. Researchers have recently developed chips capable of using much smaller, coarser numbers—four-bit floating-point values—to perform these comparisons much faster. However, simply switching to these smaller numbers is not enough; the software that manages the flow of data must also change, or the speed gains disappear.

A researcher at Graphcore has tackled this specific bottleneck with a new method they call Direct-P. Their work focuses on the "forward" pass of attention, where the model processes information to generate an answer, and the "backward" pass, where it learns from its mistakes. The researcher discovered that while the new chips can multiply numbers incredibly quickly, the step in between—converting raw scores into probabilities—was slowing everything down. It was like having a super-fast assembly line that kept stopping because a worker had to carefully measure every single part before passing it on. The researcher found that the standard way of handling this conversion, which involved complex scaling and rounding, created a traffic jam that negated the speed benefits of the new hardware.

To solve this, the researcher redesigned the conversion process to be direct and streamlined. Instead of calculating a precise probability and then rounding it, their method maps the raw scores directly to the specific codes the hardware uses. This eliminates the middleman steps that were causing delays. When they tested this approach on the latest generation of data-center chips, the results were striking. For certain common shapes of data, the new method processed information more than twice as fast as the previous standard, which relied on higher-precision numbers. It achieved this while keeping the output accurate enough for complex tasks like video generation and language understanding. In tests involving a video generation model, the new method was nearly twice as fast as the standard approach, producing results that, while finite and usable, showed measurable drift and lower similarity scores compared to the slower, more precise version.

However, the story is not just about speed; it is also about what happens when the model tries to learn. The researcher explored whether they could use these ultra-fast, low-precision numbers for the entire training process, including the backward pass where the model updates its knowledge. They found that while the forward pass could run at full speed, the backward pass required a careful compromise. When they tried to use the fastest four-bit format for the probability values during training, the learning process became unstable and the model failed to improve. The numbers became too coarse, causing the learning signal to break down. Consequently, the researcher determined that for training to remain stable, they must use a slightly more precise eight-bit format for the probability values, even if the rest of the calculation uses the faster four-bit format. This hybrid approach allowed them to speed up the entire training cycle by about 14 percent on a single powerful chip, a significant gain for large-scale models.

The researcher also looked at the physical limits of the computer chips themselves. They discovered that the speed of the calculation was no longer the main problem; the issue was how the data was stored and moved around inside the chip. The chip has a small, ultra-fast memory area where it holds the numbers while working on them. The researcher found that this memory space was completely full, leaving no room to overlap different stages of the calculation. It was like a kitchen where the counter is so crowded with ingredients that the chef cannot start chopping the next vegetable until the current one is finished, even if the knife is incredibly sharp. Because the memory space was fully occupied, the chip could not work on multiple steps at once, which limited how much faster the whole process could become.

This work highlights a crucial shift in how we think about making artificial intelligence faster. It is not enough to simply build faster calculators; the entire workflow must be redesigned to match the hardware's capabilities. The researcher showed that by changing how probabilities are calculated and by carefully managing the flow of data, they could unlock massive speed improvements. Yet, they also proved that there are hard limits. The memory constraints of the chip mean that even with the fastest possible math, there is a ceiling to how much overlap is possible. The path forward, they suggest, lies not just in faster arithmetic, but in smarter ways to organize the data so that the chip's resources are never left waiting. This balance between raw speed and careful data management is what will allow the next generation of artificial intelligence to run efficiently.

=== SUMMARY ===

Technical Summary: Hardware-Aware FP4 FlashAttention-4

Problem Statement

NVIDIA Blackwell GPUs introduce 4-bit floating-point (FP4) tensor cores that significantly accelerate matrix multiplication. However, in attention mechanisms, the speedup is often bottlenecked not by the matrix products ( $QK^T$ and $PV$ ), but by the intermediate softmax stage. This stage involves score reduction, exponential evaluation, probability construction, and scaling, which must complete before the second matrix product can begin.

Existing attempts to utilize FP4 for the entire attention block face two primary hurdles:

Timing: Constructing the probability tensor ( $P$ ) from scores ( $S$ ) creates a critical path dependency that delays the $PV$ product.
Numerical Range: FP4 formats (specifically NVFP4 and MXFP4) use a 4-bit payload (E2M1) with limited dynamic range. Standard softmax probabilities span a range that is awkward for this payload, leading to underflow, "zero scales" (where entire blocks vanish), or significant precision loss if not carefully managed.

Furthermore, in causal training, the backward pass typically requires reconstructing probabilities from saved states. Approaches that save full-precision scores or recompute them in BF16 negate the potential speedups of low-precision training.

Methodology: Direct-P and Causal Backward

The paper proposes Direct-P, a method to enable full-FP4 forward inference and a quantized causal backward path on Blackwell hardware (GB200/B300).

1. Direct-P for Forward Inference

Direct-P modifies the interval between the completion of the $QK^T$ score fragment and the issuance of the $PV$ instruction. It retains the existing HAO AI Lab two-query pipeline and Tensor Memory (TMEM) ownership protocol but optimizes the probability generation:

Direct Mapping: Instead of computing a high-precision exponential and then rounding, Direct-P treats probability construction as a code classification problem. It maps normalized log-scores directly to the 8 possible non-negative E2M1 magnitudes using an affine classifier (or selective native exponentials). This shortens the critical path.
Consistent Normalization: The denominator (normalizer) for the softmax is accumulated directly from the rounded FP4 probability codes and block scales that will be consumed by the $PV$ instruction. This ensures the numerator and denominator describe the same approximate operator, avoiding mismatches between high-precision sums and low-precision payloads.
Range Guarding: For layers with extreme logits (e.g., >500), a sampled anchor mechanism prevents subnormal underflows without requiring a full re-scan of the scores.
Format Selection: The method uses NVFP4 for $Q$ and $K$ (due to fine-grained scaling) and MXFP4 for $P$ and $V$ (due to its power-of-two block scales which align well with the $N_{32}$ producer fragments).

2. Quantized Causal Backward

For training, the paper introduces a path where the forward pass saves the quantized $Q/K$ payload, scales, and the log-sum-exp (LSE) normalizer.

Reconstruction: The backward pass reconstructs probabilities directly from these saved quantized operands rather than creating a separate BF16 score path.
Gradient Formats: The projection epilogues publish row/column-oriented FP8 views of $Q, K, V$ to feed the gradient matrix products.
Stability Constraint: While the forward path uses MXFP4 for $P/V$ , distributed training experiments revealed that MXFP4 P/V trajectories diverge. Consequently, the retained training route uses FP8 for $P$ and $V$ in the backward pass to ensure stability, while keeping $Q/K$ in NVFP4.

Key Contributions

Direct-P Implementation: A method that maps scores directly to MXFP4 probability codes and normalizes using the same represented values, achieving up to 2.13× BF16 forward throughput on GB200 for favorable shapes, with a geometric mean speedup of 2.02× across the tested D128 grid.
Quantized Causal Training Path: A backward pass that reconstructs probabilities from saved forward quantization states, accelerating a complete single-GPU 8-billion-parameter update by up to 1.14×.
Hardware Profiling: Identification of Tensor Memory (TMEM) ownership as the primary limit on overlap. The study demonstrates that simply having faster tensor cores is insufficient if the storage layout (512 columns at D128) prevents the next $QK$ operation from overlapping with the current $PV$ phase.

Results

Forward Inference Performance

Throughput: On GB200, the "fast" Direct-P route achieves a geometric mean speedup of 2.02× over HAO's BF16 baseline, peaking at 2998 TFLOP/s (S32768/H24).
Accuracy: Direct-P trades some precision for speed. On Gaussian softmax tests, the "fast" route has a mean cosine similarity of ~0.944 against BF16, compared to ~0.990 for HAO's NV/FP8 route.
Model Impact:
- ViT/BERT: On fixed-input evaluations, the "fast" route maintains task accuracy (e.g., ViT S4096 top-1 accuracy matches BF16 at 88.5%) despite lower operator-level cosine similarity.
- Wan Video Diffusion: The method remains finite and functional for 20-step diffusion. However, error accumulates over steps; the 14B model at 20 steps shows a relative-L2 of ~0.53 compared to BF16, indicating drift but functional generation.

Training Performance

Single-GPU: The quantized route (NVFP4 projections + FP8 P/V) accelerates the full attention sublayer by 1.25× and the complete 8B model update by 1.14× at batch size 4.
Distributed Training: In a matched 64-GPU experiment over 100 billion tokens, the FP8 P/V route achieved a 1.112× throughput improvement over BF16.
Stability: The paper explicitly notes that MXFP4 P/V trajectories diverged in distributed settings. The successful training route relies on FP8 for the $P/V$ product, meaning the training path is not "pure" FP4 end-to-end.

Significance and Claims

The paper claims that Direct-P successfully makes full-FP4 forward attention practical by shortening the critical path between matrix products. It demonstrates that while FP4 matrix cores are faster, the bottleneck has shifted to data movement, probability construction, and memory ownership.

Key Takeaways:

Hardware Limits: The primary barrier to further speedups is not arithmetic throughput but TMEM ownership. The current D128 layout consumes all 512 TMEM columns, preventing the next $QK$ operation from overlapping with the tail of the $PV$ phase.
Readiness Stalls: Low tensor-core activity is often a symptom of operand readiness (waiting for probability construction) rather than a lack of compute capacity.
Training Trade-offs: While FP4 forward inference is viable, stable causal training currently requires a concession to FP8 for the probability/value operands to prevent divergence. The paper suggests that future formats like Unsigned E5M3 (UE5M3) might enable fully FP4 training, but this is not evaluated here.

The work concludes that faster matrix arithmetic alone is insufficient; achieving higher efficiency requires architectural changes that allow for larger allocatable overlap windows (e.g., additional score destinations) alongside optimized low-precision algorithms.