Hardware-Aware FP4 FlashAttention-4

Abstract: Blackwell's 4-bit floating-point (FP4) tensor cores do not automatically make attention faster because softmax conversion and on-chip dependencies dominate once its matrix products shrink. We address this with \emph{Direct-P} for noncausal inference and a causal path that passes the forward quantization directly into backward. Direct-P maps scores directly to FP4 probabilities and reaches up to 2.13$\times$ the bfloat16 (BF16) forward throughput on an NVIDIA GB200. The causal path reconstructs probabilities from saved quantized queries and keys and uses 8-bit floating-point (FP8) gradient operands, accelerating a complete single-GPU 8-billion-parameter update by up to 1.14$\times$. Matched distributed training retains FP8 probabilities and values; every tested MXFP4 probability/value training trajectory diverges.

• The paper evaluates whether FP4 tensor cores can accelerate attention beyond BF16 through the development of Direct-P, a new method for efficient probability generation Direct-P achieves up to 2.125x speedup in noncausal inference but with a 5% accuracy tradeoff in certain metrics. Noncausal inference with Direct-P is reported to reach 2998 TFLOP/s on a benchmark, though at a slight trade-off of accuracy. Quantized causal backward improves projection-inclusive attention by up to 1.25x and complete single-GPU updates by up to 1.14x, but FP8 P/V is retained for training stability. Direct-Psic XP4 FlashAttention addresses the limitations of Blackwell core tensor memory

Problem formulation and central claims

“Hardware-Aware FP4 FlashAttention-4” (2609.04105) investigates whether Blackwell’s FP4 tensor cores can accelerate attention beyond the gains obtained by simply replacing BF16 matrix products with FP4 operations. The paper’s central argument is that this substitution is insufficient: attention contains two matrix multiplications, $QK^\mathsf{T}$ and $PV$, separated by online softmax, probability quantization, scale publication, synchronization, and tensor-memory ownership. Once the matrix products become faster, these intermediate operations become a critical-path bottleneck.

The work separates two objectives. For noncausal inference, it develops Direct-P, a probability-generation path that maps normalized score fragments directly to MXFP4 E2M1 codes rather than first evaluating a higher-precision exponential and subsequently discarding most of that precision. For causal training, it passes the forward quantization state into backward, allowing the backward kernel to reconstruct probabilities from saved quantized $Q/K$ payloads, scales, and log-sum-exp statistics.

The principal claims are:

• Direct-P achieves up to $2.13\times$ the BF16 forward throughput on GB200 at favorable D128 attention shapes.
• Its speed advantage is accompanied by a substantial accuracy trade-off: the mean forward cosine is approximately $0.944$ for the fast policy, compared with approximately $0.990$ for the higher-precision NVFP4-QK/FP8-PV control.
• Quantized causal backward improves projection-inclusive attention by up to $1.25\times$ and a complete single-GPU 8-billion-parameter update by up to $1.14\times$.
• FP8 P/V is retained for training because every tested MXFP4 P/V trajectory diverges in the reported long-running experiments.
• The principal hardware limitation is not nominal FP4 tensor throughput but the absence of another legally allocatable score destination in Blackwell’s tensor memory (TMEM).

These results extend the hardware–algorithm co-design perspective established by FlashAttention (Dao et al., 2022), FlashAttention-2 (Dao, 2023), FlashAttention-3 (Shah et al., 2024), and FlashAttention-4 (Zadouri et al., 5 Mar 2026).

Architectural bottleneck: online softmax and TMEM ownership

FlashAttention avoids materializing the quadratic score and probability matrices in HBM by processing key tiles incrementally. For each query row, it maintains a running maximum, denominator, and output accumulator. This eliminates large memory traffic, but introduces a strict dependency chain: a score fragment must be reduced, normalized, converted into the representation consumed by $PV$, and published before the value product can proceed.

The implementation studied in the paper inherits a two-query Blackwell schedule in which one 16-warp CTA processes two query stages. QK and PV are issued through asynchronous tensor-core operations, while specialized warpgroups perform score processing, online correction, and output publication. At D128, each FP32 score or output tile occupies 128 TMEM columns. Two score banks and two persistent output banks consume all 512 available logical TMEM columns.

This layout creates a storage-ownership dependency. A score bank is temporarily reused in the sequence

$$\text{QK score} \rightarrow \text{probability and scale overlay} \rightarrow \text{PV consumption} \rightarrow \text{bank release}.$$

