[論文レビュー] Attn-QAT: 4-Bit Attention With Quantization-Aware Training

Peiyuan Zhang, Matthew Noto+5 more|arXiv (Cornell University)|Feb 9, 2026

Advanced Neural Network Applications被引用数 0

ひとこと要約

本論文は Attn-QAT を提示する。これは注意機構の 4-bit 量子化対応訓練法で、安定した訓練と BF16 レベルのモデル品質を実現する FP4 注意を可能にし、RTX 5090 で 1.1x–1.5x のスループット向上を達成する。

ABSTRACT

Achieving reliable 4-bit attention is a prerequisite for end-to-end FP4 computation on emerging FP4-capable GPUs, yet attention remains the main obstacle due to FP4's tiny dynamic range and attention's heavy-tailed activations. This paper presents the first systematic study of 4-bit quantization-aware training (QAT) for attention. We find that "drop-in" QAT, which naively combines an FP4 forward pass with a high-precision Flash Attention (FA)-style backward pass, leads to training instability. We identify two key principles for stable FP4 attention: (1) matching low-precision recomputation of attention scores in the backward pass, and (2) resolving implicit precision assumptions in FA's gradient calculation. Based on these insights, we propose Attn-QAT and implement fused Triton kernels for training as well as FP4 inference kernels. Across diffusion and language models, Attn-QAT recovers the quality drop from FP4 attention without explicit outlier-mitigation heuristics used in prior FP4 attention, and delivers up to a 1.5x speedup on an RTX 5090. Video demos can be found at https://drive.google.com/drive/folders/190F6xbBDUF2kGQYIcXBt3ehSYij5jlim?usp=sharing.

研究の動機と目的

• FP4 対応 GPU でエンドツーエンドの FP4 計算を可能にするための 4-bit 注意の必要性を動機づける。
• 既存の後向き伝播での素朴な FP4 注意がなぜ不安定になるのかを調査する。
• QAT 中の安定した FP4 注意のための精度協調要件を特定する。
• 専門的な前方/後方の精度処理とカーネルを用いた Attn-QAT を提案する。
• Attn-QAT が拡散モデルと言語モデルの品質回復と、速度向上をもたらすことを示す。

提案手法

• FP4 注意の二つの主要課題：極端に粗い値域と裾園分布を持つ活性化を分析する。
• Forward パスで FP4 を模擬し、勾配を高精度に維持することで FlashAttention 風の結合型演算子に対して量子化対応訓練を適用する。
• Backward 再計算による注意スコアが forward と同じ低精度を使用することを保証する。
• 正しい softmax 勾配計算のための高精度の補助出力を提供する。
• 訓練用の前方・後方 Triton カーネルと、デプロイ用の FP4 推論カーネルを実装する。

Figure 1 : Both NVFP4 attention and SageAttention3 suffer from a significant quality drop on Wan 2.1 14B. Our proposed method, Attn-QAT, recovers the quality drop by using quantization-aware training. Note that temporal inconsistency is hard to visualize in sampled frames. We attach video samples in

実験結果

リサーチクエスチョン

• RQ1 FP4 における4-bit 注意を量子化対応訓練で安定化しつつ、モデル品質を保持できるか？
• RQ2 FP4 注意を FlashAttention 風のバックエンドと整合させるためには前方と後方のどの精度協調が必要か？
• RQ3 Attn-QAT は拡散モデルと言語モデルで BF16 品質を回復し、SageAttention3 と比較してどうか？
• RQ4 現代の GPU で FP4 注意によるスループットの性能向上はどの程度達成できるか？

主な発見

• Attn-QAT は FP4 注意によって生じる品質低下を回復し、評価指標全体で BF16 と同等の性能を達成する。
• QA 結果は Attn-QAT が SageAttention3 より優れており、外れ値対処のヒューリスティクスを排除できることを示す。
• 訓練の安定性は二つの設計選択に依存する： (1) 前方パスの注意スコアの再計算を FP4 で行う、(2) 正確な softmax 勾配のための高精度出力。
• Attn-QAT は RTX 5090 上で FP4 注意を用いた場合、SageAttention3 より約 1.1x–1.5x の速度向上を実現する。
• LLM の継続学習において、Attn-QAT は Qwen3-14B でほぼ BF16 品質を回復し、Llama3.1-70B の性能も一部回復させ、より長い訓練でさらなる改善の余地がある。
• カーネルベンチマークでは、Trition 訓練カーネルと CUDA 推論カーネルの前方出力がほぼ同一であることを示す。

Figure 2 : Win–Tie–Lose blind human evaluation on 99 randomly sampled VBench prompts for Wan 2.1 14B. Attn-QAT matches BF16 attention in perceived visual quality.