昇腾CANN ops-blas GEMM算子深度解读：高性能矩阵乘在昇腾 NPU 上的实现原理与调优实战

前言

矩阵乘法（GEMM，General Matrix Multiply）是深度学习里最底层也最关键的计算原语。Transformer 的 Self-Attention 需要大量 GEMM，反卷积网络里的 im2col 同样最终落地为 GEMM，甚至很多看似跟矩阵无关的算子，拆解到底层都是一组矩阵乘加操作。可以说，一个算子库的 GEMM 性能好不好，直接决定了整个训练或推理栈的天花板在哪里。

昇腾CANN 生态下负责这件事的仓库叫 ops-blas。BLAS 是 Basic Linear Algebra Subprograms 的缩写，ops-blas 的定位就是线性代数基础算子库，核心内容就是高性能 GEMM。它不是凭空从零写出来的，而是站在 catlass（昇腾算子模板库）的肩膀上，用 Ascend C 这门面向昇腾达芬奇架构的算子编程语言，把 GEMM 在 Ascend 910 NPU 上的性能压到极致。

这篇文章不会给你一段粘贴复制就能跑通的脚本——那样的内容网上一搜一大把。我要做的事是真正拆开 GEMM 在昇腾 NPU 上的实现逻辑，讲清楚它是怎么利用硬件特性的，以及你在调优自己的 GEMM 算子时可以参考哪些思路。

---

GEMM 在 CANN 五层架构中的位置

在说实现之前，先把 GEMM 放在整个昇腾CANN 的体系里定位一下，这样后面讲优化策略时你才能理解"为什么要这样做"。

昇腾CANN 是一个五层的异构计算架构，从上到下是：

• 第1层：AscendCL —— 面向应用开发者的统一编程接口，封装了推理、预处理、单算子调用等能力。
• 第2层：AOL 算子库 —— 包含 NN/BLAS/DVPP/AIPP/HCCL/融合算子，ops-blas 里的 GEMM 就归属在这一层的 BLAS 类目下。
• 第3层：Graph Compiler —— 图编译器，负责算子图层面的优化和编译。
• 第4层：Runtime 运行时 —— 真正驱动硬件执行的地方，包括 HCCL 集合通信和 DVPP 预处理。
• 第5层：计算基础层 —— 驱动、内存管理等底层支撑。

GEMM 作为 BLAS 里的 Level 3 算子，落在第2层，这意味着它被设计成可以直接被上层框架（PyTorch、TensorFlow）和图引擎调用，不需要用户写一行 Ascend C 代码。但如果你的目标是自定义融合算子，或者想在 PyTorch 里用 NPU 原生性能跑自己的矩阵乘变体，理解这一层的实现原理就是绕不过去的基本功。

---

硬件基础：Ascend 910 的矩阵计算单元

GEMM 能在昇腾 NPU 上跑得飞快，根本原因在于 Ascend 910 芯片里有一个专门负责矩阵运算的硬件单元——通常叫 Cube 单元或者矩阵计算单元。这个单元跟传统的矢量计算单元是分开的，擅长做大规模矩阵乘，吞吐远高于通用矢量核。

但硬件再强，数据送不进去也是白搭。GEMM 调优的核心，说到底就是解决一个问题：怎么把数据高效地喂给 Cube 单元，让它一直有事做，不要因为等数据而空闲。

这里涉及两个关键瓶颈：

第一个瓶颈是数据搬运。 Cube 单元要算的是 A[M×K] × B[K×N]，矩阵规模一大，全部数据根本塞不进 Cube 内部的片上存储，必须分块（Tiling）送进去。每次把一块数据从外部存储（或者 L2 缓存）搬到 Cube 附近能直接访问的存储区域，这个过程叫 Tensor Model Data Transfer，在 Ascend C 里通过 DMA 引擎实现。

第二个瓶颈是搬运和计算的Overlap。 如果等搬运完了再计算，或者算完了再搬下一块，Cube 单元在两次计算之间就会出现空闲。工程上解决这个问题的经典手法叫 双缓冲（Double Buffering）：把数据分成奇数块和偶数块，计算奇数块的时候后台异步预取偶数块；奇数块算完直接切换到偶数块，此时再预取下一批奇数块。流水线一旦跑起来，Cube 单元几乎不会闲着。

ops-blas 里 GEMM 的高性能，就是建立在这两个问题的系统性解决之上的。

---

ops-blas 仓库结构与 GEMM 实现概览

