AMLA：以加代乘，解锁昇腾 MLA 算子高性能新范式

在大模型推理与深度学习训练场景中，Multi-Head Attention（MHA）作为核心组件，其计算效率与显存占用直接决定了模型部署的性能上限。传统 Attention 计算中，Flash Attention 通过解耦计算流程降低了峰值显存占用，但在昇腾 AI 处理器等异构计算平台上，仍面临变量尺寸过大、数据搬运频繁、乘除运算开销较高等挑战。华为计算数学算法专家廖崎臣团队提出的 AMLA（Add-based Multi-Head Latent Attention）算子，创新性地采用 "以加代乘" 核心思想，从算法设计到工程实现全方位优化，为昇腾平台带来了性能突破。

一、技术背景：Flash Attention 的局限与昇腾平台的优化诉求

Flash Attention 的核心贡献是通过分块计算（tiling）将大规模矩阵运算拆解为适合 SRAM 缓存的小块，避免了 HBM 与 SRAM 之间的海量数据搬运，从而降低峰值显存占用。但其计算流程中仍存在两个关键问题，在昇腾平台上尤为突出：

1. VEC2-Rescaling 阶段的变量尺寸难题

在 Flash Attention 的归一化过程中，需要维护中间变量 O（输出累加值）和ℓ（行和累加值），并通过乘除运算进行尺度调整。当模型维度 d=512、头数 g=128 时，单个 tile 的 O 变量尺寸可达 256KB，若考虑多 tile 并行与历史数据缓存，变量总尺寸会远超昇腾处理器的 UB（Unified Buffer，192KB×2）和 L1 Cache 容量，导致变量无法常驻缓存，频繁触发 HBM 数据交互，严重影响性能。

2. 乘除运算与数据搬运的性能开销

Flash Attention 的 rescaling 过程涉及大量浮点数乘除运算，这类运算在向量核（Vector Core）上的执行效率远低于加法运算；同时，传统实现中数据需要在 GM（Global Memory）、L1 Cache、UB 等多级存储之间多次搬运，仅数据搬运操作就占据了相当比例的计算耗时，进一步限制了算子吞吐量。

昇腾 910 芯片的异构架构（包含 Cube Core、Vector Core、多级缓存与存储单元）为高性能计算提供了硬件基础，但也对算子的存储调度、计算映射提出了更高要求。如何充分利用硬件特性，解决变量存储与运算效率问题，成为突破 Attention 性能瓶颈的关键。

二、AMLA 核心创新：以加代乘的算法设计

AMLA 算子的核心突破在于用加法运算替代传统浮点数乘除，同时优化数据搬运流程，从算法根源上解决性能痛点。其核心设计思路可概括为以下两点：

1. 基于 2 的整数次幂放缩：加法替代乘法的数学基础

浮点数与 2 的整数次幂之间的乘法 / 除法，可通过二进制移位操作实现 ------ 而移位操作本质上可转化为指数位的加法运算。AMLA 利用这一特性，将 Flash Attention 中复杂的 rescaling 乘除运算，转化为基于 2 的整数次幂的放缩操作，具体原理如下：

引理支撑：给定实数 x，令 n=round (-m₂)（m₂为中间变量的指数偏移），可证明√2/2 ≤ 2^exp (x) ≤ √2，即基于 2 的整数次幂放缩能达到与传统乘除相同的防溢出效果，且误差控制在可接受范围内。
运算转化：通过将浮点数指针重定义为整型，对指数位执行加法操作（左移 / 右移），实现 2 的整数次幂乘法，完全规避了耗时的浮点数乘除运算，运算效率提升显著。

2. 逆向更新与原子累加：减少数据搬运

传统 Attention 计算中，数据需要在 Vector Core 与 Cube Core 之间往返搬运（数据搬运 ×4），而 AMLA 通过逆向更新策略重构计算流程：

简化数据流向：将中间结果 P_iV_i 直接通过原子累加操作更新至输出 O，无需额外的临时存储与数据拷贝，数据搬运次数减少 50%（降至 ×2）。
存储友好设计：优化后的变量尺寸更适配昇腾处理器的 UB 与 L1 Cache 容量，实现核心计算数据的常驻缓存，彻底解决了 VEC2-Rescaling 阶段的存储瓶颈。

