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 依赖通常与BlockMmadGroupedMatmulWithScaleMx..., ...搭配使用。调度器限制使用BlockSchedulerGmmSwatWithTailSplit负责 grouped matmul 的 tile 分发和 tail split。Scale 类型ScaleA 和 ScaleB 固定按fp8_e8m0_t解释。特殊类型别名别名含义ProblemShape整体问题规模TupleShape当前 group 的问题规模BlockShape单 tile 形状SchedulerShapescheduler 使用的问题规模BlockCoordtile 坐标特殊数据结构GMMTilingstruct 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; };参数说明参数说明groupNumgroup 数量m/n/k初始问题规模baseM/baseN/baseK基础 tile 大小kAL1/kBL1A/B 的 L1 K 轴切分scaleKAL1/scaleKBL1ScaleA/ScaleB 的 L1 K 轴切分 tiling 字段在 MX 量化中两者必须一致并作为 BlockMmad 的共享scaleKL1isBias是否启用 biasdbL0CL0C 双缓冲模式当前仅值2启用其余值均视为禁用groupTypeGMM tiling 兼容保留字段当前 kernel 未读取该字段split 方向由LayoutA对应的transA编译期路径决定!transA按 MtransA按 KgroupListTypeoffset、length 或 sparsekAL1/kBL1/scaleKAL1/scaleKBL1需满足 BlockMmad 的 L1 参数约束kAL1与kBL1的较大值作为外层 K 窗口较大值需为较小值的整数倍tiling 需传入一致的scaleKAL1和scaleKBL1且不小于该外层窗口并为该外层窗口的整数倍。Kernel 使用scaleKAL1作为共享scaleKL1。Paramsstruct 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 listProcessSingleGroup 函数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::Shapeint64_t, int64_t, int64_t, int64_t; using DispatchPolicy Blaze::Gemm::GroupedMatmulWithScaleMx0; 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 matmulgroup 间m或k动态变化需要 sparse group list / tail split 的 grouped matmul 场景【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考