ops-blas 的代码结构很干净，主目录里包含 GEMM 相关的核心实现，底层依赖 catlass 提供的模板框架。从仓库组织来看，GEMM 在 ops-blas 里并不是一个单一文件，而是一套按精度和存储布局分叉的体系，支持 FP16、FP32 等不同数据精度，以及行主序/列主序等不同内存排布。

要理解 ops-blas 的 GEMM 实现，先得知道 catlass 是什么。catlass 是昇腾的算子模板库，提供了 GEMM 算子开发所需的各种基础组件：数据分块的逻辑、内存拷贝的模板、以及跟硬件特性的适配。你不需要从零写 DMA 逻辑或者手算 tiling 尺寸，catlass 已经把最优实践封装成了模板，你只需要填进去自己的计算表达式就行。

ops-blas 就在 catlass 的基础上，构建了面向昇腾 NPU 优化过的 GEMM 算子。跟直接用 Ascend C 从零实现相比，ops-blas 的 GEMM 在 tiling 策略、内存访问模式、数据预取时机上都做了深度调优。

---

核心实现：从分块策略到双缓冲

1. Tiling 策略——把大矩阵切成 Cube 能吃下的小块

Cube 单元一次能处理的矩阵块大小是有限的，太大了放不下，太小了利用不起来。ops-blas 里的 tiling 参数不是拍脑袋定的，而是结合 Ascend 910 的片上存储容量和各维度步长精心选择的。

// 代码示意（基于 catlass 模板逻辑）
// 假设我们有 M×K 和 K×N 两个矩阵，分块思路是：
// - M 维度按 mt 块大小切分
// - N 维度按 nt 块大小切分
// - K 维度按 kt 块大小切分（kt 通常受片上可用存储限制）

const int mt = 64;   // M 方向块大小，与 Cube 单元的行处理能力匹配
const int nt = 64;   // N 方向块大小，与 Cube 单元的列处理能力匹配
const int kt = 32;   // K 方向块大小，由片上存储容量决定，不能太大

// 外层两层循环遍历 M 和 N 的块
for (int m_block = 0; m_block < M; m_block += mt) {
    for (int n_block = 0; n_block < N; n_block += nt) {
        // 内层 K 循环，每次处理一个 K 方向的块
        // 每次计算 C[m_block:mt][n_block:nt] += A[m_block:mt][k:kt] × B[k:kt][n_block:nt]
        for (int k = 0; k < K; k += kt) {
            // 这里触发 Cube 单元的矩阵乘法指令
            // Cube 一次吞吐极大，mt/nt/kt 的选择直接决定利用率
        }
    }
}

为什么要把 K 维度单独拆出来放在最内层？因为片上存储有限，A 和 B 的一个 K 向块必须同时在存储里，才能让 Cube 单元完成一次乘加。如果 K 维度不单独循环，而是跟 M、N 一起用三重循环逐元素算，Cube 根本发挥不出吞吐优势。这个分层结构是 GEMM 在各类硬件（NVIDIA GPU、昇腾 NPU、甚至 CPU 上的 BLAS 库）上通用的核心思路。

2. 双缓冲——把数据搬运和计算并行起来

// 代码示意（简化版双缓冲逻辑）
// 预取信号量，0=缓冲块A有效，1=缓冲块B有效
int cur_buf = 0;

// 第一次：同步加载第一块数据到缓冲A，然后启动计算
LoadTileToBuffer(buf_A, tile_0);  // 同步等待完成，因为还没有东西可算
ComputeWithBuffer(buf_A);         // 计算时 CPU/AI Core 指令流水线是满的

// 后续：计算当前块时异步预取下一块，循环往复
for (int i = 1; i < num_tiles; ++i) {
    int next_buf = 1 - cur_buf;   // 切换到另一块缓冲
    // 上一块计算完成后立即加载下一块，不需要额外的同步等待
    LoadTileToBuffer(next_buf, tile_i);  
    ComputeWithBuffer(next_buf);  // 同时计算
    cur_buf = next_buf;
}

// 最后一块可能没有对应的预取，直接算完即可
// 通过双缓冲，计算和搬运的耗时在稳态下完全重叠

这个逻辑看起来简单，但在实际实现里有一个关键细节：搬运和计算的切换时机必须精确对齐。如果 LoadTileToBuffer 花的时间比 ComputeWithBuffer 还长，那双缓冲就失效了，反而会拖慢整体吞吐。所以 tiling 尺寸的选择不仅要考虑片上存储能不能放得下，还要确保 DMA 带宽和 Cube 计算吞吐大致匹配。这个匹配关系跟具体芯片的内存层级和总线带宽有关，通常需要实测调优。

