在这里插入图片描述

乐高积木与算子工厂的奇妙相遇

想象一下,你是一个乐高积木工厂的设计师。你的任务不是为每个孩子造一个成品城堡,而是设计一套"基础砖块+连接件+装饰件"的组合系统——孩子们可以用这套系统拼出他们想要的任何东西。catlass 就是昇腾矩阵乘算子的"乐高工厂"。

这不是一个普通的算子库,而是一个模板工厂。它的核心哲学不是"给你一个成品算子",而是"给你一套可复用、可替换、可局部修改的模板组件"。就像乐高没有"标准城堡积木",只有"2×4红色砖、1×2灰色砖、斜坡件、窗户件"——组合的自由在孩子手里。

catlass 的全称是 CATLASS(Compute Architecture for Tensor Linear Algebra Subroutines on Ascend),聚焦高性能矩阵乘类算子的基础模板。它由华南理工陆璐教授团队与华为 CANN 团队联合开发,从 v1.0 到最新的 v1.5.0,已经支持 Ascend 910 和 Ascend 950 两代芯片。这个仓库的存在,让算子开发从"从零手写"进化到"模板组装"——性能不降,开发效率翻倍。

三层抽象架构:从基础砖块到成品算子

catlass 的核心设计理念是"分层抽象"。就像乐高分"基础砖块层→功能组件层→成品模型层"三层,catlass 也把自己的模板分成了三层:

第一层:基础层——“基础砖块"与"连接件”

基础层提供的是最原始的积木块,包括:

  • GMpr(Global Memory Provider):全局内存管理器,就像乐高的"底板"——所有积木都要搭在上面。它负责与昇腾 NPU 的全局内存(HBM)交互,决定数据放哪、怎么搬、何时释放。

  • VReg(Vector Register):向量寄存器抽象,就像乐高的"1×1小砖"——最小的工作单元。昇腾 NPU 的 Vector 单元靠它来存储和计算中间结果。

  • TPipe(Tensor Pipe):数据流水线抽象,就像乐高的"传送带"——把数据从全局内存搬运到 SRAM(片上高速缓存),再从 SRAM 搬到 Vector 寄存器,算完后再写回去。TPipe 负责编排这条流水线,让搬运和计算尽可能重叠。

这一层不关心你算的是矩阵乘还是向量加,只关心"数据怎么搬、怎么存"。它是所有上层模板的基础,就像所有乐高模型都离不开 2×4 红色砖。

第二层:特化层——“功能组件”

特化层在基础层之上,针对具体计算场景做特化。这一层的关键是"策略":

  • 矩阵乘策略(MatMul Policy):决定矩阵乘怎么算。分块大小多大?循环怎么展开?要不要做多级缓存?这些策略决定了矩阵乘的性能上限。

  • 数据类型策略(Datatype Policy):决定用什么精度算。FP32?FP16?BF16?还是 HiFloat8?不同的数据类型对应不同的硬件执行路径。

  • 硬件特化策略(Hardware Policy):针对不同芯片做特化。Ascend 910 的 Cube 单元和 Ascend 950 的 Cube 单元在延迟、吞吐、缓存大小上都有差异,这一层负责"对硬件谈判"——把计算逻辑映射到最合适的硬件执行路径。

这一层就像乐高的"轮子组件"“窗户组件”——不再是原始砖块,而是有特定功能、可以复用的模块。一个"轮子组件"可以用在赛车上,也可以用在卡车上;同样,一个"矩阵乘策略"可以用在 FlashAttention 算子里,也可以用在 MoE 算子里。

第三层:实例化层——“成品模型”

实例化层是最终交付给开发者的东西——具体的矩阵乘算子。它不是手写的,而是通过组合第二层的策略、复用第一层的基础能力,"组装"出来的。

比如,一个"FP16 矩阵乘算子"的实例化过程是这样的:

