mlsanityData to 7 Oct 2026

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

arXiv 2603.00040 ↗First posted 9 Feb 2026cs.LG cs.AIcited by 5

Not on the map. Why?

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/.

Key facts. Cited by 5 papers in the map (data through 7 Oct 2026). Builds on 35 of its 44 references that are in the map. First posted on arXiv 9 Feb 2026.

Builds on

35 of 44 references.

References · 35

  1. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
  2. 2024The Llama 3 Herd of Modelscited by 14,192
  3. 2021Training Verifiers to Solve Math Word Problemscited by 9,743
  4. 2025Qwen3 Technical Reportcited by 8,195
  5. 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,767
  6. 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,203
  7. 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
  8. 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
  9. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  10. 2016Pointer Sentinel Mixture Modelscited by 2,992
  11. 2019PIQA: Reasoning about Physical Commonsense in Natural Languagecited by 2,818
  12. 2023GPQA: A Graduate-Level Google-Proof Q&A Benchmarkcited by 2,748
23 more
  1. 2025Wan: Open and Advanced Large-Scale Video Generative Modelscited by 2,536
  2. 2013Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computationcited by 2,031
  3. 2023VBench: Comprehensive Benchmark Suite for Video Generative Modelscited by 1,755
  4. 2017Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inferencecited by 1,548
  5. 2022SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Modelscited by 1,338
  6. 2023AWQ: Activation-aware Weight Quantization for LLM Compression and Accelerationcited by 1,234
  7. 2023Instruction-Following Evaluation for Large Language Modelscited by 1,143
  8. 2018HAQ: Hardware-Aware Automated Quantization with Mixed Precisioncited by 426
  9. 2019Data-Free Quantization Through Weight Equalization and Bias Correctioncited by 328
  10. 2023LLM-QAT: Data-Free Quantization Aware Training for Large Language Modelscited by 320
  11. 2019Differentiable Soft Quantization: Bridging Full-Precision and Low-Bit Neural Networkscited by 234
  12. 2019Understanding Straight-Through Estimator in Training Activation Quantized Neural Netscited by 225
  13. 2024Are We Done with MMLU?cited by 178
  14. 2024SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Accelerationcited by 158
  15. 2024SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantizationcited by 113
  16. 2023FP8-LM: Training FP8 Large Language Modelscited by 75
  17. 2025Pretraining Large Language Models with NVFP4cited by 67
  18. 2024Scaling FP8 training to trillion-token LLMscited by 63
  19. 2025SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Trainingcited by 63
  20. 2025Optimizing Large Language Model Training Using FP4 Quantizationcited by 53
  21. 2025FP4 All the Way: Fully Quantized Training of LLMscited by 37
  22. 2024BitNet a4.8: 4-bit Activations for 1-bit LLMscited by 16
  23. 2025Towards Fully FP8 GEMM LLM Training at Scalecited by 11

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 5.

5 papers.

  1. 2026Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attentioncited by 7
  2. 2026MXAttention: Data-Free Optimal Scaling and Pre-Normalization Quantization for MXFP4 Attentioncited by 2
  3. 2026VC-Attention: Value Smoothing and Softmax Casting for Low-bit Attentioncited by 1
  4. 2026Activation Outliers Matter: Robust Recovery for Quantized Multimodal LLMscited by 0
  5. 2026JustQuant: You Don't Need Smoothing, SVD, or Rotation for 4-Bit Activation Quantizationcited by 0