3. L1/L2 内存层级利用

Ascend 910 的存储层级里，L1 Cache 离 Cube 单元最近，延迟最低，但容量也最小；L2 Cache 容量大一些但延迟稍高。ops-blas 的 GEMM 实现会尽量把正在使用的那批数据块保留在 L1 里，减少对 L2 或者外部存储的访问次数。

// 代码示意（L1/L2 数据排布策略）
// A 矩阵的 M×K 块被切分后，每个 mt×kt 的子块优先驻留在 L1
// B 矩阵的 K×N 块被切分后，每个 kt×nt 的子块优先驻留在 L2
// 这样做是因为 B 通常访问模式更连续（列方向顺序访问），适合用稍大容量的 L2 吸收

// 数据排布的核心原则：
// 1. A 矩阵块按行主序存，保证每次访问一个 K 向条带时地址连续
// 2. B 矩阵块按列主序存（如果硬件支持），保证 Cube 乘法时数据复用
// 3. 避免跨行/跨列的随机访问，那样会导致 DDR 带宽利用率急剧下降

这个原则反过来也是调优 GEMM 时要时刻记住的：输入矩阵的排布方式（行主序还是列主序）会直接影响 GEMM 的性能。如果你在 PyTorch 里用 NPU 插件跑 GEMM，确保输入张量在 NPU 上的存储布局跟算子期望的一致，中间的转置操作会凭空带来额外开销。

---

PyTorch 集成：从 Python 调用原生 GEMM

ops-blas 的 GEMM 最终要能被上层框架用起来。昇腾 PyTorch NPU 插件（Ascend PyTorch Adapter）提供了矩阵乘的公共接口，用户在 Python 端写代码时感知不到底层是调用了 ops-blas 的 GEMM 算子还是别的什么，但性能差距是真实存在的。

# PyTorch NPU 调用 GEMM 的标准方式
import torch

# 确保输入在 NPU 上
a = torch.randn(1024, 512, device="npu")
b = torch.randn(512, 2048, device="npu")

# torch.mm 底层会分发给 ops-blas 的 GEMM 算子
c = torch.mm(a, b)

这段代码表面上看不出任何昇腾特有的痕迹，但当你用 torch.npu.synchronize() 确认执行完成后，内部发生的事情是：昇腾CANN 的 PyTorch 适配层收到这个 matmul 请求，根据算子类型把它路由到第2层 AOL 算子库的 BLAS 模块，那里运行的就是 ops-blas 里经过 catlass 模板调优过的 GEMM 实现。如果你的矩阵规模足够大，Tiling 和双缓冲的效果会直接体现在端到端延迟上。

# 融合算子的场景：把 GEMM 和随后的激活函数合并执行
# 这是 ops-blas 高性能策略里很重要的一环——融合算子可以省掉中间结果的写回和重新加载
class GemmReLU(torch.nn.Module):
    def forward(self, x):
        # 这里 x @ weight 的 GEMM 和随后的 relu
        # 在昇腾上可以被融合为单个算子，只产生一次数据搬运
        return torch.nn.functional.relu(torch.mm(x, self.weight))

融合算子的收益在于减少了中间结果的 HBM（High Bandwidth Memory）访存次数。每次 GEMM 的中间结果不需要写回 DDR，直接流入激活函数继续计算，这个优化在 Transformer 里多层堆叠时效果会叠加放大。一个合理的估算（仅供参考，实际性能受矩阵规模和芯片利用率影响）是，融合后的 GEMM+ReLU 组合相比分离执行在端到端延迟上可能有 10%~20% 的改善空间。

---

调优实战：决定 GEMM 性能的关键参数

理解了原理之后，具体调优时应该关注哪些维度？根据 ops-blas 和 catlass 的设计经验，以下几个参数的影响最显著。

M/N/K 维度比例

Cube 单元的利用率跟矩阵形状强相关。一个极端的例子：M=1（向量）跟 M=4096（大批量）的 GEMM，在同一个硬件上跑出来的吞吐可能差出几倍。原因是 M 和 N 太小的时候，能分出来的并行块数量有限，Cube 单元的行和列都没法充分展开。调优的第一步就是确认你的矩阵形状是否在硬件友好的范围内——通常 M 和 N 大于等于 64 时效果会比较理想。