FP16MatMul 算子 = 
  GMpr(内存管理)
  + TPipe(数据搬运流水线)
  + MatMulPolicy(矩阵乘分块策略)
  + FP16DatatypePolicy(FP16 数据类型策略)
  + Ascend910HardwarePolicy(Ascend 910 硬件特化)

这就像孩子用"基础砖块+轮子组件+窗户组件"拼出一个成品赛车。他不需要知道窗户是怎么注塑出来的,只需要知道"窗户组件往这里一插就行"。同样,开发者不需要知道内存管理的细节,只需要选择合适的策略组合。

Ascend 910 vs Ascend 950:硬件特化的"两套积木"

catlass 的一个关键能力是"硬件特化"——同一套模板代码,通过不同的策略配置,适配不同的芯片。v1.5.0 新增对 Ascend 950 的支持,让这套积木工厂能产出"两套不同的成品模型"。

Ascend 910:训练芯片的"重型积木"

Ascend 910 是昇腾的训练芯片,Cube 单元(矩阵计算单元)的吞吐大、延迟低,但 SRAM 容量相对小(每核 1MB 级别)。这意味着矩阵乘策略要"大开大合"——分块要大,循环展开要狠,尽可能让 Cube 单元满载跑起来。

在 catlass 的模板里,Ascend 910 的硬件特化策略是这样的:

// Ascend 910 矩阵乘策略示例(概念代码)
template<>
struct MatMulPolicy<Ascend910> {
    static constexpr int kM_Block = 128;  // M 维分块大小
    static constexpr int kN_Block = 128;  // N 维分块大小
    static constexpr int kK_Block = 64;   // K 维分块大小
    static constexpr bool kUseDoubleBuffer = true;  // 开启双缓冲
};

这几行代码的含义是:对于 Ascend 910,矩阵乘采用 128×128×64 的分块策略,同时开启双缓冲(计算和搬运重叠)。这是针对训练场景的"重型配置"——吞吐优先。

Ascend 950:推理芯片的"精密积木"

Ascend 950 是昇腾的推理芯片,Cube 单元的吞吐略低,但 SRAM 容量大(每核 2MB 级别),延迟控制更精细。这意味着矩阵乘策略要"小步快跑"——分块可以小一点,流水线层级多一点,让数据在 SRAM 里待的时间更长,减少对 HBM 的访问。

在 catlass 的模板里,Ascend 950 的硬件特化策略是另一套:

// Ascend 950 矩阵乘策略示例(概念代码)
template<>
struct MatMulPolicy<Ascend950> {
    static constexpr int kM_Block = 64;   // M 维分块更小
    static constexpr int kN_Block = 64;   // N 维分块更小
    static constexpr int kK_Block = 32;   // K 维分块更小
    static constexpr bool kUseTripleBuffer = true;  // 三缓冲,更好利用大 SRAM
};

这里的变化:分块从 128×128×64 降到 64×64×32,同时从双缓冲升级到三缓冲。这是针对推理场景的"精密配置"——延迟优先,充分利用大 SRAM 减少内存访问。

同一套模板,两套策略,一份代码

关键在于:开发者在写矩阵乘算子时,不需要写两份代码。他只需要这样实例化:

// Ascend 910 版本的矩阵乘算子
using MatMul910 = GemmKernel<
    GMpr,                                    // 内存管理
    TPipe<Ascend910>,                        // Ascend 910 的流水线
    MatMulPolicy<Ascend910>,                 // Ascend 910 的分块策略
    FP16DatatypePolicy,                      // FP16 数据类型
    Ascend910HardwarePolicy                  // Ascend 910 硬件特化
>;

// Ascend 950 版本的矩阵乘算子
using MatMul950 = GemmKernel<
    GMpr,                                    // 内存管理(复用)
    TPipe<Ascend950>,                        // Ascend 950 的流水线
    MatMulPolicy<Ascend950>,                 // Ascend 950 的分块策略
    FP16DatatypePolicy,                      // FP16 数据类型(复用)
    Ascend950HardwarePolicy                  // Ascend 950 硬件特化
