CANN/ops-tensor矩阵乘Kernel框架

Kernel基础框架

【免费下载链接】ops-tensorops-tensor 是 CANN （Compute Architecture for Neural Networks）算子库中提供张量类计算的基础算子库，采用模块化设计，支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor

公共接口说明

概述

矩阵乘 Kernel 基础框架，提供统一的模板参数、数据结构和核心流程。不同实现（Basic、QBMM MX、StreamK）在此基础上扩展特定功能，其中 QBMM MX 支持 MxFP4/MxFP8 量化、Scale 因子处理、多 Batch 维度。

详见：README.md 查看 API 清单和实现对比。

类模板概述

模板参数

特殊模板参数（量化 Kernel）

说明：QBMM MX Kernel 支持isAtomicAdd参数，用于多核并行累加场景。

类模板概述

模板参数

类型别名

Layout 构建类型

核心数据结构

Params

struct Params { ProblemShape problemShape; // 问题 shape (m, n, k, batch) BlockMmadParams mmadParams; // mmad 参数（包含 GM 地址） BlockEpilogueParams epilogueParams; // epilogue 参数 BlockSchedulerParams schedulerParams; // scheduler 参数 };

Arguments

struct Arguments { ProblemShape problemShape; // 问题 shape (m, n, k, batch) BlockMmadArguments mmadArgs; // mmad 参数 BlockEpilogueArguments epilogueArgs; // epilogue 参数 };

核心成员变量

核心成员方法

构造函数

__aicore__ inline KernelMatmul()

功能：构造 Kernel 对象。

析构函数

__aicore__ inline ~KernelMatmul()

功能：析构 Kernel 对象。

Init函数

__aicore__ inline void Init(Params const& params)

功能：初始化 Kernel，提取问题规模和 GM 地址。 执行流程：

1. 设置问题规模problemShape_
2. 提取 BlockMmad 参数mmadParams
3. 设置 A、B、C 的 GM 地址
4. 判断 bias 地址是否为 nullptr，设置isBias_和biasGmAddr_

operator函数

__aicore__ inline void operator()(Params const& params)

功能：执行矩阵乘 Kernel 计算。

公共执行流程：

1. 初始化：调用Init(params)设置参数
2. BlockScheduler 初始化：创建调度器，获取 tile 信息
3. Layout 构建：构建 A、B、C、Bias 的 ND layout
4. GM Tensor 创建：创建 A、B、C、Bias 的 GM Tensor
5. Tile 循环处理：遍历 tile 执行矩阵乘计算
6. 清理：关闭 HF32/MM Layout Transform

公共调用示例

组件组装模板

// 定义数据类型和布局 using AType = half; using BType = half; using CType = float; using BiasType = float; using LayoutA = AscendC::Te::Layout::RowMajor; using LayoutB = AscendC::Te::Layout::ColMajor; using LayoutC = AscendC::Te::Layout::RowMajor; using LayoutBias = LayoutC; // 定义问题 shape using ProblemShape = AscendC::Te::Shape<int64_t, int64_t, int64_t, int64_t>; // 定义 BlockScheduler using BlockScheduler = Blaze::Gemm::Block::BlockSchedulerMatmulBasic<ProblemShape, FULL_LOAD_MODE>; // 定义 Kernel（根据需求选择 Basic 或 StreamK） using MatmulKernel = Blaze::Gemm::Kernel::KernelMatmulBasic<...>; // 或 using MatmulKernel = Blaze::Gemm::Kernel::KernelMatmulStreamK<...>;

参数准备模板

using Params = typename MatmulKernel::Params; Params params = { {m, n, k, batch}, // problem shape {aGM, bGM, cGM, biasGM}, // mmad args {...}, // epilogue args {mL1, nL1, kL1, baseM, baseN, baseK, ...} // scheduler params };

Kernel 执行模板

MatmulKernel mm; mm(params);

公共约束

1. 模板参数要求：

   • ProblemShape 必须为AscendC::Te::Shape<int64_t, int64_t, int64_t, int64_t>类型
   • BlockMmad 必须继承自相应的 BlockMmad 基类
   • BlockEpilogue 必须与 Kernel 类型匹配
   • BlockScheduler 必须提供 tile 切分和调度功能

2. 数据格式：

   • A、C、Bias：支持 ND 格式
   • B 矩阵：支持 ND 和 NZ 格式

3. Bias 支持：可选 bias 输入，通过biasGmAddr是否为 nullptr 判断

4. Layout 构建：使用FrameLayoutFormat根据数据类型自动适配 layout

性能优化建议（公共）

1. Tile 大小选择：

   • L1 tile：充分利用 L1 容量（通常 1MB）
   • L0 tile：匹配 L0A/L0B 容量（各 128KB）

2. Block 数量配置：根据问题规模合理设置 block 数量

3. HF32 模式：FP16 输入 + FP32 输出场景建议启用

4. 数据布局：权重矩阵（B）优先使用 NZ 格式

【免费下载链接】ops-tensorops-tensor 是 CANN （Compute Architecture for Neural Networks）算子库中提供张量类计算的基础算子库，采用模块化设计，支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor