Attn-QAT: 4-Bit Attention With Quantization-Aware Training

TL;DR

Attn-QAT is the first systematic study of quantization-aware training for the attention operator on FP4 hardware (NVIDIA Blackwell). Naively quantizing attention to FP4 — even with prior outlier-mitigation heuristics like SageAttention3’s Q/K smoothing and two-level P quantization — visibly degrades video generation and language-model benchmark scores. Attn-QAT identifies two precision-consistency fixes that make standard QAT actually work for FlashAttention-style fused kernels: (1) storing an auxiliary high-precision output that restores the FlashAttention backward identity when the forward uses fake-quantized P and V, and (2) recomputing attention probabilities in the same low precision during the backward pass. The resulting model matches BF16 attention quality on Wan-2.1-14B video diffusion and on Qwen3-14B / Llama-3.1-70B continued-pretraining and SFT, while delivering 1.1×–1.5× higher throughput than SageAttention3 on RTX 5090 and up to 1.39× over FlashAttention-4 on B200 via a paired NVFP4 FA4 kernel.

Key claims

Training-free FP4 attention methods (e.g. SageAttention3) still degrade Wan-2.1-14B video quality compared to BF16 attention even with Q/K smoothing and two-level P quantization, motivating a training-time fix rather than yet more inference-side outlier heuristics [Blog §“4-bit attention remains unsolved”].
Naive QAT for attention fails because FlashAttention’s memory-efficient backward uses the algebraic identity $\mathbf{P}_i^\top \mathbf{dP}_i = \mathbf{dO}_i^\top \mathbf{O}_i$ , which implicitly assumes the forward computed $\mathbf{O}_i = \sum_j \mathbf{P}_{ij}\mathbf{V}_j$ ; under fake-quantized $\mathbf{P}^F, \mathbf{V}^F$ this identity breaks and gradients become inconsistent [Blog §“The two fixes that make Attn-QAT work” — §1].
Fix 1: in the forward pass, materialize an auxiliary output $\mathbf{O}'_i = \sum_j \mathbf{P}_{ij}\mathbf{V}^F_j$ (high-precision $\mathbf{P}$ , fake-quantized $\mathbf{V}^F$ stored in BF16) and use $\mathbf{O}'$ in the backward identity; this adds ~25% storage and restores the exact identity without materializing the full $N\times N$ attention matrix [Blog §1].
Fix 2: during the backward recomputation of $\mathbf{P}$ from saved logsumexp statistics, fake-quantize the recomputed probabilities to match the low-precision forward; empirically this stabilizes gradient norms [Blog §“Recompute attention probabilities in the same low precision used in the forward pass”].
Together these two fixes eliminate the need for SageAttention3-style outlier-mitigation heuristics: Attn-QAT recovers BF16 quality without Q/K smoothing or two-level P quantization [Blog §“Final thoughts”].
On Wan-2.1-14B video diffusion (480p / 720p synthetic latents), human evaluators on 99 VBench prompts cannot distinguish BF16-attention from FP4-attention-after-QAT videos, while SageAttention3 videos show visible artifacts [Blog §“Results: Quality is Recovered”].
For LLM continued pretraining on C4, Attn-QAT recovers most of the FP4-attention quality drop on Qwen3-14B and partially on Llama-3.1-70B; the remaining 70B gap is attributed to training-budget limits rather than method failure [Blog §“Results”].
For LLM SFT on Dolci-Instruct-SFT, Attn-QAT serves as a drop-in BF16-attention replacement: nearly identical downstream benchmark performance on Qwen3-14B, small gap on Llama-3.1-70B — i.e. usable in standard fine-tuning pipelines, not only as a recovery stage [Blog §“Results”].
Removing the SageAttention3 smoothing + two-level-P preprocessing yields 1.1×–1.5× higher throughput than SageAttention3 on RTX 5090 [Blog §“Faster Inference on an RTX 5090”].
FlashAttention-4 FP4 — a paired NVFP4-quantized FA4 kernel in CuTeDSL — reaches up to 1801 TFLOPs/s and 1.39× over FA4 BF16 on B200 by running block-scaled NVFP4 on the $\mathbf{QK}$ MMA while leaving $\mathbf{PV}$ in BF16 [Blog Part 2, Table — kernel benchmarks].
Quantizing $\mathbf{PV}$ to NVFP4 in addition to $\mathbf{QK}$ actually slows the kernel on B200 because softmax is already the bottleneck (MUFU exp throughput and CUDA-core scale factor computation) and the extra NVFP4 group-quantization adds to that bottleneck — the design choice is “FP4 only where it’s free” [Blog §“TMEM overlap schedule”].
The Attn-QAT kernel achieves 2–2.5× lower max absolute error vs. an FP8 non-block-scaled FA4 variant despite using NVFP4, due to a finer group size (16 vs 128) [Blog §“Precision Results”, precision table].
The TMEM-overlap schedule for B200 keeps scale-factor loads non-stalling by reusing the S1/S2 TMEM regions for sfqk scale factors (8.5KB total) under the constraint that 128×128 tiles already saturate FA4’s TMEM budget [Blog §“TMEM overlap schedule”].
Blackwell’s headline FP4 PFLOPS is bottlenecked by softmax exp throughput in attention workloads: BF16/FP8 MMA throughput roughly doubles vs Hopper while CUDA cores and MUFU exp throughput stay flat, so quantizing the matmul further past FP4-QK gives diminishing returns until B300/Rubin double or quadruple exp throughput [Blog §“Bloated Tensor Cores and the Softmax Bottleneck”, Table].