The next QK operation cannot overwrite the bank while its probability overlay is still needed by PV. Additional barriers can expose readiness, but they cannot create another legal destination. The paper therefore distinguishes a synchronization problem from an allocation problem: the limiting resource is not merely barrier expressiveness or shared-memory capacity, but the number of usable TMEM score destinations.

The K64 granularity of the scaled-FP4 PV instruction compounds this constraint. Four N32 score quarters are generated, but two adjacent quarters must be complete before the first K64 PV operation can be issued. Consequently, accelerating one probability fragment does not necessarily advance tensor-core work unless its paired fragment is also ready.

Direct-P and the forward FP4 path

Direct-P changes only the interval between a completed score fragment and a legal PV operand. The outer CTA schedule, score-bank reuse, two-query pipeline, and publication protocol remain unchanged. Its design has three components.

First, Direct-P treats probability formation as E2M1 code classification rather than approximate exponential evaluation. Given a score transformed into base-two coordinates, it selects one of the nonnegative E2M1 magnitudes

$PV$0

The implementation uses packed FFMA2 operations followed by Blackwell’s native floating-point conversion instruction. The fitted affine map is selected to place the E2M1 decision boundaries efficiently, not to minimize error against the continuous exponential function. This distinction is important: the value eventually consumed by PV is a four-bit code, so optimizing an intermediate FP32 exponential is computationally wasteful when its precision will be discarded.

Second, Direct-P normalizes using the represented FP4 probabilities actually consumed by PV. For each N32 block, it accumulates the MXFP4 amplitude multiplied by the sum of the emitted E2M1 codes. The numerator and denominator therefore describe the same approximate operator. This avoids a mismatch in which the numerator uses rounded FP4 values while the denominator is computed from an unrelated higher-precision exponential approximation.

Third, the implementation includes a selective guard for extreme model logits. The fast shiftless path is finite on the synthetic benchmark grid but encounters logits above 500 and sometimes 1000 in later Wan layers. Rather than scanning every score row, the guarded path samples fixed key rows to obtain an anchor, applies a compile-time margin and scale floor, and reassociates the denominator computation to avoid subnormal intermediate underflow. This guard is not a general stable-softmax fallback; it is a targeted correction for observed model distributions.

The format choice reflects a latency–range trade-off. Stabilized NVFP4 provides better probability fidelity but requires finer scale handling and additional range correction. MXFP4 offers power-of-two E8M0 scales aligned with the N32 producer granularity, but its E2M1 representation is coarser. In the paper’s Gaussian diagnostic at sequence length 4096, MXFP4 produces a probability relative-$PV$1 error of approximately $PV$2 and a $PV$3 cosine of approximately $PV$4, whereas unscaled NVFP4 exhibits catastrophic underflow. Stabilized NVFP4 achieves a probability relative-$PV$5 error near $PV$6 and $PV$7 cosine near $PV$8.

Forward performance and accuracy

Across the nine-row GB200 D128 benchmark suite, the fast Direct-P policy is reported as $PV$9 faster than the HAO BF16 baseline, reaching a maximum of 2998 TFLOP/s. The accurate policy reaches $Q/K$0 and 2416 TFLOP/s. The largest explicitly reported relative speedup is $Q/K$1 at S8192/H64, where Direct-P takes 0.758336 ms compared with 1.611488 ms for BF16.

B300 results show 5.6–7.7% lower latency than GB200 on S4096–S8192 rows, with 3116 TFLOP/s at S8192/H64 and 3159 TFLOP/s at the wave-aligned S9472/H64 shape. This improvement is shape-dependent rather than a uniform consequence of the newer GPU. At S32768/H24, B300 reaches 2945 TFLOP/s, below the 2998 TFLOP/s reported on GB200. The paper attributes such variation to launch geometry, persistent-grid width, and the fixed TMEM allocation.

The performance gain is not free. The fast Direct-P policy has mean cosine approximately $Q/K$2 and mean relative-$Q/K$3 approximately $Q/K$4 against BF16 in the principal operator suite. The accurate policy improves these values to approximately $Q/K$5 and $Q/K$6, respectively, but remains considerably less accurate than the NVFP4-QK/FP8-PV route, whose mean cosine is approximately $Q/K$7.

The paper’s accuracy-matched control clarifies the source of this trade-off. An exact local NVFP4-QK/FP8-PV route reaches cosine $Q/K$8 but only 1490 TFLOP/s, whereas Direct-P reaches 2945 TFLOP/s at cosine $Q/K$9 on the corresponding B300 long-sequence case. Thus the reported speed is not simply a consequence of FP4 QK; it depends materially on reducing the probability-construction path and accepting a coarser P representation.