三、工程实现：分层调度与硬件协同优化

算法的高性能落地离不开与硬件架构的深度协同，AMLA 在工程实现上采用了多项关键技术，充分挖掘昇腾平台的计算潜力：

1. Preload Pipeline：流水线预加载优化

针对 Cube Core 与 Vector Core 的计算依赖关系，AMLA 设计了多阶段流水线预加载机制：

构建依赖链：Cube Core 负责矩阵乘（A、B、C），Vector Core 负责向量运算（a、b、c），形成 "A→a→B→b→C→c" 的依赖链。
重叠执行：通过（Aa）→（AaBb）→（AaBcCc）的流水线调度，将数据预加载、计算、结果写回三个阶段重叠执行，隐藏了数据搬运 latency，提升了硬件利用率。

2. 分层 Tiling：多级存储适配

结合昇腾处理器的 GM→L1 Cache→UB→Buffer 多级存储架构，AMLA 采用分层分块策略：

宏观分块（baseBlock）：将大规模矩阵按（baseM、baseN、baseK）划分为适配 L1 Cache 的块，减少 GM 与 L1 之间的数据交互。
微观分块（singleBlock）：将 baseBlock 进一步划分为（singleM、singleN、singleK）的小块，适配 UB 与 LOA/LOB Buffer 容量，最大化 Cube Core 的计算并行度。

3. 硬件资源高效映射

充分利用昇腾芯片的硬件特性：Cube Core 专注于高性能矩阵乘（MMAD）运算，Vector Core 负责向量累加与指数运算，MTE（Memory Transfer Engine）负责多通道数据搬运，实现计算、存储、搬运的并行协同，避免单一硬件资源成为性能瓶颈。

四、性能验证：昇腾平台的显著优势

为验证 AMLA 算子的性能，团队在昇腾 910 芯片与主流 GPU 平台上进行了对比测试，测试场景覆盖不同序列长度（S_k=1024~16384），核心性能指标如下：

1. latency 优化显著

在昇腾 910 平台上，AMLA 算子的计算延迟较传统实现大幅降低：

短序列（S_k=1024）：延迟降至 95μs，性能提升 40.9%；
中长序列（S_k=4096）：延迟 241μs，性能提升 64.1%；
超长序列（S_k=16384）：延迟 830μs，性能提升 74.5%，且序列越长，性能优势越明显。

2. 硬件利用率领先

AMLA 算子在昇腾 910 上的硬件利用率最高达 86.8%，远超同类实现：

Cube Core 利用率提升 30% 以上，充分发挥矩阵乘并行计算能力；
数据搬运效率提升 50%，多级缓存命中率保持在 95% 以上，大幅减少 HBM 访问次数。

3. 跨场景泛化性强

AMLA 算子已适配 LLM 推理、多模态模型训练等多种场景，支持不同头数、维度的模型配置，且与 CANN 平台的 Transformer 仓深度融合，可直接集成到主流大模型训练 / 推理框架中，迁移成本低。

五、结语与开源邀请

AMLA 算子通过 "以加代乘" 的算法创新与深度硬件协同优化，成功解决了传统 Attention 算子在昇腾平台上的存储与性能瓶颈，为大模型高效部署提供了新的技术路径。该成果已集成至华为 CANN 异构计算架构，成为昇腾 AI 处理器的核心优化算子之一。

CANN 作为华为面向 AI 场景的开源异构计算平台，不仅提供了丰富的算子库（包括 Transformer、NN、CV、Math 等系列），还构建了完善的开发者生态，涵盖调试工具、优化样例、社区支持等全方位资源。目前，CANN 开源社区已拥有 3.8K Star、1.74K Fork，累计下载量超 37 万次，成为 AI 异构计算领域的重要开源力量。

欢迎大家关注 CANN 开源主仓（AtomGit）：https://atomgit.com/cann ，参与算子开发、贡献优化方案，与全球开发者一同构建高效、开放、共赢的 AI 计算生态！