MatMul 算子在昇腾 NPU 上的优化实践：从原理到实战

前言

刚接触昇腾CANN那会，我被 MatMul 算子砸懵了。

不是因为它难——矩阵乘法谁没写过？问题在于，同样的矩阵乘法，在 Ascend 910 上跑出来的性能，能差出 3 倍。后来才明白，昇腾 NPU 的达芬奇架构不是让你把 CPU 上的 MatMul 代码原样搬过来就能跑快的。它的内存层级、向量计算单元、Cube 单元的配合方式，跟 GPU 完全不是一个路数。

这篇文章把我踩过的坑、测过的数据、翻过的源码捋一遍。读完你应该能回答三个问题：

1. MatMul 在 ops-nn 里到底长什么样？
2. 昇腾 NPU 上有哪些可以挖的性能点？
3. 融合算子为什么能把 MatMul + Activation 一起做了？

---

MatMul 算子的本质

矩阵乘法听起来简单：两个矩阵 A(M×K) 和 B(K×N)，得到 C(M×N)。

但"简单"是个陷阱。

在昇腾 NPU 上，MatMul 不是一条指令搞定的事。它涉及三个层面的协作：

数据搬运：A 和 B 从系统内存搬到昇腾 NPU 的片上内存（L1 Buffer），再搬到计算单元附近的 Local Buffer。搬运路径不对，Cube 单元饿死，计算单元空转。

分块计算：达芬奇架构的 Cube 单元一次能算一个 16×16×16 的块（FP16 场景下）。M、K、N 三个维度都要切成块，分块大小直接影响 Cube 利用率。

尾块处理：当 M、K、N 不是 16 的倍数时，边缘位置的那些元素要单独处理。处理不好，性能掉 30% 很正常。

ops-nn 仓库里的 MatMul 算子，核心就是把这些事做对。

---

昇腾 NPU 的硬件特性

写 MatMul 优化之前，得先搞清楚对手盘——Ascend 910 的达芬奇架构到底长什么样。

Cube 单元

Cube 单元是昇腾 NPU 的矩阵计算核心。它专门算矩阵乘法，吞吐量远高于向量单元。FP16 场景下，每个时钟周期能完成 16×16×16 次乘加运算。

但 Cube 单元有个特点：它只认分块后的数据。你不能直接扔一个任意形状的矩阵进去，必须按 16×16×16 的块组织数据。

内存层级

Ascend 910 的内存层级大概是这么个结构：

系统内存（Host DDR）
    ↓ PCIe 搬运
全局内存（Global Memory，设备侧）
    ↓ 高带宽总线
L1 Buffer（片上，大小有限）
    ↓ 快速通路
Local Buffer（每个 AI Core 独享）
    ↓
Cube 单元 / Vector 单元

问题在哪？Global Memory 到 L1 Buffer 的带宽是瓶颈。如果分块策略导致频繁搬运、重复读取，MatMul 的吞吐直接被带宽卡死。

多 AI Core 并行

Ascend 910 有几十个 AI Core。MatMul 要把输出矩阵 C 按行或者按块切分，分到不同 AI Core 上算。切分策略要考虑两点：负载均衡和数 据复用。

如果切得太细，每个 AI Core 处理的块太小，Cube 单元利用率上不去。切得太粗，部分 AI Core 闲着，浪费算力。

---

ops-nn 中的 MatMul 实现

ops-nn 是昇腾CANN开源的基础算子库，matmul 和 activation 类是它的核心内容，支持算子融合。

目录结构

ops-nn 仓库里跟 MatMul 相关的代码主要在这几个位置（基于开源仓库的公开目录结构）：

• matmul/：MatMul 算子主实现
• activation/：Activation 算子（ReLU、GELU、SiLU 等）
• fusion/：融合算子实现（MatMul + Activation 融合）

具体文件和函数签名以仓库实际代码为准，这里不做猜测。

Ascend C 编程模型

ops-nn 的算子用 Ascend C 编写。Ascend C 是昇腾CANN的算子编程语言，它提供了一套 C++ 模板库，让你可以直接操作 AI Core 的 Cube 单元、Vector 单元和内存层级。

一个典型的 Ascend C 算子包含几个部分：

• Init()：初始化，设置输入输出张量的形状、数据类型、内存布局
• Process()：主计算循环，分块、搬运、计算、写回
• 数据搬运用 DataCopy 系列接口
• Cube 计算用 MatMul 模板类
• Vector 计算用 UnaryOps / BinaryOps 系列接口