Model-level fixed-input results are more favorable than the standalone operator metrics. On ViT S4096, fast Direct-P preserves BF16 top-1 accuracy at 88.5%, with 95.5% prediction agreement, while the accurate policy reaches 89.0% and 98.5% agreement. Across 2272 classification examples, 32 fast-policy predictions change, and 31 of those occur in the lowest quartile of BF16 top-two logit margins. This supports a margin-sensitive interpretation of the error, although the evaluation is not sufficient to establish general inference or training safety.

In the Wan2.1 diffusion evaluation, all self-attention layers use Direct-P while the remainder of the model remains BF16. Fast Direct-P is 1.75× faster for the 1.3-billion-parameter model and 2.09× faster for the 14-billion-parameter model. However, error accumulates over diffusion steps: for Wan2.1-14B, the 20-step output has cosine 0.8496 and relative-$2.13\times$0 0.5337 against BF16. The selective guard enables all 1600 attention calls to complete, but guarded layers are 21–23% slower individually; because only a minority of layers are guarded, the aggregate penalty is approximately 2.2–2.3%.

The ViT-MAE reconstruction experiment shows a smaller end-task displacement. After replacing all twelve encoder attention layers, fast Direct-P obtains reconstruction cosine 0.99973 and relative-$2.13\times$1 0.0203, with a PSNR decrease of $2.13\times$2 dB. The result indicates that residual attention error can be attenuated by downstream network structure in some vision workloads, but this does not contradict the larger accumulated drift observed in long diffusion trajectories.

Quantized causal backward and training

The causal training method saves the forward NVFP4 Q/K payloads, block and global scales, and per-row LSE values. Backward reconstructs the represented scores and probabilities rather than generating a separate BF16 score path. Projection epilogues publish row- and column-oriented FP8 views needed by different gradient products, avoiding standalone transpose and quantization kernels.

The backward computation retains the standard dependency structure:

$2.13\times$3

followed by softmax centering and the $2.13\times$4 and $2.13\times$5 products. The implementation reuses the represented P state, publishes both physical $2.13\times$6 layouts, and uses E5M2 rather than E4M3 for $2.13\times$7. This choice is motivated by range rather than precision: E4M3 rounded approximately 97% of observed $2.13\times$8 values to zero in a failing diagnostic, while E5M2 reduced the zero fraction to approximately 14% with less than 1% publisher overhead.

The isolated reconstruction core reduces D128 causal-backward latency from 0.501 ms for BF16 to 0.356 ms, a $2.13\times$9 speedup. Once the E5M2 publisher and row-statistics publisher are included, latency increases to 0.508 ms, slightly slower than BF16. This result is important because it demonstrates that a faster inner backward kernel does not imply a faster training path; producer and publication overhead can erase the kernel-level gain.

The gain reappears when the method is integrated with projections and output gradients. At B1/S4096/D128, the projection-inclusive attention sublayer reaches:

Boundary	BF16	Quantized route	Speedup
Backward only	1.572 ms	1.397 ms	$0.944$0
Forward plus backward	2.656 ms	2.133 ms	$0.944$1

For a complete 8.03-billion-parameter update at S4096, speedup increases with local batch size: $0.944$2 at batch 1, $0.944$3 at batch 2, and $0.944$4 for FP8 P/V at batch 4. The MXFP4 P/V arm reaches $0.944$5 at batch 4, but FP8 and MXFP4 complete updates differ by at most 0.31%, so the isolated MXFP4 forward advantage does not produce a material end-to-end improvement.

The paper explicitly separates these timing results from training quality. Initial-logit cosine is only 0.416–0.426 for FP8 P/V and 0.373–0.374 for MXFP4 P/V in the short update experiment, and fixed-token timing does not establish convergence.

Why FP8 P/V is retained for training

The strongest negative result concerns MXFP4 P/V in longer training. In the reported four-arm diagnostic, both FP8-P/V trajectories remain non-divergent, while both MXFP4-P/V trajectories separate from their FP8 controls near 0.1 billion tokens and later exhibit rising loss and very large pre-clipping gradient norms. A matched B4 experiment with NVFP4 projections and MXFP4 P/V shows loss rising from 7.01 at update 325 to 16.25 at update 350.

