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

arXiv: https://arxiv.org/abs/2603.00040 · HTML: https://mlsanity.com/paper/2603.00040/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 44 references that are in the map.
- First posted on arXiv 9 Feb 2026.

## Builds on (most cited first, 30 of 35)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [Wan: Open and Advanced Large-Scale Video Generative Models](https://mlsanity.com/paper/2503.20314/) (2025)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [VBench: Comprehensive Benchmark Suite for Video Generative Models](https://mlsanity.com/paper/2311.17982/) (2023)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2211.10438/) (2022)
- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](https://mlsanity.com/paper/2306.00978/) (2023)
- [Instruction-Following Evaluation for Large Language Models](https://mlsanity.com/paper/2311.07911/) (2023)
- [HAQ: Hardware-Aware Automated Quantization with Mixed Precision](https://mlsanity.com/paper/1811.08886/) (2018)
- [Data-Free Quantization Through Weight Equalization and Bias Correction](https://mlsanity.com/paper/1906.04721/) (2019)
- [LLM-QAT: Data-Free Quantization Aware Training for Large Language Models](https://mlsanity.com/paper/2305.17888/) (2023)
- [Differentiable Soft Quantization: Bridging Full-Precision and Low-Bit Neural Networks](https://mlsanity.com/paper/1908.05033/) (2019)
- [Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets](https://mlsanity.com/paper/1903.05662/) (2019)
- [Are We Done with MMLU?](https://mlsanity.com/paper/2406.04127/) (2024)
- [SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration](https://mlsanity.com/paper/2410.02367/) (2024)
- [SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization](https://mlsanity.com/paper/2411.10958/) (2024)
- [FP8-LM: Training FP8 Large Language Models](https://mlsanity.com/paper/2310.18313/) (2023)
- [Pretraining Large Language Models with NVFP4](https://mlsanity.com/paper/2509.25149/) (2025)
- [SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training](https://mlsanity.com/paper/2505.11594/) (2025)

## Built on by (most cited first, 8 of 8)

- [Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention](https://mlsanity.com/paper/2607.04422/) (2026)
- [MXAttention: Data-Free Optimal Scaling and Pre-Normalization Quantization for MXFP4 Attention](https://mlsanity.com/paper/2607.24377/) (2026)
- [Hardware-Aware FP4 FlashAttention-4](https://mlsanity.com/paper/2609.04105/) (2026)
- [VC-Attention: Value Smoothing and Softmax Casting for Low-bit Attention](https://mlsanity.com/paper/2609.15810/) (2026)
- [Activation Outliers Matter: Robust Recovery for Quantized Multimodal LLMs](https://mlsanity.com/paper/2608.26581/) (2026)
- [JustQuant: You Don't Need Smoothing, SVD, or Rotation for 4-Bit Activation Quantization](https://mlsanity.com/paper/2609.33601/) (2026)
- [Delta-Matching: Closing the Final Gap of Native 8-bit Training for LLMs](https://mlsanity.com/paper/2609.37852/) (2026)
- [TurboPairFormer: Fast and Stable Protein Folding Model Training with an Optimized Triangle Attention Kernel](https://mlsanity.com/paper/2610.05854/) (2026)

## Abstract (arXiv)

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 ``drop-in'' QAT -- naively combining an FP4 forward pass with 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 and 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 over SageAttention3 and up to a 1.74x speedup over FA4 on a GB300. Video demos can be found at https://drive.google.com/drive/folders/190F6xbBDUF2kGQYIcXBt3ehSYij5jlim?usp=sharing. Code can be found at https://haoailab.com/FastVideo/training/attn_qat/.