---

代码实战：基础 MatMul

下面给一个简化版的 MatMul 算子框架，展示 Ascend C 的基本写法。

// 这是一个教学框架，展示 Ascend C MatMul 算子的基本结构
// 实际 ops-nn 代码以开源仓库为准
#include "matmul_kernel.h"
#include "kernel_operator.h"

// 模板参数：输入类型、输出类型、是否启用融合
template <typename InT, typename OutT, bool FUSION_ENABLED>
class MatMulKernel {
public:
    __aicore__ inline void Init(
        GM_ADDR aGM, GM_ADDR bGM, GM_ADDR cGM,
        const MatMulParams& params)
    {
        // WHY: 先把 Global Memory 地址映射到局部指针
        // 这样后续 DataCopy 才能知道从哪搬数据
        aGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ InT*>(aGM), params.M * params.K);
        bGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ InT*>(bGM), params.K * params.N);
        cGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ OutT*>(cGM), params.M * params.N);

        // WHY: 根据 M、K、N 计算分块数
        // 每个块 16x16x16（FP16），分块数决定循环次数
        blockM = (params.M + 15) / 16;
        blockK = (params.K + 15) / 16;
        blockN = (params.N + 15) / 16;
    }

    __aicore__ inline void Process()
    {
        for (int i = 0; i < blockM; ++i) {
            for (int j = 0; j < blockN; ++j) {
                // WHY: 每次只搬一个块到 Local Buffer
                // 这样 L1 Buffer 不会被撑爆
                CopyABlock(i, j);
                ComputeBlock(i, j);
                WriteBackBlock(i, j);
            }
        }
    }

private:
    __aicore__ inline void CopyABlock(int bi, int bj)
    {
        // WHY: DataCopy 是异步的，后面要跟 SetFlag/WaitFlag
        // 否则 Cube 单元可能读到旧数据
        DataCopy(aLocal, aGlobal[/* offset */], /* len */);
        DataCopy(bLocal, bGlobal[/* offset */], /* len */);
    }

    __aicore__ inline void ComputeBlock(int bi, int bj)
    {
        // WHY: MatMul 模板类封装了 Cube 单元的调用
        // 第一个参数是输出，后面是左右操作数
        mmObject.MatMul(cLocal, aLocal, bLocal);
    }

    // 局部缓冲区，存在 AI Core 的 Local Memory 里
    LocalTensor<InT> aLocal, bLocal;
    LocalTensor<OutT> cLocal;
    GlobalTensor<InT> aGlobal, bGlobal;
    GlobalTensor<OutT> cGlobal;

    MatMul<InT, OutT> mmObject;

    int blockM, blockK, blockN;
};

这段代码有几个点值得说：

分块循环：blockM × blockN 的双重循环决定了多少个块需要计算。每个块独立计算，可以并行到不同 AI Core。

数据搬运：DataCopy 是 Ascend C 的数据搬运接口。它从 Global Memory 搬数据到 Local Buffer。这里的关键是搬运和计算重叠——用双缓冲（ping-pong buffer）可以让 Cube 单元不停工。

MatMul 模板类：mmObject.MatMul(...) 最终会编译成 Cube 单元的机器指令。你不用手写汇编，但分块大小、数据对齐方式会影响最终生成的指令序列。

---

优化点一：双缓冲与流水

上面那个基础版本有个问题：搬运一个块的时候，Cube 单元只能干等。

解决办法是双缓冲（ping-pong buffer）：准备两块 Local Buffer，一块在计算的时候，另一块在搬运下一个块的数据。这样搬运和计算可以重叠。

// 双缓冲版本的 Process 核心逻辑
__aicore__ inline void ProcessWithDoubleBuffer()
{
    // WHY: ping 和 pong 两块缓冲区交替使用
    // cur 表示当前正在计算的块，next 表示正在搬运的下一个块
    int cur = 0, next = 1;

    // 预热：先搬第一个块（没有计算可以重叠，必须单独搬）
    CopyABlockWithBuf(cur, 0, 0);

    for (int i = 0; i < blockM; ++i) {
        for (int j = 0; j < blockN; ++j) {
            // WHY: 搬运下一个块（如果还有的话）
            // 这跟当前块的计算是并行的
            if (HasNextBlock(i, j)) {
                CopyABlockWithBuf(next, NextI(i, j), NextJ(i, j));
            }

            // WHY: WaitFlag 确保当前块的数据已经搬完
            // 否则 Cube 单元可能算到一半发现数据还没到位
            WaitFlag(cur);
            ComputeBlockWithBuf(cur, i, j);

            // WHY: SetFlag 通知搬运单元可以开始搬下一个块了
            // 这个 flag 是 AI Core 内部的同步原语
            SetFlag(next);

            // 交换 ping/pong 缓冲区角色
            std::swap(cur, next);
        }
    }
}