>;

同一个 GemmKernel 模板,通过传入不同的策略参数,生成了两个完全不同的算子实现。这就是 catlass 的核心能力——“白盒化组装”:模板是透明的,策略是可替换的,性能是可预期的。

一个完整实例:从模板到 FP16 矩阵乘算子

理论讲完了,来个实战。下面是一个完整的 FP16 矩阵乘算子模板实例化代码(简化版,保留核心逻辑):

// 文件:examples/gemm_fp16_910.cpp
#include "catlass/gemm_kernel.h"          // 1️⃣ 引入矩阵乘核心模板

// 2️⃣ 定义 Ascend 910 硬件特化策略
using Hardware910 = catlass::HardwarePolicy<
    catlass::Ascend910,                   // 芯片型号
    catlass::CubeUnit,                    // 使用 Cube 矩阵计算单元
    catlass::HBM                          // 全局内存类型
>;

// 3️⃣ 定义 FP16 数据类型策略
using FP16Policy = catlass::DatatypePolicy<
    catlass::FP16,                        // 输入数据类型
    catlass::FP16,                        // 输出数据类型
    catlass::FP32                         // 累加精度(避免精度损失)
>;

// 4️⃣ 定义矩阵乘分块策略(Ascend 910 优化版)
using BlockPolicy = catlass::MatMulPolicy<
    128,                                  // M 维分块大小
    128,                                  // N 维分块大小
    64,                                   // K 维分块大小
    true                                  // 开启双缓冲
>;

// 5️⃣ 实例化最终算子
using GemmFP16_910 = catlass::GemmKernel<
    catlass::GMpr,                        // 全局内存管理
    catlass::TPipe<Hardware910>,          // 数据流水线
    BlockPolicy,                          // 分块策略
    FP16Policy,                           // 数据类型策略
    Hardware910                           // 硬件特化
>;

// 6️⃣ 主函数:调用算子
extern "C" __global__ void gemm_fp16_kernel(
    half* A, half* B, half* C,           // 输入输出矩阵指针
    int M, int N, int K                   // 矩阵维度
) {
    GemmFP16_910::execute(A, B, C, M, N, K);
}

逐行解析

第 1 行:引入 catlass 的矩阵乘核心模板 GemmKernel。这个模板是所有矩阵乘算子的"骨架",定义了算子的整体执行流程。

第 2 行Hardware910 定义):硬件特化策略。这里指定三个关键信息:

  • Ascend910:目标芯片是 Ascend 910 训练芯片
  • CubeUnit:使用 Cube 单元做矩阵乘(不是 Vector 单元,Cube 单元是昇腾 NPU 的矩阵计算专用硬件)
  • HBM:全局内存是 HBM 类型(高带宽内存)

这个策略会让 GemmKernel 在编译时选择针对 Ascend 910 Cube 单元的指令序列。

第 3 行FP16Policy 定义):数据类型策略。关键点在第三行:FP32 累加精度。这意味着虽然输入输出是 FP16,但中间的累加(矩阵乘的 dot product 累加)用 FP32,避免 FP16 精度损失导致数值误差累积。这是高性能矩阵乘的标准做法。

第 4 行BlockPolicy 定义):分块策略。这是性能的关键:

  • 128×128×64 的分块意味着:每次从 HBM 加载 128×64 的 A 矩阵块和 64×128 的 B 矩阵块到 SRAM,算出一个 128×128 的 C 矩阵块
  • true 表示开启双缓冲:在计算当前块时,同时加载下一块的数据到另一个缓冲区,让搬运和计算重叠

