CANN/ops-tensor MX量化Batch Matmul Kernel

Kernel Qbmm Mx

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

代码位置

功能说明

MX 量化 Batch Matmul Kernel，仅支持 AIC 计算，支持 MxFP4/MxFP8 量化格式。集成 Scale 因子（pertokenScale、scale）处理、多 Batch 维度支持、L2 Cache 动态配置，适用于量化推理场景。

继承自：Kernel Matmul 基础框架

特殊约束

量化格式支持

支持以下量化数据类型：

• MxFP4：fp4x2_e2m1_t、fp4x2_e1m2_t（4-bit 浮点）
• MxFP8：fp8_e5m2_t、fp8_e4m3fn_t（8-bit 浮点）

Scale 因子要求

必须提供两个 Scale 因子：

• pertokenScaleGmAddr：A 矩阵的 per-token scale（fp8_e8m0_t类型）
• scaleGmAddr：B 矩阵的 per-group scale（fp8_e8m0_t类型）

计算模式

仅支持 AIC 模式，不支持 AIV 计算（AIV 核直接返回）。

BlockScheduler 限制

仅支持BlockSchedulerQbmm调度器，支持多 Batch 维度切分。

L2 Cache 动态配置

根据 tile 形状动态启用/禁用 L2 Cache：

• 大 tile（curBaseM >= problemShape.m）：禁用 L2 Cache
• 小 tile：启用 L2 Cache

Batch 维度限制

支持 4 维 Batch（batchA1/A2/A3/A4、batchB1/B2/B3/B4、batchC1/C2/C3/C4），需满足广播规则。

Atomic Add 模式

可选 Atomic Add 模式（isAtomicAdd = true），用于多核并行累加场景。

特殊成员方法

构造函数

__aicore__ inline QuantBatchMmMx()

功能：构造 QuantBatchMmMx 对象。

析构函数

__aicore__ inline ~QuantBatchMmMx()

功能：析构 QuantBatchMmMx 对象。

特殊模板参数

template < class ProblemShape, // 问题形状类型 class BlockMmad, // BlockMmadMX 组件 class BlockEpilogue, // 后处理组件（通常为 void） class BlockScheduler, // BlockSchedulerQbmm 调度器 bool isAtomicAdd> // 是否启用 Atomic Add

模板参数说明

特殊类型别名

特殊数据结构

Params

