Grouped MatMul 开发指南【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills适用场景M 轴分组矩阵乘按groupList将 M 维拆分为 E 个子问题。路径blaze_customtensor_api 手写 kernel/block融合跳转如果需求是 Grouped MatMul Mul/Add/Bias/GELU 等 Vector Epilogue转到 CV 融合 MatMul 开发指南 的 Grouped CV 小节。Grouped CV 是正式支持场景但 Epilogue 必须使用TileContext接口。§1 场景背景数学定义将 C[M,N] A[M,K] × B[K,N] 沿 M 轴按groupList拆分为 E 个子问题groupList[e] m_e, sum(groupList) M C_e[m_e, N] A_e[m_e, K] × B_e[K, N] (e 0..E-1)输入输出张量shape说明A(M, K)左输入按 prefixM 切取 A_eB(E, K, N) 或共享右输入按 groupIdx 选取切片C(M, N)输出按 prefixM 写回 C_egroupList(E,) int64_tGM 缓冲区每个元素为该 group 的 M 大小prefixM前 e 个 group 的 M 累加和用于 A/C 的 GM 偏移定位。§2 组件选择本节默认描述纯 Grouped MatMul 单算子。Grouped CV 融合仍使用GroupMatmulKernel但模板参数需传入非 void Epilogue并采用 context-based epilogue 接口。组件选择来源路径blaze_customop_kernel/include/blaze_custom/KernelKernel::GroupMatmulKernelkernel/group_matmul_kernel.hBlockMmadBlock::BlockMmad(NO_FULL_LOAD)block/matmul_block_mmad.hSchedulerBlock::GroupMatmulBlockSchedulerSplitMblock/group_matmul_block_scheduler.hPolicyMatmulMultiBlockPolicyNO_FULL_LOAD_MODEpolicy/dispatch_policy.hTiling复用对应非 groupedMatmulTilingSwatassets/op_tiling/matmul/§3 组装代码Kernel 入口函数#include kernel/group_matmul_kernel.h #include block/matmul_block_mmad.h #include block/group_matmul_block_scheduler.h #include policy/dispatch_policy.h #include tiling/blaze_matmul_tiling.h #include tiling/blaze_matmul_tiling_data.h template typename LayoutA, typename LayoutB __global__ __aicore__ __cube__ void grouped_matmul_kernel( GM_ADDR dA, GM_ADDR dB, GM_ADDR dC, GM_ADDR dGroupList, const MatmulTilingData tilingData) { if (ASCEND_IS_AIV) return; using AType bfloat16_t; using BType bfloat16_t; using CType bfloat16_t; using layoutA LayoutA; using layoutC AscendC::Te::NDExtLayoutPtn; using DispatchPolicy MatmulMultiBlockPolicyNO_FULL_LOAD_MODE; using ProblemShape MatmulShape; using BlockMmad Block::BlockMmadDispatchPolicy, AType, layoutA, BType, LayoutB, CType, layoutC; using BlockScheduler Block::GroupMatmulBlockSchedulerSplitM; using MatmulKernelImpl Kernel::GroupMatmulKernelProblemShape, BlockMmad, BlockScheduler; using Params typename MatmulKernelImpl::Params; using BlockMmadParams typename BlockMmad::Params; using L1Params typename MatmulKernelImpl::L1Params; using BlockSchedulerParams typename BlockScheduler::Params; using GroupTiling typename MatmulKernelImpl::GroupTiling; ProblemShape problemShape{ static_castint64_t(tilingData.m), static_castint64_t(tilingData.n), static_castint64_t(tilingData.k), 1L}; BlockMmadParams mmadParams{dA, dB, dC}; L1Params l1Params{static_castuint64_t(tilingData.kL1)}; BlockSchedulerParams schedulerParams{ static_castint32_t(tilingData.baseM), static_castint32_t(tilingData.baseN)}; GroupTiling groupTiling{ tilingData.groupNum, tilingData.baseM, tilingData.baseN, tilingData.baseK, tilingData.l0cDB}; Params params{problemShape, mmadParams, l1Params, schedulerParams, groupTiling, {}, dGroupList}; MatmulKernelImpl kernel; kernel(params); }TilingData → Params 映射Params 字段TilingData 来源说明problemShape{m, n, k, 1}总 M所有 group 之和mmadParams{dA, dB, dC}3 个 GM 地址l1Params{kL1}K 方向 L1 tile 尺寸schedulerParams{baseM, baseN}仅 2 个字段groupTiling{groupNum, baseM, baseN, baseK, l0cDB}5 个字段dGroupListGM 指针groupList 缓冲区地址§4 Tiling 参数Tiling 引擎直接复用对应非 groupedMatmulTilingSwat不新增 grouped tiling 算法和 TilingData。MatmulTilingSwat tilingEngine; MatmulTilingData tilingData; tilingEngine.GetTilingData(totalM, N, K, inputElemBytes, tilingData, transA, transB, isANz, isBNz, hasBias);关键输出字段字段含义Kernel 端使用usedCoreNum启动核数usedCoreNum, ...baseM/baseN/baseKL0 切分颗粒groupTilingkL1L1 K 方向窗口l1Paramsl0cDBL0C ping-pong 级数groupTilingTiling 使用总 M所有 group 之和计算切分不读取groupList内容。groupNum/groupList是 grouped kernel 参数不写入MatmulTilingData。§5 关键约束Scheduler 生命周期覆盖整个 group loopScheduler 在 group loop 外构造每个 group 只调用UpdateNextProblem刷新 shape。禁止在 loop 内重建 scheduler否则丢失跨 group 轮转状态导致逐核负载不均衡。每 group 刷新 problem shapem_e可能不同必须对每个 group 调用 scheduler 刷新(m_e, n, k)不能复用上一个 group 的派生字段。per-group GM 偏移A按 prefixM 偏移前 e 个 group 的 M 累加 × KB按groupIdx * N * K偏移C按 prefixM 偏移前 e 个 group 的 M 累加 × N尾块处理m_e 0时 device 侧直接跳过不参与 tile 调度但 offset 仍按groupList[e]更新。Epilogue 约束Epiloguevoid走 pure AIC direct writebackGrouped CV 是正式支持场景传非 void Epilogue 时必须使用 context-based view hook签名为operator()(BlockShape, TileContext, flagId)不能在 epilogue 内保存 group 状态或读取groupList。§6 常见陷阱#现象根因修复1group 输出错位用 aggregate shape 推导 offset使用 prefixM 构造 per-group view2zero group 后结果错乱跳过计算但未更新 offsetoffset 始终按groupList[e]更新3小 group 精度/越界异常scheduler 未刷新当前 group shape每个 group 调用UpdateNextProblem4many-small group 核间不均衡group loop 内重建 schedulerscheduler 在 loop 外构造每组只刷新 shape5epilogue 与 group 强绑定epilogue 保存 group 状态kernel 层切好 viewepilogue 保持通用 hook【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考