HiFA4：升腾 HIF4 NPU 上的 无需训练 4 位 FlashAttention，用于 LLM 推理

我们提出了 HiFA4，这是一种 后训练 操作员级设计，它在 FlashAttention 中执行 QK^T 和 PV 作为 4 位 HIF4 Cube GEMM，用于 Ascend NPU 上的 LLM 推理，同时保持 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 轨迹的影响，其中所有 3.6M 测量的注意力图块均表现出净概率质量损失，中值 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 中，关键路径延迟相对于 BF16 减少了 35.4%；硬件验证留给未来的工作。