The factorial structure strengthens the interpretation. The experiments vary learned projection precision while holding the attention backward path fixed. Both projection formats fail with MXFP4 P/V, whereas both remain stable with FP8 P/V. The evidence identifies P/V representation, or state changes induced by it, as the common separator. It does not establish whether the cause is forward probability quantization, the saved V payload, backward use of the representation, or an interaction among them.

The retained matched distributed study uses 64 GPUs, local batch four, four accumulation steps, and effective global batch 1024. The BF16 and NVFP4-projection/FP8-PV routes share model, optimizer, tokenizer, sample order, and token schedule over approximately 100 billion tokens. At the final scheduled training report, BF16 loss is 2.3095 and FP8-route loss is 2.3613. At the final same-update held-out validation report, losses are 2.3048 and 2.3948, respectively, yielding a gap of 0.0900.

The FP8 route is stable and descending but not numerically equivalent to BF16. Its median throughput is 24,303 tokens/s/GPU versus 21,853 for BF16, corresponding to a $0.944$6 median speedup. The paired throughput ratio has a 10th–90th percentile range of 1.080–1.114×.

Figure 1: Token-aligned training and held-out validation for the matched B4 experiment over the completed 100-billion-token schedule; the plot shows stable but non-identical FP8-P/V and BF16 trajectories without uncertainty estimates.

The training result therefore supports a narrower claim than “FP4 training.” The retained route uses NVFP4 in learned projections and Q/K attention operands, but FP8 P/V and FP8 gradient operands remain necessary under the reported stability evidence.

Hardware implications

The paper’s hardware analysis argues that Direct-P has already removed most of the exposed probability arithmetic. A historical matched diagnostic executes the same 98,304 tensor instructions with either real or fixed probability construction. Tensor-pipe activity rises from 18.8% to 26.6% when probability work is removed, but the tensor instruction count is unchanged. Source sampling attributes most not-issued samples to long-scoreboard dependencies involving final statistics, score readiness, and output publication.

A final speed-of-light diagnostic reinforces this conclusion. Relative to a valid 0.092448-ms kernel, retaining only raw score packing saves 1.38%, retaining row maxima saves 2.53%, and using a fixed probability tile saves 5.23%. Thus even eliminating nearly all probability construction cannot approach the nominal fourfold FP4 matrix-throughput ratio.

The paper proposes several hardware and instruction-set directions, but does not measure them as improvements:

• an additional allocatable score bank with compatible issue semantics;
• K32 scaled-FP4 PV instructions so each N32 probability fragment can be consumed immediately;
• scale storage outside TMEM;
• wider tiles with a compatible score, probability, and output lifecycle.

The first proposal is the most consequential. The paper’s analysis indicates that a larger usable overlap window would matter more than another polynomial approximation once Direct-P has shortened probability generation.

Limitations and open questions

The noncausal evidence consists primarily of fixed-input inference evaluations. It does not establish finetuning behavior, pretraining stability, or broad task robustness. The Wan results show substantial multi-step drift, especially for the 14-billion-parameter model, despite successful finite execution.

The distributed training result is based on one trajectory per route. It supports stability and throughput over the measured 100-billion-token schedule but provides no run-to-run variance estimate or statistical-equivalence claim. The FP8-versus-MXFP4 comparison also does not isolate the precise source of MXFP4 divergence. Determining whether the failure arises from forward P quantization, V quantization, backward reconstruction, or their interaction remains open.

The comparison boundaries also matter. Learned projection precision is separate from attention precision, and the reported end-to-end training route is not pure FP4. D64 and other head dimensions use different tile and TMEM ownership regimes, so the D128 conclusions cannot be transferred directly. Several B300 comparisons use different binaries or cross-run published baselines, which limits causal interpretation of hardware-generation differences. Finally, the paper’s hardware proposals are hypotheses: no additional score bank, K32 PV instruction, or external scale path is evaluated.

Conclusion

The paper establishes that FP4 attention performance on Blackwell is governed by the full score-to-PV pipeline rather than by tensor-core matrix throughput alone. Direct-P obtains up to $0.944$7 BF16 forward speed by mapping scores directly to MXFP4 probability codes and normalizing the represented probabilities consumed by PV. The resulting accuracy is materially below FP8-PV controls, although several fixed-input model evaluations show limited task-level impact.

For causal training, reusing the forward quantized state enables $0.944$8 projection-inclusive attention speedup and up to $0.944$9 complete-update speedup. However, the long-running experiments reject MXFP4 P/V for the reported training recipe and retain FP8 P/V. The dominant unresolved systems issue is TMEM ownership: at D128, the existing score and output allocation leaves insufficient space for a larger overlap window.