第 5 行GemmFP16_910 定义):这是整个模板实例化的核心。GemmKernel 是 catlass 的核心模板,它接收 5 个策略参数:

  1. GMpr:全局内存管理(负责与 HBM 交互)
  2. TPipe<Hardware910>:数据流水线(负责 HBM→SRAM→寄存器的搬运编排)
  3. BlockPolicy:分块策略(决定矩阵乘怎么切)
  4. FP16Policy:数据类型策略(决定用什么精度算)
  5. Hardware910:硬件特化策略(决定用什么硬件指令)

编译时,C++ 模板会把这 5 个策略"融合"成一个完整的算子实现,包含具体的指令序列、内存访问模式、循环展开策略。

第 6 行gemm_fp16_kernel 函数):这是 Ascend C 的算子入口函数。extern "C" 保证 C 链接,__global__ 表示这是 NPU 核函数(类似 CUDA 的 __global__)。函数内部只需一行:调用实例化后的 GemmFP16_910::execute,传入矩阵指针和维度。

这段代码的魔力在哪?

如果你用传统方式手写这个矩阵乘算子,需要做这些事情:

  1. 手写 HBM→SRAM 的数据搬运代码(处理对齐、bank conflict)
  2. 手写 Cube 单元的矩阵乘指令(查阅昇腾指令集手册)
  3. 手写双缓冲的流水线编排(计算何时加载下一块)
  4. 手写 FP16→FP32 的类型转换和累加逻辑
  5. 针对 Ascend 910 和 Ascend 950 分别写两套

用 catlass 模板,你只需要 6 行策略配置代码。性能?不低于手写。为什么?因为 catlass 的策略就是从手写优化代码里抽象出来的"最佳实践"——开发者已经踩过坑了,你直接用就行。

与 ops-blas 的关系:模板工厂 vs 成品货架

catlass 和 ops-blas 经常被混淆。简单说:catlass 是模板工厂,ops-blas 是成品货架

ops-blas:开箱即用的算子封装

ops-blas 是 CANN AOL 算子库中的线性代数算子库,提供直接可用的算子封装:

// ops-blas 的调用方式(示意)
#include "ops_blas.h"
ops_blas::Gemm(A, B, C, M, N, K);  // 直接调用,不用关心实现细节

ops-blas 的定位是"让普通开发者不碰算子细节"。它封装了矩阵乘、向量点乘、矩阵分解等常用线性代数操作,提供类似 BLAS 库的标准 API。开发者只需要调用,不需要知道内部用的是 catlass 模板还是其他实现。

catlass:给算子开发者的"乐高工坊"

catlass 的定位完全不同。它不是给普通开发者用的,而是给算子开发者用的。如果你:

  • 需要写一个自定义矩阵乘算子(比如特殊的分块策略、特殊的融合逻辑)
  • 需要针对新芯片适配矩阵乘(比如下一代 Ascend 芯片)
  • 需要在矩阵乘里嵌入特殊计算(比如量化、稀疏、混合精度)

这时候,你不需要从零手写,而是用 catlass 的模板组装。

两者关系:ops-blas 可能调用 catlass

在 CANN 的生态里,ops-blas 底层可能使用 catlass 模板实现矩阵乘。这意味着:

普通开发者 → 调用 ops-blas → ops-blas 内部调用 catlass 实例化的算子
算子开发者 → 直接用 catlass → 自己组装算子

catlass 是"底层引擎",ops-blas 是"上层封装"。普通开发者用 ops-blas 就够了;需要定制化时,才下到 catlass 层。

旅程终点站:一个留给你思考的问题

catlass 的三层抽象设计,本质上是把"算子开发"从"手工作坊"变成了"流水线工厂"。基础层提供砖块,特化层提供功能组件,实例化层组装成品——开发者的角色从"工匠"变成了"设计师"。

但这里有一个取舍:模板化意味着"有限自由度"。你能组装的算子,受限于 catlass 提供的策略组合。如果你想要的矩阵乘策略不在模板库里(比如极端的分块策略、特殊的融合逻辑),你只能:

  1. 扩展 catlass 的策略(写新的 Policy 类)
  2. 退回到手写 Ascend C 代码
Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