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

• URL: http://arxiv.org/abs/2603.00040v1
• Date: Mon, 09 Feb 2026 04:46:21 GMT
• Title: Attn-QAT: 4-Bit Attention With Quantization-Aware Training
• Abstract summary: This paper presents the first systematic study of 4-bit quantization-aware training (QAT) for attention.<n>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.<n>We propose Attn-QAT and implement fused Triton kernels for training as well as FP4 inference kernels.
• Score: 10.003298547571447
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• HiFloat4 Format for End-To-End Reinforcement Learning Post-Training of Large Language Models [9.149747042283197]
  We present the first end-to-end FP4 RL post-training, in which both the rollout and training policies operate at 4-bit precision.<n>A lightweight correction, Rollout-ResQ, recovers most of the precision lost to outlier-driven underflow without inflating the rollout's compute footprint.
  arXiv  Detail & Related papers  (2026-07-29T06:28:26Z)
• Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention [34.07319488267708]
  Recent NVFP4 pretraining methods mainly target transformer linear layers, leaving computation states, arithmetic and attention underexplored in 4-bit pipelines.<n>This critical gap blocks stable full-stack 4-bit pretraining, as the three core modules exhibit unique numerical failure patterns.<n>We propose Full-Stack FP4, the first complete NVFP4 pretraining framework resolving all three stability bottlenecks.
  arXiv  Detail & Related papers  (2026-07-05T17:31:36Z)
• ReQAT: Achieving Full-Precision Reasoning Accuracy with 4-bit Floating-Point Quantization-Aware Training [12.549469252572079]
  Large Reasoning Models (LRMs) achieve strong problem-solving through long chain-of-thought.<n>LRMs are constrained by the high cost of full-precision inference and growing KV cache footprints.<n>We propose ReQAT, a reasoning-centric FP4 training framework with three components.
  arXiv  Detail & Related papers  (2026-06-14T09:00:06Z)
• ThriftAttention: Selective Mixed Precision for Long-Context FP4 Attention [0.12691047660244334]
  We propose ThriftAttention, a low-bit attention variant that delivers near-FP16 long-context quality at FP4 inference efficiency.<n>We demonstrate across long-context benchmarks and model families that by computing only 5% of query-key blocks in FP16, ThriftAttention recovers on average 89.1% of the FP4-to-FP16 performance gap.
  arXiv  Detail & Related papers  (2026-05-21T22:28:27Z)
• Dissecting Outlier Dynamics in LLM NVFP4 Pretraining [46.10969678564592]
  This study conducts a longitudinal analysis of outlier dynamics across architecture during NVFP4 pretraining.<n>We find that, compared with Softmax Attention (SA), Linear Attention (LA) reduces per-tensor heavy tails but still exhibits persistent block-level spikes under block quantization.<n>We then develop CHON, an NVFP4 training recipe integrating with post-QK operation protection.
  arXiv  Detail & Related papers  (2026-02-02T12:50:27Z)
• Quartet II: Accurate LLM Pre-Training in NVFP4 by Improved Unbiased Gradient Estimation [40.140261007984215]
  We improve the state of the art for quantized training in NVFP4 via a novel unbiased quantization routine for micro-scaled formats.<n>We show analytically that Quartet II achieves consistently better gradient estimation across all major matrix multiplications.
  arXiv  Detail & Related papers  (2026-01-30T10:39:11Z)
• What Makes Low-Bit Quantization-Aware Training Work for Reasoning LLMs? A Systematic Study [59.44848132298657]
  Post-training quantization (PTQ) usually comes with the cost of large accuracy drops, especially for reasoning tasks under low-bit settings.<n>In this study, we present a systematic empirical study of quantization-aware training (QAT) for reasoning models.
  arXiv  Detail & Related papers  (2026-01-21T11:22:29Z)
• Bridging the Gap Between Promise and Performance for Microscaling FP4 Quantization [77.67818998672516]
  We present the first comprehensive study of MXFP4 and NVFP4 for post-training quantization.<n>We introduce Micro-Rotated-GPTQ (MR-GPTQ), a variant of the classic GPTQ quantization algorithm.<n>We show that MR-GPTQ matches or outperforms state-of-the-art accuracy.
  arXiv  Detail & Related papers  (2025-09-27T09:22:21Z)
• Quartet: Native FP4 Training Can Be Optimal for Large Language Models [27.800012997794987]
  Training large language models (LLMs) models directly in low-precision offers a way to address computational costs.<n> NVIDIA's recent Blackwell architecture facilitates very low-precision operations using FP4 variants.<n>We introduce a new approach for accurate, end-to-end FP4 training with all the major computations in low precision.
  arXiv  Detail & Related papers  (2025-05-20T17:55:50Z)
• Mitigating the Impact of Outlier Channels for Language Model Quantization with Activation Regularization [62.15918574997175]
  It is known that language models contain outlier channels whose values on average are orders of magnitude higher than other channels. We propose a strategy which regularizes a layer's inputs via quantization-aware training (QAT) and its outputs via activation kurtosis regularization. We show that regularizing both the inputs and outputs is crucial for preventing a model's "migrating" the difficulty in input quantization to the weights.
  arXiv  Detail & Related papers  (2024-04-04T17:25:30Z)
• QUIK: Towards End-to-End 4-Bit Inference on Generative Large Language Models [57.04178959678024]
  We show that the majority of inference computations for large generative models can be performed with both weights and activations being cast to 4 bits. We achieve this via a hybrid quantization strategy called QUIK, which compresses most of the weights and activations to 4-bit. We provide GPU kernels matching the QUIK format with highly-efficient layer-wise runtimes, which lead to practical end-to-end throughput improvements of up to 3.4x.
  arXiv  Detail & Related papers  (2023-10-13T17:15:05Z)
• Towards Accurate Post-Training Quantization for Vision Transformer [48.779346466374406]
  Existing post-training quantization methods still cause severe performance drops. APQ-ViT surpasses the existing post-training quantization methods by convincing margins.
  arXiv  Detail & Related papers  (2023-03-25T03:05:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.