这个优化在实测中能带来多少收益？取决于矩阵大小和分块策略。小矩阵（M、N 都小于 512）上，双缓冲的收益可能只有 10~15%；大矩阵（M、N 大于 2048）上，收益能到 30~40%。

数据来源：基于 Ascend 910 上 FP16 MatMul 的估算，实际性能跟输入形状、Batch 大小、内存对齐都有关系。

---

优化点二：尾块处理

当 M、K、N 不是 16 的倍数时，边缘块（tail block）要特殊处理。

最直接的做法是补零（padding）——把矩阵补成 16 的倍数，算完再把多余的部分裁掉。但补零有代价：搬运更多的数据、占用更多 Local Buffer、计算无效结果。

ops-nn 里的做法是动态分块：先算完整的 16×16 块，最后单独处理尾块。

// 尾块处理的核心逻辑
__aicore__ inline void ComputeTailBlock(
    int bi, int bj, int actualM, int actualN)
{
    // WHY: 尾块的 actualM 和 actualN 可能小于 16
    // 直接调 MatMul 会越界，必须用小矩阵专用路径
    if (actualM < 16 || actualN < 16) {
        // WHY: 小矩阵走 Vector 单元而不是 Cube 单元
        // Cube 单元的最小块是 16x16，小矩阵用 Cube 反而慢
        ComputeSmallMatMul(bi, bj, actualM, actualN);
    } else {
        // 正常大小的块，走标准 Cube 路径
        ComputeBlock(bi, bj);
    }
}

这里有个取舍：尾块处理让代码变复杂了，但要不要做取决于你的场景中非对齐矩阵的频率。如果 80% 的 MatMul 调用都是对齐的（比如 Transformer 里的 hidden_size 通常是 16 的倍数），尾块处理的收益有限，反而让代码难维护。

---

优化点三：MatMul + Activation 融合

这是 ops-nn 支持融合的实际价值所在。

MatMul 后面跟 Activation（ReLU、GELU、SiLU 等）是很常见的模式，比如 Transformer 的 FFN 层：Linear → GELU → Linear。

如果不融合，流程是这样的：

MatMul 算完 → 结果写回 Global Memory
    ↓
Activation 读 Global Memory → 计算 → 结果写回 Global Memory

两次 Global Memory 的读写，带宽浪费。

融合之后：

MatMul 算完 → 结果留在 Local Buffer
    ↓
Activation 直接用 Local Buffer 的数据算 → 结果写回 Global Memory

省了一次写回和一次读取。

// MatMul + GELU 融合的核心逻辑
template <typename InT, typename OutT>
class MatMulGELUFusion {
public:
    __aicore__ inline void Process()
    {
        for (int i = 0; i < blockM; ++i) {
            for (int j = 0; j < blockN; ++j) {
                // WHY: MatMul 结果不写回 Global Memory
                // 直接存在 cLocal 里，给 GELU 用
                mmObject.MatMul(cLocal, aLocal, bLocal);

                // WHY: GELU 用 Vector 单元算
                // cLocal 还在 Local Buffer 里，不需要再搬一次
                GELU(cLocal, cLocal);

                // WHY: 最后才写回 Global Memory
                // 整个融合算子只有一次写回
                WriteBackBlock(i, j);
            }
        }
    }

private:
    __aicore__ inline void GELU(LocalTensor<OutT>& dst, LocalTensor<OutT>& src)
    {
        // GELU 近似公式：x * sigmoid(1.702 * x)
        // 用 Vector 单元的 UnaryOps 和 BinaryOps 实现
        VectorMul(dst, src, sigmoidResult);
    }
};

融合算子的性能收益取决于矩阵大小。小矩阵上，Global Memory 读写的开销占比高，融合的收益更明显。大矩阵上，Cube 单元的计算时间占主导，融合的收益相对小一些。

性能数据（仅供参考）：在 Ascend 910 上，FP16 MatMul(1024×1024, 1024×1024) + GELU 融合相比分开执行，吞吐提升约 15~25%。数据来源：基于 ops-nn 仓库 README 中提及的融合能力所做的估算，具体数值取决于输入形状和运行环境。

---

多 AI Core 并行策略

MatMul 的输出矩阵 C(M×N) 可以按行切分，也可以按列切分，还可以按二维块切分。

ops-nn 用的是按行切分：把 M 个输出行分到多个 AI Core 上，每个 AI Core 算其中的一部分。

切分的时候要考虑两件事：

负载均衡：每个 AI Core 分到的行数尽量相等。如果 M=513，16 个 AI Core，不能前面 15 个分 32 行、最后一个分 33 行——这种不均衡在 AI Core 数量多的时候会被放大。

数据复用：A 矩阵的同一行可能被 C 矩阵的多行复用（因为 B 矩阵不变）。如果切分策略让同一个 A 的行被多个 AI Core 重复搬运，带宽就浪费了。按行切分的好处是 A 的搬运可以广播，但实现起来需要处理 AI Core 间的同步。

---

与 PyTorch NPU 插件的对接

在 PyTorch 里调用昇腾 NPU 上的 MatMul，走的是 PyTorch NPU 插件的公共接口。

典型的调用路径：

import torch

# 在 NPU 上创建输入张量
a = torch.randn(1024, 2048, device="npu", dtype=torch.float16)
b = torch.randn(2048, 512, device="npu", dtype=torch.float16)

# PyTorch 的 torch.matmul 会自动调度到 ops-nn 的 MatMul 算子
c = torch.matmul(a, b)

# 如果要用融合算子，需要通过 PyTorch NPU 插件提供的融合接口
# 具体接口以 PyTorch NPU 插件官方文档为准

这里的调度逻辑是：PyTorch NPU 插件把 torch.matmul 映射到昇腾CANN的 MatMul 算子实现。如果输入满足融合条件（比如后面紧跟 GELU），插件会自动选择融合算子。

不编造具体 API 名称，因为 PyTorch NPU 插件的公共接口随版本变化，具体函数名和参数以官方文档为准。

---

调试技巧

写 Ascend C 算子的时候，调试是个麻烦事——你不能在 AI Core 上直接 printf。

几个实用的调试方法：

CPU 模式仿真：Ascend C 提供了 CPU 模式，可以在 x86 上跑算子的仿真版本。虽然性能数据没有参考意义，但正确性验证可以提前做。

Dump 中间结果：在算子的关键点（搬运完、计算完）把 Local Buffer 的数据拷出来，存到文件里用 NumPy 对比。

小规模测试：先拿 16×16 的小矩阵验证正确性，再扩大到实际大小。小矩阵的问题好定位。

性能 Profile：用昇腾CANN的调优引擎 AOE 做性能分析，看 Cube 利用率、带宽利用率、同步开销各占多少。

---

常见误区

误区一：分块越大越好

不是。分块大，每个 AI Core 的局部性更好，但 Local Buffer 大小有限。分块超过 Local Buffer 容量，就得频繁置换数据，反而慢。

误区二：Cube 利用率 100% 就是最优

Cube 利用率高不代表整体性能好。如果数据搬运成为瓶颈，Cube 单元利用率再高也没用。要看的是端到端的吞吐，不是单个计算单元的指标。

误区三：融合算子一定更快

不一定。如果 Activation 的计算量很小（比如 ReLU 就是一条指令），融合的收益可能覆盖不了融合算子带来的代码复杂度和编译开销。GELU 这种计算密集的 Activation，融合的收益才明显。

---

结尾

MatMul 看起来是个"已经解决的问题"，但在昇腾 NPU 上把它跑好，涉及的层面不少：分块策略、双缓冲流水、尾块处理、算子融合、多 AI Core 并行，每个环节都有可以挖的性能点。

ops-nn 作为昇腾CANN开源的基础算子库，把这些优化都封装好了。你直接用 torch.matmul 就能享受到。但知道底层发生了什么，出了问题才知道去哪找。

如果要做定制优化——比如你的模型里 MatMul 的形状有特殊规律，或者你要融合一个不常见的 Activation——就得自己下场改 Ascend C 代码了。希望这篇文章能帮你少踩几个坑。

仓库地址：https://atomgit.com/cann/ops-nn