AMLA通过加法替代乘法，大幅提升昇腾芯片上大模型推理性能。
AMLA: MUL by ADD in FlashAttention Rescaling
- 用整数加法替代浮点乘法，降低计算开销并减少中间变量
- 在昇腾910上实测达614 TFLOPS，利用率86.8%，超越竞品
- 适合追求高算力利用率的AI推理部署者，尤其昇腾生态开发者
多头隐空间注意力（MLA）显著降低大语言模型中键值缓存内存占用，但带来巨大计算开销和中间变量膨胀，对硬件高效实现构成挑战，尤其在解码阶段。本文提出针对华为昇腾NPU优化的高性能内核AMLA，基于两大创新：(1) 一种新型基于FlashAttention的算法，将输出块缩放中的浮点乘法替换为整数加法，利用FP32与INT32表示的二进制对应关系；(2) 采用分层分块预加载流水线策略，最大化浮点运算利用率：预加载流水线达到立方体边界性能，分层分块在立方体核心内重叠数据搬运与计算。实验表明，在集成于CloudMatrix384的昇腾910 NPU上，AMLA最高达614 TFLOPS，达到理论峰值的86.8%，优于开源最优方案FlashMLA（在NVIDIA H800 SXM5上最高利用率为66.7%）。该内核已集成至华为CANN，即将发布。
原文摘要 · Abstract (English)
Multi-head Latent Attention (MLA) significantly reduces KVCache memory usage in Large Language Models while introducing substantial computational overhead and intermediate variable expansion. This poses challenges for efficient hardware implementation -- especially during the decode phase. This paper introduces Ascend MLA (AMLA), a high-performance kernel specifically optimized for Huawei's Ascend NPUs. AMLA is built on two core innovations: (1) A novel FlashAttention-based algorithm that replaces floating-point multiplications with integer additions for output block rescaling, leveraging binary correspondence between FP32 and INT32 representations; (2) A Preload Pipeline strategy with hierarchical tiling that maximizes FLOPS utilization: the Preload Pipeline achieves Cube-bound performance, while hierarchical tiling overlaps data movement and computation within the Cube core. Experiments show that on Ascend 910 NPUs (integrated in CloudMatrix384), AMLA achieves up to 614 TFLOPS, reaching 86.8% of the theoretical maximum FLOPS, outperforming the state-of-the-art open-source FlashMLA implementation, whose FLOPS utilization is up to 66.7% on NVIDIA H800 SXM5. The AMLA kernel has been integrated into Huawei's CANN and will be released soon.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。