HiFA4: Training-Free 4-bit FlashAttention on Ascend HIF4 NPUs for LLM Inference

2026-07-05

AI Summary

This work addresses the accuracy degradation and decision drift in 4-bit quantized large language model (LLM) inference on the Ascend HIF4 NPU, proposing HiFA4—a training-free, operator-level optimization framework. HiFA4 introduces the first 4-bit FlashAttention implementation on the HIF4 architecture by shifting quantization difficulty from the QK^T computation to post-RoPE scaling via Smooth-QK, and incorporates a P-Reordering mechanism that fuses the softmax normalization factor into the 4-bit Cube GEMM for PV computation. While retaining softmax in FP16, this approach injects a controlled scaling error to enhance output consistency. Experiments show that HiFA4 recovers 37.5% of the accuracy loss on Qwen3-8B, reduces MMLU prediction inconsistency from 16.3% to 8.2%, and decreases regressive samples by 57%. Across multiple models, MMLU regression is reduced by 27%–52%, with an estimated 35.4% reduction in critical-path latency.

Abstract

We present HiFA4, a post-training operator-level design that executes both QK^T and PV in FlashAttention as 4-bit HIF4 Cube GEMMs for LLM inference on Ascend NPUs, while maintaining the online softmax state in FP16. To our knowledge, HiFA4 is the first Ascend-HIF4-targeted design of this kind evaluated on standard NLP benchmarks. HiFA4 combines two mechanisms. Smooth-QK applies a calibration-static per-channel equivalent rescaling to Q and K after RoPE, transferring quantization difficulty from K to Q without per-tile online reduction at inference. P-Reordering accumulates the softmax normalizer from the same quantized attention weights P_hat used in the PV GEMM, rather than from a higher-precision reconstruction. We show that this inconsistent formulation introduces a coherent output-scaling error, and validate the effect on a Qwen3-8B Layer-0 MMLU trace, where all 3.6M measured attention tiles exhibit net probability-mass loss with median epsilon_bar = -0.064. P-Reordering also allows the normalizer to be fused into the PV Cube GEMM. Across five LLMs, HiFA4 reduces quantization-induced decision drift. On Qwen3-8B, it recovers 37.5% of the accuracy gap introduced by direct HIF4 quantization, narrows the sample-weighted accuracy loss from 1.12 pp to 0.70 pp, reduces BF16-inconsistent MMLU predictions from 16.3% to 8.2%, and cuts MMLU accuracy regressions by 57% (1071 to 465). On Gemma2-9B, mild smoothing keeps HiFA4 within 0.7 pp of BF16 while reducing MMLU regressions by 27%. On LLaMA3.1-8B, Mistral-7B, and Phi-4B, where Smooth-QK is disabled, P-Reordering with the adopted Q-Mean auxiliary still reduces full-set MMLU regressions by 41-52%. A preliminary instruction-scheduling analysis projects a 35.4% critical-path latency reduction relative to BF16 by fusing the softmax normalizer into the PV Cube GEMM; on-hardware validation is left to future work.

Problem

• LLM inference
• 4-bit quantization
• FlashAttention
• Ascend NPU
• quantization-induced decision drift

Innovation

• 4-bit quantization
• FlashAttention
• Ascend NPU
• post-training quantization
• softmax normalization