Method

The forward pass is standard fake-quantized FlashAttention: $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ are fake-quantized to FP4 (NVFP4 E2M1, group size 16); the row-wise softmax over $\mathbf{S} = \mathbf{Q}^F (\mathbf{K}^F)^\top$ is held in FP32, then $\mathbf{P}$ is fake-quantized to $\mathbf{P}^F$ and combined with $\mathbf{V}^F$ to produce the low-precision output $\mathbf{O}$ . The novelty is the auxiliary high-precision output $\mathbf{O}' = \mathbf{P}\mathbf{V}^F$ materialized alongside $\mathbf{O}$ (~25% extra storage), which is fed into the backward pass so the FlashAttention identity $\mathbf{P}_i^\top \mathbf{dP}_i = \mathbf{dO}_i^\top \mathbf{O}'_i$ remains exact under fake quantization. The backward additionally fake-quantizes the recomputed $\mathbf{P}$ to match the forward’s low precision, stabilizing gradient norms. The training-side recipe is paired with a CuTeDSL B200 kernel (FlashAttention-4 FP4) that runs block-scaled NVFP4 on $\mathbf{QK}$ + BF16 on $\mathbf{PV}$ ; the design choice is dictated by Blackwell’s softmax bottleneck, where additional MMA-side quantization adds CUDA-core overhead that net slows the kernel. A TMEM-overlap schedule reuses the FA4 $\mathbf{S}_1/\mathbf{S}_2$ TMEM regions for the sfqk scale factors and inserts only minimal barriers around T2R copies to keep the warp-specialized pipeline from stalling. The whole stack ships in FastVideo and a dedicated flash-attention-fp4 repo.

Results

On Wan-2.1-14B (text-to-video), human evaluators on 99 randomly sampled VBench prompts could not distinguish BF16-attention from FP4-Attn-QAT video. SageAttention3 — the strongest prior training-free FP4 attention — visibly degrades quality on the same model. For LLMs, Attn-QAT recovers “most” of the FP4-attention quality drop on Qwen3-14B continued pretraining (C4) and partially recovers it on Llama-3.1-70B; for SFT on Dolci-Instruct-SFT, Qwen3-14B benchmark scores under Attn-QAT match BF16 closely, with a small Llama-3.1-70B gap. Kernel-level: 1.1×–1.5× over SageAttention3 on RTX 5090; FA4-FP4 on B200 reaches 1801 TFLOPs/s and 1.39× over FA4 BF16 at b=1 s=32768 h=16 d=128, with 1.2–1.3× across most large-context configurations. NVFP4 max absolute error is 2–2.5× lower than an FP8 non-block-scaled FA4 baseline at the same configurations.

Why it’s interesting

Attn-QAT is the first filed paper to argue that training-time and low-precision kernel design must be co-designed for attention, where for matmuls the QAT recipe has been mostly transferable. The two precision-consistency fixes are the kind of finding that doesn’t show up unless you actually trace which algebraic identities the fused backward kernel quietly assumes — exactly the same flavor of insight that drives FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling‘s shift from “tensor cores are scarce” to “softmax + SMEM are scarce.” This paper is the training-side complement to that kernel story from the same broader Dao-AILab / Hao AI Lab orbit, and the B200 kernel it ships is a literal FP4 variant of FA4 built in CuTeDSL. Practically, the result that an FP4 attention pass can survive both video diffusion (Wan-2.1-14B) and 70B LLM fine-tuning without outlier heuristics is a step toward end-to-end FP4 serving for Luma-relevant workloads — long-context video DiTs in particular, where the attention bill dominates wall-clock. The B200 design choice “FP4 only where it’s free” (QK only, leave PV in BF16) also operationalizes the softmax-bottleneck finding from FA4 into a concrete kernel: the marketed FP4 PFLOPS won’t translate to attention speedups until B300 / Rubin lift the exp-throughput ceiling. Worth contrasting with MonarchRT: Efficient Attention for Real-Time Video Generation, which tackles the same “attention is no longer matmul-bound on Blackwell” observation by changing the algorithm (Monarch sparsity, 95% effective sparsity at quality parity) instead of the precision; Attn-QAT and MonarchRT are orthogonal levers on the same bottleneck.