[논문 리뷰] Attn-QAT: 4-Bit Attention With Quantization-Aware Training

Peiyuan Zhang, Matthew Noto외 5명|arXiv (Cornell University)|2026. 02. 09.

Advanced Neural Network Applications인용 수 0

한 줄 요약

이 논문은 Attn-QAT를 제시합니다. 이는 FP4 주의(attention)에 대한 4비트 양자화 인식 훈련(quantization-aware training) 방법으로, 안정적인 교육과 BF16 수준의 모델 품질을 가능하게 하며 RTX 5090에서 1.1배~1.5배의 처리량 향상을 제공합니다.

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비트 주의의 필요성 동기화.
• 기존 역전파를 가진 단순한 FP4 주의가 왜 불안정한지 조사.
• QAT 중 안정적인 FP4 주의에 대한 정밀도 조정 요구 사항 파악.
• 특수화된 순전파/역전파 정밀도 처리 및 커널과 함께 Attn-QAT를 제안.
• Attn-QAT가 확산 모델과 언어 모델의 품질을 회복하고 속도 향상을 제공하는지 입증.

제안 방법

• FP4 주의의 두 가지 주요 도전 과제: 극히 거친 값 범위와 무한정 꼬리 활성화(heavy-tailed activations)를 분석.
• 정전도(FlashAttention) 스타일의 융합 연산자에 양자화 인식 학습을 적용하기 위해 순전파에서 FP4를 시뮬레이션하고 높은 정밀도 그래디언트를 유지.
• 역전파에서 주의 점수의 재계산이 순전파와 동일한 저정밀도를 사용하도록 보장.
• 소프트맥스 그라디언트를 정확히 계산하기 위한 고정밀 보조 출력 제공.
• 학습용 트라이톤(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비트 주의를 안정화하면서 모델 품질을 보존할 수 있는가?
• RQ2 Forward 및 backward 패스 간의 정밀도 조정이 FlashAttention 스타일 백엔드와의 정합성을 위해 어떤 요구를 가지는가?
• RQ3 Attn-QAT가 확산 모델 및 언어 모델에서 BF16 품질을 회복하는가, 그리고 SageAttention3와 비교했을 때 어떤 차이가 있는가?
• RQ4 현대 GPU에서 FP4 주의로 달성 가능한 성능 향상(처리량)은 어느 정도인가?

주요 결과

• Attn-QAT는 FP4 주의로 인해 발생한 품질 저하를 회복하고 평가된 모든 지표에서 BF16 성능과 일치합니다.
• QA 결과는 Attn-QAT가 SageAttention3를 능가하며 이상값 완화 휴리스틱의 필요성을 없애줍니다.
• 훈련 안정성은 두 가지 설계 선택에 의존합니다: (1) 순전파에서 FP4를 사용하는 주의 점수 재계산과 (2) 정확한 소프트맥스 그래디언트를 위한 고정밀 출력.
• Attn-QAT는 RTX 5090에서 FP4 주의 대비 SageAttention3 대비 약 1.1배~1.5배의 속도 향상을 제공합니다.
• LLM 지속 학습에서 Attn-QAT는 Qwen3-14B에서 거의 BF16 수준의 성능을 회복하고, Llama3.1-70B의 성능도 부분적으로 회복시키며 더 긴 학습으로 추가 이득이 있을 여지가 있습니다.
• 커널 벤치마크는 Triton 학습 커널과 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.