4-bit量化加速大模型推理，保持高精度且减少误差
HiFA4: Training-Free 4-bit FlashAttention on Ascend HIF4 NPUs for LLM Inference
- 通过动态重排与通道校准，实现4比特量化下的高效注意力计算
- 在Qwen3-8B上恢复37.5%的精度损失，准确率下降从1.12pp降至0.70pp
- 适合追求低精度推理性能优化的芯片部署开发者
我们提出HiFA4，一种无需训练的算子级设计，可在昇腾HIF4 NPU上以4比特HIF4立方体GEMM执行FlashAttention中的QK^T和PV计算，同时保持softmax状态为FP16。据我们所知，这是首个针对昇腾HIF4架构的标准NLP基准评估。该设计结合两项机制：Smooth-QK在RoPE后对Q、K进行静态逐通道等效缩放，将量化难度从K转移至Q，避免推理时每块的在线归约；P-Reordering利用与PV GEMM中相同的量化注意力权重P_hat累积softmax归一化因子，而非高精度重构。我们验证其不一致形式引入系统性输出缩放误差，在Qwen3-8B Layer-0 MMLU trace中，全部360万次注意力块均出现概率质量净损失，中位ε̄ = -0.064。该方法还使归一化因子可融合进PV立方体GEMM。在五款大模型上，HiFA4显著降低量化导致的决策漂移。在Qwen3-8B上，恢复37.5%精度损失，样本加权准确率损失从1.12个百分点降至0.70，BF16不一致预测从16.3%降至8.2%，MMLU准确率退化减少57%（1071降至465）。在Gemma2-9B上，轻微平滑使精度保持在0.7个百分点内，同时减少27%的退化。在LLaMA3.1-8B、Mistral-7B和Phi-4B上，即使关闭Smooth-QK，P-Reordering配合Q-Mean辅助仍使整体MMLU退化减少41%-52%。初步指令调度分析显示，融合归一化因子可使关键路径延迟降低35.4%，硬件验证留待后续。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。