HiFA4：在 Ascend HIF4 NPU 上实现免训练 4-bit FlashAttention 用于 LLM 推理

我们提出 HiFA4，一种训练后算子级设计，在 Ascend NPU 上进行 LLM 推理时，将 FlashAttention 中的 QK^T 和 PV 均作为 4-bit HIF4 Cube GEMM 执行，同时以 FP16 维护在线 softmax 状态。据我们所知，HiFA4 是首个在标准 NLP 基准上评估的此类面向 Ascend-HIF4 的设计。HiFA4 结合了两种机制。Smooth-QK 在 RoPE 之后对 Q 和 K 应用校准静态的逐通道等效重缩放，将量化难度从 K 转移到 Q，且在推理时无需逐分块在线归约。P-Reordering 从用于 PV GEMM 的相同量化注意力权重 P_hat 中累积 softmax 归一化因子，而非从更高精度的重构中获取。我们表明，这种不一致的表述引入了相干的输出缩放误差，并在 Qwen3-8B Layer-0 MMLU 追踪上验证了该效应，其中所有 360 万个测量的注意力分块均表现出净概率质量损失，中位数 epsilon_bar = -0.064。P-Reordering 还允许将归一化因子融合到 PV Cube GEMM 中。在五个 LLM 上，HiFA4 减少了由量化引起的决策漂移。在 Qwen3-8B 上，它恢复了由直接 HIF4 量化引入的 37.5% 的准确率差距，将样本加权准确率损失从 1.12 pp 缩小至 0.70 pp，将 BF16 不一致的 MMLU 预测从 16.3% 降至 8.2%，并将 MMLU 准确率退化减少了 57%（从 1071 降至 465）。在 Gemma2-9B 上，轻度平滑使 HiFA4 保持在 BF16 的 0.7 pp 范围内，同时将 MMLU 退化减少 27%。在禁用 Smooth-QK 的 LLaMA3.1-8B、Mistral-7B 和 Phi-4B 上，采用 Q-Mean 辅助的 P-Reordering 仍将全量 MMLU 退化减少了 41-52%。初步的指令调度分析预计，通过将 softmax 归一化因子融合到 PV Cube GEMM 中，关键路径延迟可减少 35.4%；硬件验证留待未来工作。