arXiv:2603.00040cs.LGcs.AI2026-03被引 3

首次系统解决4比特注意力训练不稳问题,实现高效低精度推理。

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

  • 通过匹配反向传播的低精度重计算,稳定4比特注意力训练
  • 在扩散与语言模型上恢复FP4注意力性能,无需额外异常值处理
  • 支持新显卡的4比特端到端推理,速度提升1.5倍以上

实现可靠的4比特注意力是实现在新兴FP4支持显卡上端到端FP4计算的前提,但受限于FP4极小动态范围及注意力层重尾激活分布,注意力仍是主要障碍。本文首次系统研究了注意力的4比特量化感知训练(QAT)。发现‘即插即用’的QAT——将FP4前向传播与高精度Flash Attention(FA)风格反向传播直接结合——会导致训练不稳定。我们识别出两个关键原则以实现稳定FP4注意力:(1)反向传播中注意分值的低精度重计算需与前向一致;(2)解决FA梯度计算中隐含的精度假设。基于此,提出Attn-QAT,并实现用于训练的融合Triton内核与用于FP4推理的内核。在扩散模型和语言模型上,Attn-QAT在无需先前方法中使用的显式异常值缓解启发式的情况下,恢复了FP4注意力的性能损失,并在RTX 5090上相比SageAttention3实现高达1.5倍加速,在GB300上相比FA4实现高达1.74倍加速。视频演示见:https://drive.google.com/drive/folders/190F6xbBDUF2kGQYIcXBt3ehSYij5jlim?usp=sharing。代码见:https://haoailab.com/FastVideo/training/attn_qat/。

原文摘要 · Abstract (English)

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

量化训练注意力机制4比特计算GPU优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。