数据排布与步长

输入张量是否经过优化排布（类似 NVIDIA 的 Tensor Core 要求的数据排布格式），对 Ascend 910 的 Cube 单元能否高效执行影响很大。如果输入数据的 stride（步长）和 tiling 参数不匹配，硬件在读取数据时会产生大量非连续访问，严重拖累有效带宽。昇腾CANN 提供的 AscendCL 接口允许用户在创建张量时指定排布格式，调优时建议优先确认排布是否对齐。

融合策略

在 ops-blas 的设计里，GEMM 很少单独裸跑。常见的融合模式包括：

• GELU 融合：矩阵乘结果直接送入 GELU 激活函数，省掉中间结果的 DDR 往返。
• Softmax 融合：Attention 里的 QK^T 矩阵乘后面紧跟 Softmax，在 GEMM 层面直接融合省掉显存的读写。
• Bias 融合：在矩阵乘输出上直接加上 Bias 向量再激活，是卷积层里极其常见的模式。

融合的收益主要来自两部分：省掉中间结果的存储搬运，以及减少 kernel launch 的开销。融合算子在 ops-blas 里以子算子的形式存在，catlass 提供了把这些子算子组合起来的模板框架，让开发者不需要从零设计融合策略。

计算精度与累加策略

// Ascend C 中 GEMM 累加精度示意
// Cube 单元内部以 FP16 计算乘，累加时用 FP32 防止溢出
// 这是 Transformer 类模型训练时的标准配置，在数值稳定性和硬件效率间取平衡

using算子计算精度 = CubeAccumulator<FP32>;  // 乘用 FP16，累加用 FP32
// 如果做推理，可以全部用 FP16，把累加也切成 FP16 以换取更高吞吐
// 切换精度只影响数据排布和寄存器分配，不影响 tiling 逻辑本身

// 以下是 catlass 模板里注册算子精度时的标准写法
// 精度选择发生在算子实例化阶段，不影响上层的 PyTorch 调用接口
gemm_operator = GeluGemm<half, half, half,  // A/B/C 的数据类型
                         64, 64, 32>;      // mt/nt/kt tiling 参数
// 参数不对外暴露，由模板根据芯片型号和矩阵规模自动选择

在 Ascend 910 上做推理场景调优时，把 FP32 GEMM 降级为 FP16 或 BF16 通常能获得接近翻倍的吞吐提升（仅供参考，实际倍数取决于矩阵规模和内存带宽）。ops-blas 的算子实例支持在创建时指定精度类型，catlass 模板负责把精度参数一路透传到 Cube 指令生成层，不需要开发者手动写底层汇编。

---

性能对比与硬件适配考量

同样的 GEMM 算子在 Ascend 910 上的表现，跟在 NVIDIA V100/A100 上相比，有共性也有差异。共性在于：分块、双缓冲、内存层级利用这些基本策略在任何加速器上都是有效的。差异在于不同硬件的存储容量、Cube 单元的并行度、以及内存总线带宽各不相同，所以最优的参数配置会随硬件型号变化。

ops-blas 本身是为 Ascend 910 设计的，它充分利用了这块芯片上 Cube 单元的吞吐优势和昇腾达芬奇架构的存储层次。如果你的目标是让 GEMM 在昇腾 NPU 上跑出接近理论峰值的性能，建议先用较小的矩阵规模验证算子正确性，再逐步增大到能充分利用硬件并行的规模，过程中监控各个阶段的耗时占比——是计算耗时还是搬运耗时占主导，决定了后续优化的方向。

---

结尾

GEMM 之所以值得单独用一篇文章来拆解，是因为它是深度学习算子栈里少有的"懂原理就能调优"的地带。硬件特性的理解（Tiling、双缓冲、Cube 单元）、软件层面的排布策略、以及融合算子的选择，这些维度叠加在一起，决定了最终的性能数字。

ops-blas 作为昇腾CANN 线性代数基础算子库，把这些复杂逻辑封装成可直接调用的 GEMM 算子，让大多数开发者不需要碰 Ascend C 也能用上经过 catlass 模板调优过的矩阵乘性能。但如果你有自定义融合算子的需求，或者想深入理解 GEMM 在昇腾 NPU 上是怎么跑起来的，ops-blas 的源码和 catlass 的模板框架就是最好的参考资料。

昇腾CANN ops-blas 仓库：https://atomgit.com/cann/ops-blas