HiFA4: Ascend HIF4 NPU에서 LLM 추론을 위한 학습이 필요 없는 4-bit FlashAttention
요약
Ascend NPU에서 LLM 추론을 최적화하기 위해 학습 없이 4-bit FlashAttention을 구현하는 HiFA4 기술을 제안합니다. Smooth-QK와 P-Reordering 메커니즘을 통해 양자화로 인한 정확도 손실을 최소화하고 연산 효율을 높입니다.
핵심 포인트
- Ascend HIF4 NPU 타겟의 최초 4-bit FlashAttention 설계
- Smooth-QK를 통한 양자화 난이도 조절 및 정확도 유지
- P-Reordering으로 소프트맥스 정규화 인자 누적 및 GEMM 융합
- Qwen3-8B 등 주요 모델에서 MMLU 정확도 퇴보 대폭 감소
- 임계 경로 지연 시간(latency) 약 35.4% 감소 기대
우리는 Ascend NPU에서 LLM 추론을 위해 FlashAttention의 QK^T와 PV를 모두 4-bit HIF4 Cube GEMM으로 실행하면서, FP16에서 온라인 소프트맥스(softmax) 상태를 유지하는 사후 학습(post-training) 연산자 수준 설계인 HiFA4를 제시합니다. 우리가 알기로, HiFA4는 표준 NLP 벤치마크에서 평가된 이와 같은 종류의 최초 Ascend-HIF4 타겟 설계입니다. HiFA4는 두 가지 메커니즘을 결합합니다. Smooth-QK는 RoPE 이후 Q와 K에 대해 보정-정적(calibration-static) 채널별 등가 재스케일링(rescaling)을 적용하여, 추론 시 타일별 온라인 리덕션(online reduction) 없이 양자화 난이도를 K에서 Q로 전달합니다. P-Reordering은 고정밀 재구성(reconstruction) 대신 PV GEMM에서 사용되는 것과 동일한 양자화된 어텐션 가중치 $\hat{P}$로부터 소프트맥스 정규화 인자(softmax normalizer)를 누적합니다. 우리는 이러한 불일치하는 공식이 일관된 출력 스케일링 오차(output-scaling error)를 유발한다는 것을 보여주며, Qwen3-8B Layer-0 MMLU 트레이스에 대한 효과를 검증합니다. 여기서 측정된 360만 개의 어텐션 타일 모두가 중앙값 $\epsilon_{bar} = -0.064$와 함께 순 확률 질량 손실(net probability-mass loss)을 나타냅니다. P-Reordering은 또한 정규화 인자가 PV Cube GEMM으로 융합(fuse)될 수 있도록 합니다. 5개의 LLM에 걸쳐 HiFA4는 양자화로 인한 결정 드리프트(decision drift)를 줄입니다. Qwen3-8B의 경우, 직접적인 HIF4 양자화로 인해 발생한 정확도 격차의 37.5%를 회복하고, 샘플 가중 정확도 손실을 1.12 pp에서 0.70 pp로 좁히며, BF16과 불일치하는 MMLU 예측을 16.3%에서 8.2%로 줄이고, MMLU 정확도 퇴보를 57%(1071에서 465로) 절감합니다. Gemma2-9B에서는 가벼운 스무딩(smoothing)을 통해 HiFA4를 BF16의 0.7 pp 이내로 유지하면서 MMLU 퇴보를 27% 줄입니다. Smooth-QK가 비활성화된 LLaMA3.1-8B, Mistral-7B, Phi-4B의 경우에도, 채택된 Q-Mean 보조 도구를 사용한 P-Reordering은 전체 세트 MMLU 퇴보를 여전히 41-52% 줄여줍니다. 예비 명령어 스케줄링(instruction-scheduling) 분석에 따르면, 소프트맥스 정규화 인자를 PV Cube GEMM으로 융합함으로써 BF16 대비 임계 경로 지연 시간(critical-path latency)을 35.4% 감소시킬 것으로 예상됩니다. 하드웨어 상의 검증은 향후 과제로 남겨둡니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.AR의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기