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

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.

• The paper presents HiFA4, a training-free quantization scheme enabling QK^T and PV multiplications as 4-bit Cube GEMMs on Ascend NPUs.
• It employs Smooth-QK and P-Reordering to calibrate K activation outliers and restore decision-level consistency, recovering up to 37.5% accuracy loss.
• Experimental results on various LLMs show substantial latency reduction and reduced accuracy loss, all without the need for retraining.

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

Overview and Motivation

The paper "HiFA4: Training-Free 4-bit FlashAttention on Ascend HIF4 NPUs for LLM Inference" (2607.04302) introduces HiFA4, a post-training quantization scheme for attention computation in LLMs on Ascend NPUs using the HIF4 block-floating-point format. HiFA4 enables both $QK^T$ and $PV$ matrix multiplications in FlashAttention to be executed as 4-bit Cube GEMMs, while maintaining the online softmax state in FP16. The main goal is to recover decision-level consistency and accuracy lost due to aggressive quantization, without requiring any retraining or fine-tuning. The technical design addresses accuracy degradation specific to activation quantization in attention, distinguishing itself from prior work primarily focused on weight quantization or NVIDIA-specific formats.

Origin and Structure of K-Activation Outliers

The root cause of quantization-induced degradation is traced to sparse, parameter-driven activation outliers in the K tensor, primarily induced by QK-RMSNorm affine weights. In Qwen3-8B, Layer~0, channel~51 exhibits an extreme outlier due to a large fixed parameter $\gamma_{K,51}$, amplifying the otherwise normal projection output and saturating the HIF4 quantizer's group scale, which disrupts downstream attention computation.

Figure 1: Computation chain leading to the K-activation outlier at Layer~0, channel~51 of Qwen3-8B; the outlier arises entirely within QK-RMSNorm via a small RMS normalizer and large affine weight $\gamma_{K,51}$.

The outlier is stable across tokens and heads, demonstrating its parameter-level, broadcast nature.

Figure 2: QK-RMSNorm affine weights across all layers and channels; Layer~0 K weight (approx 34) is an isolated outlier among K weights (mean = 2.87, Q layer mean = 1.79).

Smooth-QK: Calibration-Static Outlier Suppression

Smooth-QK applies a per-channel scaling to Q and K after RoPE, determined from calibration statistics (RMS aggregate over 1024 training samples), transferring quantization difficulty from K to Q. The transformation is controlled by a gating mechanism:

• Strong smoothing ($\alpha=0.5$): When intra-K outlier concentration ($\rho_K$) and K-to-Q dynamic range ($T_{K/Q}$) exceed empirical thresholds.
• Mild smoothing ($\alpha=0.25$): When the outlier is less pronounced.

Figure 3: The Smooth-QK applicability gate as a decision plane over $\rho_K$ and $T_{K/Q}$; smoothing is triggered only for concentrated and dominant outliers.

The method is deployment-friendly, as all scale factors are determined offline and applied as simple per-channel multiplies/divides, with no per-tile online reduction. Channel-51 outliers are effectively suppressed ($PV$0 for Qwen3-8B, compressing max K from 220.8 to $PV$159).

Figure 4: K activations before Smooth-QK for a typical slice; channel-51 stripe is outside the quantizer's dynamic range.

P-Reordering: Algebraic Consistency and Latency

P-Reordering addresses the mismatched normalization in quantized attention: previous approaches accumulate the softmax normalizer from a high-precision reconstruction $PV$2, while the $PV$3 GEMM consumes the quantized $PV$4. Theorem~\ref{thm:main} proves this inconsistency introduces a coherent output-scaling error, with median $PV$5 (radial contraction) on empirical traces (100% of Layer~0 tiles in Qwen3-8B).

Figure 5: Theorem characterizing systematic normalization-induced error; P-Reordering eliminates this via consistent quantized paths.

Figure 6: Projected critical-path latency reduction; P-Reordering enables fused softmax normalizer in Cube GEMM, reducing vector-path bottleneck (estimated 35.4% reduction relative to BF16).

Experimental Results and Numerical Evidence

Evaluation on Qwen3-8B, Gemma2-9B, LLaMA3.1-8B, Mistral-7B, and Phi-4B demonstrates HiFA4's efficacy:

• On Qwen3-8B: HiFA4 narrows sample-weighted accuracy loss ($PV$6) from 1.12 pp (direct HIF4) to 0.70 pp (recovering 37.5%); MMLU Quant Hurts cut from 1071 to 465; flip rate and Hurts halved across benchmarks.
• On Gemma2-9B: HiFA4 stays within 0.7 pp of BF16; Quant Hurts reduced by 27% (MMLU).
• On bypass models (no concentrated K-outlier): HiFA4's P-Reordering + Q-Mean still reduces Quant Hurts by 41–52% (LLaMA3.1-8B MMLU: 976 → 468).

Top-1/top-2 logit margins, used as a robustness proxy, show HiFA4 recovers much of the decision margin lost under direct HIF4 quantization.

Figure 7: Top-1/top-2 logit margin on correct predictions; HiFA4 recovers robustness lost to quantization.

Long-context retrieval remains saturated (100%), but attention output relative error shows Direct HIF4 accumulates error with context length (51.5% → 63.1% from 4k to 40k), whereas HiFA4 holds steady ($PV$712%).

Ablation and Generalization

Component ablations confirm that Smooth-QK and P-Reordering both contribute, with combined effects outperforming either alone. On models where Smooth-QK is not triggered (diffuse outlier), P-Reordering and Q-Mean alone still recover substantial accuracy and decision consistency.

Practical and Theoretical Implications

• Practical: HiFA4 enables high-throughput LLM inference on Ascend NPUs with HIF4, with significant latency reduction and without retraining.
• Theoretical: The parameter-level analysis of activation outliers exposes root causes of quantization failure and highlights the necessity of algebraically consistent normalization.
• Future Directions: Empirical gate thresholds need validation across broader model suites; on-hardware latency measurements are pending NPU availability; deeper study of outlier prevalence and cross-layer/coherent error accumulation is warranted.

Conclusion

HiFA4 is a robust, post-training quantization scheme for FlashAttention in LLMs, proven on five models and grounded in parameter-level analysis and algebraic consistency. The scheme supports 4-bit Cube GEMMs for both $PV$8 and $PV$9 on Ascend NPUs, recovers accuracy losses and decision consistency, and projects substantial latency reductions. The core mechanisms are calibration-static correction of K-outlier activations (Smooth-QK) and normalization consistency (P-Reordering), generally applicable without retraining and demonstrably effective even on models without concentrated outlier structures. The approach sets a formal and practical baseline for quantized attention on hierarchical block-floating-point hardware.