CANN ops-tensor QGMM MX文档

Kernel Qgmm Mx

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

代码位置

功能说明

MX 量化 Grouped Matmul 的 Kernel 组件，基于 Tensor API 实现，仅支持 AIC 计算。 组件负责 group list 解析、group 间m/n/k更新、A/B/Scale/Bias/C 的 GM 偏移维护，以及 tail split 调度。

实现框架参考：Kernel 公共框架

特殊约束

计算模式

仅支持 AIC，不支持 AIV 参与主计算流程。

Block 依赖

通常与BlockMmad<GroupedMatmulWithScaleMx<...>, ...>搭配使用。

调度器限制

使用BlockSchedulerGmmSwatWithTailSplit负责 grouped matmul 的 tile 分发和 tail split。

Scale 类型

ScaleA 和 ScaleB 固定按fp8_e8m0_t解释。

特殊类型别名

特殊数据结构

GMMTiling

struct GMMTiling { uint32_t groupNum; int64_t m; int64_t n; int64_t k; uint32_t baseM; uint32_t baseN; uint32_t baseK; uint32_t kAL1; uint32_t kBL1; uint32_t scaleKAL1; uint32_t scaleKBL1; uint8_t isBias; uint8_t dbL0C; int8_t groupType; uint8_t groupListType; };

参数说明：

kAL1/kBL1/scaleKAL1/scaleKBL1需满足 BlockMmad 的 L1 参数约束：kAL1与kBL1的较大值作为外层 K 窗口，较大值需为较小值的整数倍；tiling 需传入一致的scaleKAL1和scaleKBL1，且不小于该外层窗口，并为该外层窗口的整数倍。Kernel 使用scaleKAL1作为共享scaleKL1。

Params

struct Params { ProblemShape problemShape; BlockMmadParams mmadParams; BlockEpilogueParams epilogueParams; GM_ADDR groupListGmAddr; GMMTiling gmmParams; };

特殊成员方法

operator() 函数

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

功能：

• 作为 kernel 入口，内部调用Run(params)

Init 函数

__aicore__ inline void Init(const Params& params)

功能：

• 初始化 GM 基地址
• 读取 grouped matmul tiling 参数
• 构造首个 group 的problemShape_

Run 函数

__aicore__ inline void Run(const Params& params)

功能：

• 遍历所有 group
• 逐 group 更新m/n/k
• 逐 group 更新地址偏移
• 调用 scheduler 和 block 完成 tile 级计算

SetMNK 函数

__aicore__ inline void SetMNK(uint32_t groupIdx)

功能：

• 从 group list 中提取当前组切分值
• 当前 QGMM MX scalar 路径按LayoutA对应的transA编译期路径更新当前组的problemShape_；groupType在当前 kernel 中未参与该判断

UpdateBaseOffsets 函数

__aicore__ inline void UpdateBaseOffsets(uint32_t groupIdx)

功能：

• 根据上一组的m/n/k更新 A/B/ScaleA/ScaleB/Bias/C 偏移
• 兼容普通 group list 与 sparse group list

ProcessSingleGroup 函数

template <bool isLastGroupAndNeedSplit> __aicore__ inline void ProcessSingleGroup(BlockScheduler& scheduler, uint32_t groupIdx)

功能：

• 构造当前 group 的 Tensor API GM Tensor
• 获取每个 tile 的坐标和形状
• 将 slice 后的 tensor 交给BlockMmad

调用示例

组件组装

using ProblemShape = AscendC::Te::Shape<int64_t, int64_t, int64_t, int64_t>; using DispatchPolicy = Blaze::Gemm::GroupedMatmulWithScaleMx<0>; using BlockMmad = Blaze::Gemm::Block::BlockMmad< DispatchPolicy, AType, LayoutA, BType, LayoutB, CType, LayoutC, BiasType, LayoutBias>; using BlockEpilogue = Blaze::Gemm::Block::BlockEpilogueEmpty; using BlockScheduler = Blaze::Gemm::Block::BlockSchedulerGmmSwatWithTailSplit; using QgmmKernel = Blaze::Gemm::Kernel::GemmUniversal< ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler>;

参数准备

using Params = typename QgmmKernel::Params; Params params = { {m, n, k, 0}, {aGmAddr, bGmAddr, cGmAddr, biasGmAddr, scaleAGmAddr, scaleBGmAddr}, {}, groupListGmAddr, {groupNum, m, n, k, baseM, baseN, baseK, kAL1, kBL1, scaleKAL1, scaleKBL1, isBias, dbL0C, groupType, groupListType} };

Kernel 执行

QgmmKernel kernel; kernel(params);

调度流程

读取首组 tiling -> 初始化 block scheduler -> 遍历 group -> 更新当前组的 m/n/k -> 更新 A/B/Scale/Bias/C 偏移 -> scheduler 分发 tile -> block 执行单 tile 计算 -> 末组按需执行 tail split

适用场景

• MX 量化 grouped matmul
• group 间m或k动态变化
• 需要 sparse group list / tail split 的 grouped matmul 场景

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