struct Params { ProblemShape problemShape; // 问题 shape (m, n, k, batch) BlockMmadParams mmadParams; // mmad 参数（包含 GM 地址） L1Params l1Params; // L1 参数（kL1, scaleKL1, l1BufNum） BlockSchedulerParams schParams; // scheduler 参数 QBMMTiling qbmmParams; // QBMM 特有参数 };

QBMMTiling

struct QBMMTiling { uint32_t batchA1, batchA2, batchA3, batchA4; // A 矩阵 Batch 维度 uint32_t batchB1, batchB2, batchB3, batchB4; // B 矩阵 Batch 维度 uint32_t batchC1, batchC2, batchC3, batchC4; // C 矩阵 Batch 维度 uint32_t biasThreeDim; // Bias 是否为 3 维 uint32_t baseM, baseN, baseK; // L0 tile 形状 uint32_t isBias; // 是否启用 bias uint32_t dbL0C; // L0C 双缓冲标志 };

BlockMmadParams（MX 特有）

struct Params { GM_ADDR aGmAddr; // A 矩阵 GM 地址 GM_ADDR bGmAddr; // B 矩阵 GM 地址 GM_ADDR cGmAddr; // C 矩阵 GM 地址 GM_ADDR biasGmAddr; // Bias GM 地址（可选） GM_ADDR pertokenScaleGmAddr; // A 矩阵 Scale GM 地址 GM_ADDR scaleGmAddr; // B 矩阵 Scale GM 地址 };

特殊成员方法

Init函数

__aicore__ inline void Init(const Params& params)

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

1. AIV 检查：如果是 AIV 核则直接返回
2. 设置 Bias 标志：根据qbmmParams.isBias判断
3. 设置 BiasThreeDim 标志：根据qbmmParams.biasThreeDim判断
4. 调用ResetGmAddr设置 GM 地址

Run函数

__aicore__ inline void Run(const Params& params)

功能：执行量化 Batch Matmul Kernel 计算。 执行流程：

1. Atomic Add 配置：如果isAtomicAdd = true，调用SetAtomicAdd<float>
2. 调用Init(params)设置参数
3. 创建 BlockScheduler 实例
4. 初始化 BlockMmadMX 组件
5. 判断 Batch 数量：
   • 单 Batch（b == 1）：调用ProcessSingleBatch
   • 多 Batch：调用ProcessWithBatch
6. 清理 Atomic Add：如果启用，调用SetAtomicNone

ProcessSingleBatch函数

__aicore__ inline void ProcessSingleBatch( const Params& params, BlockScheduler& bs, uint64_t batchCnt, bool isTailRound)

功能：处理单个 Batch 的矩阵乘计算。 执行流程：

1. 构建 Layout：A、B、ScaleA、ScaleB、Bias、C
2. 创建 GM Tensor
3. 动态配置 L2 Cache
4. Tile 循环处理：
   • 获取 tile 坐标和形状
   • Slice GM Tensor 到当前 tile
   • 调用 BlockMmadMX 执行量化矩阵乘

ProcessWithBatch函数

__aicore__ inline void ProcessWithBatch(const Params& params, BlockScheduler& bs)

功能：处理多 Batch 的矩阵乘计算。 执行流程：

1. 计算 Batch 广播倍数：multiA1C1、multiB1C1 等
2. 4 维 Batch 循环（batchC1/C2/C3/C4）：
   • 更新 Batch 偏移：batchCOffset_、batchAOffset_、batchBOffset_
   • 调用AddBatchOffset更新 GM 地址
   • 调用ProcessSingleBatch处理当前 Batch

SetL2Cache函数

template <typename TensorB, typename TensorScaleB, typename TensorC> __aicore__ inline void SetL2Cache( const ProblemShape& problemShape, uint64_t curBaseM, uint64_t baseN, TensorB& gmB, TensorScaleB& gmScaleB, TensorC& gmC)

功能：动态配置 L2 Cache。 说明：

• 大 tile 场景：禁用 B 和 ScaleB 的 L2 Cache
• Atomic Add 模式：禁用 C 的 L2 Cache

调用示例

组件组装

// 定义量化数据类型 using AType = fp4x2_e2m1_t; // 或 fp8_e5m2_t using BType = fp4x2_e2m1_t; // 或 fp8_e5m2_t using CType = float; using BiasType = float; using ScaleType = fp8_e8m0_t; // 定义 Layout using LayoutA = AscendC::Te::NDExtLayoutPtn; using LayoutB = AscendC::Te::NZLayoutPtn; using LayoutC = AscendC::Te::NDExtLayoutPtn; using LayoutBias = AscendC::Te::NDExtLayoutPtn; // 定义调度策略 using DispatchPolicy = Blaze::Gemm::MatmulWithScaleMx<A_FULL_LOAD_MODE>; // 定义 BlockMmadMX using BlockMmad = Blaze::Gemm::Block::BlockMmad< DispatchPolicy, AType, LayoutA, BType, LayoutB, CType, LayoutC, BiasType, LayoutBias>; // 定义 BlockScheduler using BlockScheduler = Blaze::Gemm::Block::BlockSchedulerQbmm<ProblemShape>; // 定义 Kernel using QBMMKernel = Blaze::Gemm::Kernel::QuantBatchMmMx< ProblemShape, BlockMmad, void, BlockScheduler, false>;

参数准备

using Params = typename QBMMKernel::Params; Params params = { {m, n, k, batch}, // problem shape {aGM, bGM, cGM, biasGM, scaleAGM, scaleBGM}, // mmad params {kL1, scaleKL1, l1BufNum}, // L1 params {baseM, baseN, baseK, mTailTile, nTailTile, ...}, // scheduler params {batchA1, batchA2, ... batchC4, biasThreeDim, baseM, baseN, baseK, isBias, dbL0C} // QBMM tiling };

Kernel 执行

QBMMKernel qbmm; qbmm(params); // 或 qbmm.Run(params);

数据流

存储层次

GM (A/B/ScaleA/ScaleB/Bias) → L1 (量化数据 + Scale) → L0A/L0B → L0C → GM (C)

Batch 处理流程

Batch 循环（4 维） ↓ 更新 Batch 偏移（batchA/B/COffset） ↓ 重置 GM 地址（AddBatchOffset） ↓ 处理单个 Batch（ProcessSingleBatch） ↓ Tile 循环 → BlockMmadMX 执行

量化计算流程

加载量化数据 (A/B) + Scale (ScaleA/ScaleB) ↓ Dequantize（自动在 Mmad 中完成） ↓ Mmad 计算（MX 模式：MmadTraitMX） ↓ 输出 float 结果

性能优化建议

L1 缓冲配置

• l1BufNum：建议 2 或 4，平衡 L1 容量和流水线并行度
• kL1：建议对齐到MXFP_DIVISOR_SIZE（64）

Scale KL1 配置

• scaleKL1：建议为kL1的整数倍（如 2×kL1）
• Scale 数据复用：Scale 常驻 L1，减少搬运开销

全载模式选择

• 非全载模式：每次迭代重新加载 A/B 块
• A 全载模式：A 矩阵常驻 L1，适用于大 K、小 M 场景

Batch 维度设计

• Batch 维度需满足广播规则：batchA = batchC × multiA、batchB = batchC × multiB
• 4 维 Batch 灵活支持多种广播场景

L2 Cache 配置

• 大 tile 场景自动禁用 L2 Cache，避免缓存污染
• Atomic Add 模式自动禁用 C 的 L2 Cache

适用场景

• 量化推理：MxFP4/MxFP8 量化矩阵乘
• Batch Matmul：多 Batch 维度支持
• Scale 因子处理：per-token 和 per-group scale
• 广播 Batch：A/B/C 不同 Batch 维度的广播计算

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