当前位置：首页 > news >正文

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

news 2026/5/22 6:38:33

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

前言

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

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

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

MatMul 在 ops-nn 里到底长什么样？
昇腾 NPU 上有哪些可以挖的性能点？
融合算子为什么能把 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 插件的公共接口随版本变化，具体函数名和参数以官方文档为准。