AMLA：FlashAttention 重缩放中以加代乘
机器学习
2025-10-23 v1 人工智能
摘要
多头潜在注意力显著减少了大型语言模型中的 KVCache 内存使用，但同时引入了大量的计算开销和中间变量扩展。这给高效的硬件实现带来了挑战——尤其是在解码阶段。本文介绍了 Ascend MLA (AMLA)，这是一种专门针对华为 Ascend NPU 优化的高性能 kernel。AMLA 基于两项核心创新构建：(1) 一种新颖的基于 FlashAttention 的算法，利用 FP32 和 INT32 表示之间的二进制对应关系，用整数加法代替浮点乘法来进行输出块重缩放；(2) 采用分层分块的预加载流水线策略，最大化 FLOPS 利用率：预加载流水线实现了 Cube 算子受限的性能，而分层分块则重叠了 Cube 核心内的数据移动与计算。实验表明，在 Ascend 910 NPU（集成于 CloudMatrix384）上，AMLA 达到了高达 614 TFLOPS，达到理论最大 FLOPS 的 86.8%，优于当前最先进的开源 FlashMLA 实现（后者在 NVIDIA H800 SXM5 上的 FLOPS 利用率最高仅为 66.7%）。AMLA kernel 已集成至华为 CANN 中，即将发布。
引用
@article{arxiv.2509.25224,
title = {AMLA: MUL by ADD in FlashAttention Rescaling},
author = {Qichen Liao and Chengqiu Hu and Fangzheng Miao and Bao Li and Yiyang Liu and Junlong Lyu and Lirui Jiang and Jun Wang and Lingchao Zheng and Jun Li and Yuwei Fan},
journal= {arXiv preprint arXiv:2509.25224},
year = {2025}
}
备注
21 pages, 11 figures