CANN ops-transformer GroupedMatmulAdd 算子全解析分组矩阵乘与原地累加的 K 轴分组实现与 aclnn 调用指南【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读本文围绕 CANN ops-transformer 开源仓库中gmm/grouped_matmul_add模块的 GroupedMatmulAdd 算子展开系统讲解其在 NPU 上实现“分组矩阵乘 原地累加”的计算语义、K 轴分组原理、完整参数约束以及通过 AscendCL 两段式接口aclnnGroupedMatmulAdd/aclnnGroupedMatmulAddV2进行编程调用的完整实战流程。读完本文你将掌握该算子的数学定义、groupList 两种编码格式的构造方法、各数据类型的合法组合、典型错误码排查思路并能直接基于仓库中的可运行示例完成 Host 侧代码编写与编译运行。一、算子定位与核心功能1.1 一句话理解 GroupedMatmulAddGroupedMatmulAdd 是 CANN ops-transformer 中用于 MoEMixture of Experts类大模型场景的分组矩阵乘算子。它的核心特征是将多个矩阵乘合并为一次算子调用每组矩阵乘的 K 轴维度可以不同并且结果直接原地累加到输入矩阵 y 上即“add”语义来自原地累加而非独立的偏置相加。以数学形式表达对第 i 组i 1 … gg 为分组个数$$ yRef_i x_i \times weight_i y_i $$其中yRef_i[m_i, n_i] x_i[m_i, k_i] × weight_i[k_i, n_i] y_i[m_i, n_i]m_i / k_i / n_i为对应分组的 shape。可见该算子本质是y x weight的分组版本属于带累加的分组矩阵乘Grouped MatMul with in-place accumulation。1.2 分组方式当前仅支持 K 轴分组算子当前只支持K 轴分组K-axis grouping各分组的k_i各不相同但所有分组的m_i / n_i相同。也就是说输入 x、权重 weight 只在 K第一维方向被切分成 g 段每段独立做矩阵乘后再分别叠加到 y 的对应位置。在仓库源码 op_api/aclnn_grouped_matmul_add.cpp 的参数校验CheckAttrs中明确写死了约束CHECK_COND(params.groupType 2, ACLNN_ERR_PARAM_INVALID, ... [%s] must be 2 ...); CHECK_COND(params.transposeX, ACLNN_ERR_PARAM_INVALID, ... [%s] must be True ...); CHECK_COND(!params.transposeWeight, ACLNN_ERR_PARAM_INVALID, ... [%s] must be False ...);即groupType仅支持 2K 轴分组且当前版本要求x 必须转置、weight 不能转置。1.3 产品支持情况根据 README.md 与接口文档该算子在以下产品上支持/不支持产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从算子定义文件 op_host/grouped_matmul_add_def.cpp 可以看到算子为 AiCore 核注册了ascend910b、ascend910_93、ascend950、ascend350四类 AICore 配置其中后两者启用了动态编译、动态格式、动态 Rank 与动态 Shape 支持与文档列出的 950 系列与 A2/A3 系列产品支持情况相互印证。另外在 Ascend 950 上yRef的 shape 不支持 2 维必须为 3 维。二、参数说明与数据类型约束2.1 完整参数表以下是 README.md 中列出的完整参数说明参数名输入/输出/属性描述数据类型数据格式x输入公式中的输入 xFLOAT16、BFLOAT16NDweight输入公式中的 weightFLOAT16、BFLOAT16NDgroupList输入表示输入 K 轴方向的 matmul 大小分布的 cumsum 结果累积和INT64NDy输入表示原地累加的输出矩阵FLOAT32NDtransposeX属性表示 x 矩阵是否转置BOOL-transposeWeight属性表示 weight 矩阵是否转置BOOL-groupType属性表示分组类型仅支持 2K 轴分组INT64-groupListType属性表示分组 groupList 格式V2 接口引入INT64-yRef输出表示原地累加的输出矩阵FLOAT32ND2.2 数据类型组合支持的输入类型组合严格限定为两种由 op_host/grouped_matmul_add_def.cpp 的 OpDef 注册与 API 层CheckDtype双重约束x 为 FLOAT16、weight 为 FLOAT16、y/yRef 为 FLOAT32x 为 BFLOAT16、weight 为 BFLOAT16、y/yRef 为 FLOAT32。x 与 weight 的数据类型必须一致yRef 固定为 FLOAT32groupList 固定为 INT64。在 op_api/aclnn_grouped_matmul_add.cpp 的CheckDtype中可以看到对应的运行时校验逻辑yRef 非 FLOAT32、groupList 非 INT64、x/weight 非 FP16/BF16 或二者类型不一致时都会返回ACLNN_ERR_PARAM_INVALID。2.3 转置语义与 shape 约定由于transposeX固定为 True、transposeWeight固定为 False实际计算的是$$yRef x^T \times weight y$$即 x 的 shape 为[K, M]转置后按[M, K]参与矩阵乘weight 的 shape 为[K, N]。API 层的 shape 校验aclnn_grouped_matmul_add.cpp进一步明确x、weight 均为 2 维xDimNum 2、weightDimNum 2groupList 为 1 维x 与 weight 的 K 维第一维必须相等aKDim bKDimyRef 的元素总数必须等于M × N × groupNum代码按 3 维 shape 计算期望大小并做整体尺寸校验。2.4 约束说明使用该算子前需重点核对以下约束来自 README.md 与 aclnnGroupedMatmulAdd.mdint32 边界x 和 weight 中每一组 tensor 的每一维大小在 32 字节对齐后都应小于 int32 的最大值 2147483647。最后一维上限x 和 weight 中每一组 tensor 的最后一维大小都应小于 65536。x 的最后一维指当 x 不转置时的 K 轴或 x 转置时的 M 轴weight 的最后一维指 weight 不转置时的 N 轴或转置时的 K 轴。确定性aclnnGroupedMatmulAdd与 V2 接口默认均为确定性实现deterministic。分组上限groupList 第 1 维最大支持 1024即最多支持 1024 个 group。三、groupList 分组描述从 cumsum 到每组大小groupList 是理解该算子的关键输入。它是一段 Device 侧的 INT64 aclTensor描述 K 轴方向上各分组的切分位置。V1 接口aclnnGroupedMatmulAdd只支持 cumsum累积和模式而V2 接口aclnnGroupedMatmulAddV2通过新增的groupListType属性同时支持两种格式groupListTypegroupList 语义约束0各分组 K 轴大小的cumsum累积和即单调非递减序列最后一个值不大于 x 第一维K 轴总长1各分组 K 轴的每组大小V2 新增数值非负总和不超过 x 第一维文档中的例子非常直观。假设 K 256、分组数 E 4各组大小依次为64、0、128、64groupListType 0cumsum[64, 64, 192, 256]第 2 组大小为 0因此累积和与第 1 组相同groupListType 1每组大小[64, 0, 128, 64]。从 op_api/aclnn_grouped_matmul_add.cpp 的CheckAttrs可以看到groupListType的合法取值为 0 或 1V1 接口在构造参数时未显式传入该值实际默认按 cumsum 模式0处理。在仓库的 ST 测试基准实现 tests/st/aclnnGroupedMatmulAdd/executor_aclnnGroupedMatmulAdd.py 中可以清晰看到 cumsum 模式的计算方式——按 groupList 中的累积和切分 K 轴逐组做矩阵乘再求和last 0 for i in groupList.tolist(): x_tensor x[last:i, :].to(torch.float64) weight_tensor weight[last:i, :].to(torch.float64) result.append(torch.matmul(x_tensor.t(), weight_tensor)) last i result torch.stack(result).reshape(y.shape) y.cpu()这正好对应公式yRef_i x_i^T × weight_i y_i按[last:i]切分 K 轴x 转置后与 weight 相乘最后堆叠 reshape 并与 y 相加。四、算子工程结构从 API 到 Kernel 的完整链路gmm/grouped_matmul_add是一个结构完整的自定义算子工程源码组织如下op_api/aclnn 两段式接口实现与l0op::GroupedMatmulAdd图级接口封装aclnn_grouped_matmul_add.cpp、grouped_matmul_add.cppop_host/算子注册grouped_matmul_add_def.cpp、Infershape 推导grouped_matmul_add_infershape.cpp以及 tiling 逻辑op_tiling/grouped_matmul_add_tiling.hop_kernel/NPU 侧 AiCore kernel 实现op_kernel/grouped_matmul_add.cpp及 arch35 架构的 kernel 与 tiling 数据结构examples/可直接参考的可运行 Host 侧调用示例tests/包含 utHost 侧 tiling / infershape / op_api 单测、kernel 单测与 st基于 ATK 框架的 aclnn 接口测试含 JSON 用例与 Python executor。4.1 Infershape输出与输入 y 同形在 op_host/grouped_matmul_add_infershape.cpp 中InferShapeGroupedMatmulAdd直接令输出yRefShape yShape原地累加输出即输入 y 本身InferDataTypeGroupedMatmulAdd校验 x 与 weight 数据类型一致并将输出数据类型固定为DT_FLOAT。这与“yRef 是 y 的原地别名”的语义完全一致——在示例代码中体现为yRef y;这一行赋值。4.2 Tiling 数据分组信息如何在 Host 侧编排Tiling 数据结构定义在 op_host/op_tiling/grouped_matmul_add_tiling.hBEGIN_TILING_DATA_DEF(GmmBaseParams) TILING_DATA_FIELD_DEF(int64_t, groupNum); // 分组个数 TILING_DATA_FIELD_DEF(int64_t, coreNum); // 使用的核数 TILING_DATA_FIELD_DEF(int64_t, groupType); // 分组类型仅支持2 --- K分组 TILING_DATA_FIELD_DEF(int64_t, groupListType); END_TILING_DATA_DEF其中GroupedMatmulAddTilingData由GmmBaseParams分组基础信息与TCubeTiling矩阵乘 cube 单元 tiling 信息组合而成并配套定义了MAX_TENSOR_LIST_SIZE 128的常量。可以推断Host 侧 tiling 负责把 groupList 解析为分组个数与各组的矩阵乘切分方案交给 Kernel 侧的 cube 流水执行。五、两段式 aclnn 接口与调用实战5.1 两段式接口原型与 CANN 其他算子一致GroupedMatmulAdd 采用两段式接口two-phase API第一段GetWorkspaceSize完成入参校验、构图并计算 workspace 大小第二段执行实际计算。函数原型如下见 docs/aclnnGroupedMatmulAdd.md。V1 第一段aclnnStatus aclnnGroupedMatmulAddGetWorkspaceSize( const aclTensor *x, const aclTensor *weight, const aclTensor *groupList, aclTensor *yRef, bool transposeX, bool transposeWeight, int64_t groupType, uint64_t *workspaceSize, aclOpExecutor **executor)V2 第一段多出groupListType参数aclnnStatus aclnnGroupedMatmulAddV2GetWorkspaceSize( const aclTensor *x, const aclTensor *weight, const aclTensor *groupList, aclTensor *yRef, bool transposeX, bool transposeWeight, int64_t groupType, int64_t groupListType, uint64_t *workspaceSize, aclOpExecutor **executor)第二段两版本一致aclnnStatus aclnnGroupedMatmulAdd(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)从实现源码 op_api/aclnn_grouped_matmul_add.cpp 可以看到两段接口内部的实际流程创建 OpExecutor → 参数校验CheckParams→ 空 tensor 检查任一分组轴为 0 时直接返回空输出→ 数据连续性转换Contiguous→ 调用l0op::GroupedMatmulAdd构图 → 对非连续输出做ViewCopy→ 返回 workspace 大小。5.2 错误码与排查第一段接口完成入参校验文档列出的典型错误如下返回值错误码触发场景ACLNN_ERR_PARAM_NULLPTR161001x、weight、groupList、yRef 为空指针ACLNN_ERR_PARAM_INVALID161002数据类型/格式不在支持范围或 x 与 weight 数据类型不一致ACLNN_ERR_INNER_TILING_ERROR561002x、weight、yRef 的 shape 不满足矩阵乘限制要求与源码对应CheckNotNull对应 161001CheckFormat/CheckShape/CheckDtype/CheckAttrs统一返回 161002shape 不满足矩阵乘限制时返回 561002。5.3 完整调用示例可编译运行仓库提供了可直接编译运行的示例 examples/test_aclnn_grouped_matmul_add.cpp下面是梳理后的完整调用流程V1FP16 版本V2 示例见 examples/test_aclnn_grouped_matmul_add_v2.cpp。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_grouped_matmul_add.h // ... CHECK_RET / LOG_PRINT / GetShapeSize 宏与工具函数参考仓库示例 int main() { // 1. 固定写法device/stream 初始化 int32_t deviceId 0; aclrtStream stream; Init(deviceId, stream); // 2. 构造输入与输出K轴分组K512MN256g2 std::vectorint64_t xShape {512, 256}; // x: [K, M]因 transposeXtrue std::vectorint64_t weightShape {512, 256}; // weight: [K, N] std::vectorint64_t yShape {2, 256, 256}; // y/yRef: [g, M, N] std::vectorint64_t groupListShape {2}; std::vectorint64_t groupListData {256, 512}; // cumsum: 第1组K256第2组K256 void *xDeviceAddr, *weightDeviceAddr, *yDeviceAddr, *groupListDeviceAddr; aclTensor *x nullptr, *weight nullptr, *groupedList nullptr; aclTensor *y nullptr, *yRef nullptr; bool transpose_x true; // 仅支持 True bool transpose_weight false; // 仅支持 False int group_type 2; // 仅支持 2K轴分组 // 创建 aclTensorx/weight 用 FP16uint16_t 承载y 用 FP32groupList 用 INT64 CreateAclTensoruint16_t(xShape, xDeviceAddr, aclDataType::ACL_FLOAT16, x); CreateAclTensoruint16_t(weightShape, weightDeviceAddr, aclDataType::ACL_FLOAT16, weight); CreateAclTensorfloat(yShape, yDeviceAddr, aclDataType::ACL_FLOAT, y); CreateAclTensor_Newint64_t(groupListData, groupListShape, groupListDeviceAddr, aclDataType::ACL_INT64, groupedList); yRef y; // 原地累加输出与输入 y 是同一块内存 // 3. 第一段接口计算 workspace 并获取 executor uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; aclnnGroupedMatmulAddGetWorkspaceSize(x, weight, groupedList, yRef, transpose_x, transpose_weight, group_type, workspaceSize, executor); // 按需申请 workspace void *workspaceAddr nullptr; if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 4. 第二段接口执行计算 aclnnGroupedMatmulAdd(workspaceAddr, workspaceSize, executor, stream); // 5. 同步等待执行结束拷贝结果回 Host 并打印 aclrtSynchronizeStream(stream); auto size GetShapeSize(yShape); std::vectorfloat resultData(size, 0); aclrtMemcpy(resultData.data(), size * sizeof(float), yDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); // 6. 释放 tensor、workspace、stream 与 device 资源 aclDestroyTensor(x); aclDestroyTensor(weight); aclDestroyTensor(y); aclrtFree(xDeviceAddr); aclrtFree(weightDeviceAddr); aclrtFree(yDeviceAddr); if (workspaceSize 0) aclrtFree(workspaceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }关键点说明x 的 shape 是[K, M]因为transposeX truex 实际以[K, M]存储、按[M, K]参与运算因此示例中xShape {512, 256}表示 K512、M256。groupList 传的是累积和{256, 512}表示两组 K 分别为 256 和 256合计 512与 x 的 K 维一致。yRef y输出与输入共享同一块 Device 内存实现真正的原地累加这也是 Infershape 中输出 shape 直接拷贝输入 y shape 的原因。V2 的差异V2 第一段接口多传一个groupListType并将其放入GetWorkspaceSize的形参列表中示例中int groupListType 0;表示 cumsum 模式。当需要直接传“每组大小”时置groupListType 1并传入[64, 0, 128, 64]这类数据即可。5.4 编译与运行示例编译需要 CANN 工具链与 AscendCL 运行环境具体编译命令、链接库-lascendcl等与执行步骤可参考仓库 docs/zh/QUICKSTART.md 与贡献指南 CONTRIBUTING.md 中关于算子样例工程的说明。代码中的#include aclnnop/aclnn_grouped_matmul_add.h头文件由 CANN 的 aclnn 头文件目录提供示例源码位于 examples/test_aclnn_grouped_matmul_add.cpp同一目录下还提供了 V2 的 BF16test_aclnn_grouped_matmul_add_v2_bf16bf16.cpp与 FP16test_aclnn_grouped_matmul_add_v2_fp16fp16.cpp变体分别对应aclnnGroupedMatmulAddV2两种数据类型组合可直接对照使用。六、V1 与 V2 对比及选型建议两个接口的差异可用下表概括对比项aclnnGroupedMatmulAdd (V1)aclnnGroupedMatmulAddV2 (V2)groupList 编码仅 cumsum累积和cumsumtype0或每组大小type1额外属性无 groupListType新增groupListTypeINT64取值 0/1yRef shape2、3 维950 上仅 3 维3 维数据类型组合FP16/FP16/FP32、BF16/BF16/FP32同左计算语义yRef x×weight yRef原地同左yRef 标注为输入/输出选型建议如果分组描述数据已经是累积和形式例如由上层框架直接给出 cumsum 结果直接使用 V1 即可如果需要按每组大小显式描述分组更贴合“各组 K 不等”的 MoE 语义且可表达大小为 0 的空组则应使用 V2 并设置groupListType 1。两种接口在源码中共享同一套aclnnGroupedMatmulAddGetWorkspaceSizeCommon实现op_api/aclnn_grouped_matmul_add.cpp底层算子均为grouped_matmul_add差异仅在 groupList 的解析方式。七、质量保障测试体系一览该算子在仓库中配有较完整的测试体系可作为验证与二次开发的参考Host 侧单测tests/ut/op_host覆盖 Infershapetest_grouped_matmul_add_infershape.cpp、tilingtest_grouped_matmul_add_tiling.cpp以及 op_api 层test_aclnn_grouped_matmul_add.cppKernel 侧单测tests/ut/op_kernel/test_grouped_matmul_add.cpp配套数据生成脚本 gen_data.pyST 测试tests/st基于 ATK 框架分别对 V1executor_aclnnGroupedMatmulAdd.py与 V2executor_aclnnGroupedMatmulAddV2.py进行端到端验证JSON 用例文件定义了输入输出规格executor 中给出了基于 PyTorch 的参考实现torch.matmul 分组切分 累加可作为理解算子语义最直接的“标杆代码”。结语GroupedMatmulAdd 是 CANN ops-transformer 中面向 MoE 分组推理与训练场景的高效算子它以 K 轴分组支持各专家权重 K 维不等以 FLOAT32 累加输出保证精度并通过两段式 aclnn 接口屏蔽了构图与 workspace 管理细节。掌握本文中的公式语义、groupList 两种编码、数据类型组合与示例代码即可在 Ascend 950 / Atlas A2 / A3 系列产品上快速接入该算子若需了解更细的 tiling 策略与 kernel 实现可直接深入 op_host/op_tiling 与 op_kernel 目录继续阅读源码。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-transformer GroupedMatmulAdd 算子全解析：分组矩阵乘与原地累加的 K 轴分组实现与 aclnn 调用指南
CANN ops-transformer GroupedMatmulAdd 算子全解析分组矩阵乘与原地累加的 K 轴分组实现与 aclnn 调用指南【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读本文围绕 CANN ops-transformer 开源仓库中gmm/grouped_matmul_add模块的 GroupedMatmulAdd 算子展开系统讲解其在 NPU 上实现“分组矩阵乘 原地累加”的计算语义、K 轴分组原理、完整参数约束以及通过 AscendCL 两段式接口aclnnGroupedMatmulAdd/aclnnGroupedMatmulAddV2进行编程调用的完整实战流程。读完本文你将掌握该算子的数学定义、groupList 两种编码格式的构造方法、各数据类型的合法组合、典型错误码排查思路并能直接基于仓库中的可运行示例完成 Host 侧代码编写与编译运行。一、算子定位与核心功能1.1 一句话理解 GroupedMatmulAddGroupedMatmulAdd 是 CANN ops-transformer 中用于 MoEMixture of Experts类大模型场景的分组矩阵乘算子。它的核心特征是将多个矩阵乘合并为一次算子调用每组矩阵乘的 K 轴维度可以不同并且结果直接原地累加到输入矩阵 y 上即“add”语义来自原地累加而非独立的偏置相加。以数学形式表达对第 i 组i 1 … gg 为分组个数$$ yRef_i x_i \times weight_i y_i $$其中yRef_i[m_i, n_i] x_i[m_i, k_i] × weight_i[k_i, n_i] y_i[m_i, n_i]m_i / k_i / n_i为对应分组的 shape。可见该算子本质是y x weight的分组版本属于带累加的分组矩阵乘Grouped MatMul with in-place accumulation。1.2 分组方式当前仅支持 K 轴分组算子当前只支持K 轴分组K-axis grouping各分组的k_i各不相同但所有分组的m_i / n_i相同。也就是说输入 x、权重 weight 只在 K第一维方向被切分成 g 段每段独立做矩阵乘后再分别叠加到 y 的对应位置。在仓库源码 op_api/aclnn_grouped_matmul_add.cpp 的参数校验CheckAttrs中明确写死了约束CHECK_COND(params.groupType 2, ACLNN_ERR_PARAM_INVALID, ... [%s] must be 2 ...); CHECK_COND(params.transposeX, ACLNN_ERR_PARAM_INVALID, ... [%s] must be True ...); CHECK_COND(!params.transposeWeight, ACLNN_ERR_PARAM_INVALID, ... [%s] must be False ...);即groupType仅支持 2K 轴分组且当前版本要求x 必须转置、weight 不能转置。1.3 产品支持情况根据 README.md 与接口文档该算子在以下产品上支持/不支持产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从算子定义文件 op_host/grouped_matmul_add_def.cpp 可以看到算子为 AiCore 核注册了ascend910b、ascend910_93、ascend950、ascend350四类 AICore 配置其中后两者启用了动态编译、动态格式、动态 Rank 与动态 Shape 支持与文档列出的 950 系列与 A2/A3 系列产品支持情况相互印证。另外在 Ascend 950 上yRef的 shape 不支持 2 维必须为 3 维。二、参数说明与数据类型约束2.1 完整参数表以下是 README.md 中列出的完整参数说明参数名输入/输出/属性描述数据类型数据格式x输入公式中的输入 xFLOAT16、BFLOAT16NDweight输入公式中的 weightFLOAT16、BFLOAT16NDgroupList输入表示输入 K 轴方向的 matmul 大小分布的 cumsum 结果累积和INT64NDy输入表示原地累加的输出矩阵FLOAT32NDtransposeX属性表示 x 矩阵是否转置BOOL-transposeWeight属性表示 weight 矩阵是否转置BOOL-groupType属性表示分组类型仅支持 2K 轴分组INT64-groupListType属性表示分组 groupList 格式V2 接口引入INT64-yRef输出表示原地累加的输出矩阵FLOAT32ND2.2 数据类型组合支持的输入类型组合严格限定为两种由 op_host/grouped_matmul_add_def.cpp 的 OpDef 注册与 API 层CheckDtype双重约束x 为 FLOAT16、weight 为 FLOAT16、y/yRef 为 FLOAT32x 为 BFLOAT16、weight 为 BFLOAT16、y/yRef 为 FLOAT32。x 与 weight 的数据类型必须一致yRef 固定为 FLOAT32groupList 固定为 INT64。在 op_api/aclnn_grouped_matmul_add.cpp 的CheckDtype中可以看到对应的运行时校验逻辑yRef 非 FLOAT32、groupList 非 INT64、x/weight 非 FP16/BF16 或二者类型不一致时都会返回ACLNN_ERR_PARAM_INVALID。2.3 转置语义与 shape 约定由于transposeX固定为 True、transposeWeight固定为 False实际计算的是$$yRef x^T \times weight y$$即 x 的 shape 为[K, M]转置后按[M, K]参与矩阵乘weight 的 shape 为[K, N]。API 层的 shape 校验aclnn_grouped_matmul_add.cpp进一步明确x、weight 均为 2 维xDimNum 2、weightDimNum 2groupList 为 1 维x 与 weight 的 K 维第一维必须相等aKDim bKDimyRef 的元素总数必须等于M × N × groupNum代码按 3 维 shape 计算期望大小并做整体尺寸校验。2.4 约束说明使用该算子前需重点核对以下约束来自 README.md 与 aclnnGroupedMatmulAdd.mdint32 边界x 和 weight 中每一组 tensor 的每一维大小在 32 字节对齐后都应小于 int32 的最大值 2147483647。最后一维上限x 和 weight 中每一组 tensor 的最后一维大小都应小于 65536。x 的最后一维指当 x 不转置时的 K 轴或 x 转置时的 M 轴weight 的最后一维指 weight 不转置时的 N 轴或转置时的 K 轴。确定性aclnnGroupedMatmulAdd与 V2 接口默认均为确定性实现deterministic。分组上限groupList 第 1 维最大支持 1024即最多支持 1024 个 group。三、groupList 分组描述从 cumsum 到每组大小groupList 是理解该算子的关键输入。它是一段 Device 侧的 INT64 aclTensor描述 K 轴方向上各分组的切分位置。V1 接口aclnnGroupedMatmulAdd只支持 cumsum累积和模式而V2 接口aclnnGroupedMatmulAddV2通过新增的groupListType属性同时支持两种格式groupListTypegroupList 语义约束0各分组 K 轴大小的cumsum累积和即单调非递减序列最后一个值不大于 x 第一维K 轴总长1各分组 K 轴的每组大小V2 新增数值非负总和不超过 x 第一维文档中的例子非常直观。假设 K 256、分组数 E 4各组大小依次为64、0、128、64groupListType 0cumsum[64, 64, 192, 256]第 2 组大小为 0因此累积和与第 1 组相同groupListType 1每组大小[64, 0, 128, 64]。从 op_api/aclnn_grouped_matmul_add.cpp 的CheckAttrs可以看到groupListType的合法取值为 0 或 1V1 接口在构造参数时未显式传入该值实际默认按 cumsum 模式0处理。在仓库的 ST 测试基准实现 tests/st/aclnnGroupedMatmulAdd/executor_aclnnGroupedMatmulAdd.py 中可以清晰看到 cumsum 模式的计算方式——按 groupList 中的累积和切分 K 轴逐组做矩阵乘再求和last 0 for i in groupList.tolist(): x_tensor x[last:i, :].to(torch.float64) weight_tensor weight[last:i, :].to(torch.float64) result.append(torch.matmul(x_tensor.t(), weight_tensor)) last i result torch.stack(result).reshape(y.shape) y.cpu()这正好对应公式yRef_i x_i^T × weight_i y_i按[last:i]切分 K 轴x 转置后与 weight 相乘最后堆叠 reshape 并与 y 相加。四、算子工程结构从 API 到 Kernel 的完整链路gmm/grouped_matmul_add是一个结构完整的自定义算子工程源码组织如下op_api/aclnn 两段式接口实现与l0op::GroupedMatmulAdd图级接口封装aclnn_grouped_matmul_add.cpp、grouped_matmul_add.cppop_host/算子注册grouped_matmul_add_def.cpp、Infershape 推导grouped_matmul_add_infershape.cpp以及 tiling 逻辑op_tiling/grouped_matmul_add_tiling.hop_kernel/NPU 侧 AiCore kernel 实现op_kernel/grouped_matmul_add.cpp及 arch35 架构的 kernel 与 tiling 数据结构examples/可直接参考的可运行 Host 侧调用示例tests/包含 utHost 侧 tiling / infershape / op_api 单测、kernel 单测与 st基于 ATK 框架的 aclnn 接口测试含 JSON 用例与 Python executor。4.1 Infershape输出与输入 y 同形在 op_host/grouped_matmul_add_infershape.cpp 中InferShapeGroupedMatmulAdd直接令输出yRefShape yShape原地累加输出即输入 y 本身InferDataTypeGroupedMatmulAdd校验 x 与 weight 数据类型一致并将输出数据类型固定为DT_FLOAT。这与“yRef 是 y 的原地别名”的语义完全一致——在示例代码中体现为yRef y;这一行赋值。4.2 Tiling 数据分组信息如何在 Host 侧编排Tiling 数据结构定义在 op_host/op_tiling/grouped_matmul_add_tiling.hBEGIN_TILING_DATA_DEF(GmmBaseParams) TILING_DATA_FIELD_DEF(int64_t, groupNum); // 分组个数 TILING_DATA_FIELD_DEF(int64_t, coreNum); // 使用的核数 TILING_DATA_FIELD_DEF(int64_t, groupType); // 分组类型仅支持2 --- K分组 TILING_DATA_FIELD_DEF(int64_t, groupListType); END_TILING_DATA_DEF其中GroupedMatmulAddTilingData由GmmBaseParams分组基础信息与TCubeTiling矩阵乘 cube 单元 tiling 信息组合而成并配套定义了MAX_TENSOR_LIST_SIZE 128的常量。可以推断Host 侧 tiling 负责把 groupList 解析为分组个数与各组的矩阵乘切分方案交给 Kernel 侧的 cube 流水执行。五、两段式 aclnn 接口与调用实战5.1 两段式接口原型与 CANN 其他算子一致GroupedMatmulAdd 采用两段式接口two-phase API第一段GetWorkspaceSize完成入参校验、构图并计算 workspace 大小第二段执行实际计算。函数原型如下见 docs/aclnnGroupedMatmulAdd.md。V1 第一段aclnnStatus aclnnGroupedMatmulAddGetWorkspaceSize( const aclTensor *x, const aclTensor *weight, const aclTensor *groupList, aclTensor *yRef, bool transposeX, bool transposeWeight, int64_t groupType, uint64_t *workspaceSize, aclOpExecutor **executor)V2 第一段多出groupListType参数aclnnStatus aclnnGroupedMatmulAddV2GetWorkspaceSize( const aclTensor *x, const aclTensor *weight, const aclTensor *groupList, aclTensor *yRef, bool transposeX, bool transposeWeight, int64_t groupType, int64_t groupListType, uint64_t *workspaceSize, aclOpExecutor **executor)第二段两版本一致aclnnStatus aclnnGroupedMatmulAdd(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)从实现源码 op_api/aclnn_grouped_matmul_add.cpp 可以看到两段接口内部的实际流程创建 OpExecutor → 参数校验CheckParams→ 空 tensor 检查任一分组轴为 0 时直接返回空输出→ 数据连续性转换Contiguous→ 调用l0op::GroupedMatmulAdd构图 → 对非连续输出做ViewCopy→ 返回 workspace 大小。5.2 错误码与排查第一段接口完成入参校验文档列出的典型错误如下返回值错误码触发场景ACLNN_ERR_PARAM_NULLPTR161001x、weight、groupList、yRef 为空指针ACLNN_ERR_PARAM_INVALID161002数据类型/格式不在支持范围或 x 与 weight 数据类型不一致ACLNN_ERR_INNER_TILING_ERROR561002x、weight、yRef 的 shape 不满足矩阵乘限制要求与源码对应CheckNotNull对应 161001CheckFormat/CheckShape/CheckDtype/CheckAttrs统一返回 161002shape 不满足矩阵乘限制时返回 561002。5.3 完整调用示例可编译运行仓库提供了可直接编译运行的示例 examples/test_aclnn_grouped_matmul_add.cpp下面是梳理后的完整调用流程V1FP16 版本V2 示例见 examples/test_aclnn_grouped_matmul_add_v2.cpp。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_grouped_matmul_add.h // ... CHECK_RET / LOG_PRINT / GetShapeSize 宏与工具函数参考仓库示例 int main() { // 1. 固定写法device/stream 初始化 int32_t deviceId 0; aclrtStream stream; Init(deviceId, stream); // 2. 构造输入与输出K轴分组K512MN256g2 std::vectorint64_t xShape {512, 256}; // x: [K, M]因 transposeXtrue std::vectorint64_t weightShape {512, 256}; // weight: [K, N] std::vectorint64_t yShape {2, 256, 256}; // y/yRef: [g, M, N] std::vectorint64_t groupListShape {2}; std::vectorint64_t groupListData {256, 512}; // cumsum: 第1组K256第2组K256 void *xDeviceAddr, *weightDeviceAddr, *yDeviceAddr, *groupListDeviceAddr; aclTensor *x nullptr, *weight nullptr, *groupedList nullptr; aclTensor *y nullptr, *yRef nullptr; bool transpose_x true; // 仅支持 True bool transpose_weight false; // 仅支持 False int group_type 2; // 仅支持 2K轴分组 // 创建 aclTensorx/weight 用 FP16uint16_t 承载y 用 FP32groupList 用 INT64 CreateAclTensoruint16_t(xShape, xDeviceAddr, aclDataType::ACL_FLOAT16, x); CreateAclTensoruint16_t(weightShape, weightDeviceAddr, aclDataType::ACL_FLOAT16, weight); CreateAclTensorfloat(yShape, yDeviceAddr, aclDataType::ACL_FLOAT, y); CreateAclTensor_Newint64_t(groupListData, groupListShape, groupListDeviceAddr, aclDataType::ACL_INT64, groupedList); yRef y; // 原地累加输出与输入 y 是同一块内存 // 3. 第一段接口计算 workspace 并获取 executor uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; aclnnGroupedMatmulAddGetWorkspaceSize(x, weight, groupedList, yRef, transpose_x, transpose_weight, group_type, workspaceSize, executor); // 按需申请 workspace void *workspaceAddr nullptr; if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 4. 第二段接口执行计算 aclnnGroupedMatmulAdd(workspaceAddr, workspaceSize, executor, stream); // 5. 同步等待执行结束拷贝结果回 Host 并打印 aclrtSynchronizeStream(stream); auto size GetShapeSize(yShape); std::vectorfloat resultData(size, 0); aclrtMemcpy(resultData.data(), size * sizeof(float), yDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); // 6. 释放 tensor、workspace、stream 与 device 资源 aclDestroyTensor(x); aclDestroyTensor(weight); aclDestroyTensor(y); aclrtFree(xDeviceAddr); aclrtFree(weightDeviceAddr); aclrtFree(yDeviceAddr); if (workspaceSize 0) aclrtFree(workspaceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }关键点说明x 的 shape 是[K, M]因为transposeX truex 实际以[K, M]存储、按[M, K]参与运算因此示例中xShape {512, 256}表示 K512、M256。groupList 传的是累积和{256, 512}表示两组 K 分别为 256 和 256合计 512与 x 的 K 维一致。yRef y输出与输入共享同一块 Device 内存实现真正的原地累加这也是 Infershape 中输出 shape 直接拷贝输入 y shape 的原因。V2 的差异V2 第一段接口多传一个groupListType并将其放入GetWorkspaceSize的形参列表中示例中int groupListType 0;表示 cumsum 模式。当需要直接传“每组大小”时置groupListType 1并传入[64, 0, 128, 64]这类数据即可。5.4 编译与运行示例编译需要 CANN 工具链与 AscendCL 运行环境具体编译命令、链接库-lascendcl等与执行步骤可参考仓库 docs/zh/QUICKSTART.md 与贡献指南 CONTRIBUTING.md 中关于算子样例工程的说明。代码中的#include aclnnop/aclnn_grouped_matmul_add.h头文件由 CANN 的 aclnn 头文件目录提供示例源码位于 examples/test_aclnn_grouped_matmul_add.cpp同一目录下还提供了 V2 的 BF16test_aclnn_grouped_matmul_add_v2_bf16bf16.cpp与 FP16test_aclnn_grouped_matmul_add_v2_fp16fp16.cpp变体分别对应aclnnGroupedMatmulAddV2两种数据类型组合可直接对照使用。六、V1 与 V2 对比及选型建议两个接口的差异可用下表概括对比项aclnnGroupedMatmulAdd (V1)aclnnGroupedMatmulAddV2 (V2)groupList 编码仅 cumsum累积和cumsumtype0或每组大小type1额外属性无 groupListType新增groupListTypeINT64取值 0/1yRef shape2、3 维950 上仅 3 维3 维数据类型组合FP16/FP16/FP32、BF16/BF16/FP32同左计算语义yRef x×weight yRef原地同左yRef 标注为输入/输出选型建议如果分组描述数据已经是累积和形式例如由上层框架直接给出 cumsum 结果直接使用 V1 即可如果需要按每组大小显式描述分组更贴合“各组 K 不等”的 MoE 语义且可表达大小为 0 的空组则应使用 V2 并设置groupListType 1。两种接口在源码中共享同一套aclnnGroupedMatmulAddGetWorkspaceSizeCommon实现op_api/aclnn_grouped_matmul_add.cpp底层算子均为grouped_matmul_add差异仅在 groupList 的解析方式。七、质量保障测试体系一览该算子在仓库中配有较完整的测试体系可作为验证与二次开发的参考Host 侧单测tests/ut/op_host覆盖 Infershapetest_grouped_matmul_add_infershape.cpp、tilingtest_grouped_matmul_add_tiling.cpp以及 op_api 层test_aclnn_grouped_matmul_add.cppKernel 侧单测tests/ut/op_kernel/test_grouped_matmul_add.cpp配套数据生成脚本 gen_data.pyST 测试tests/st基于 ATK 框架分别对 V1executor_aclnnGroupedMatmulAdd.py与 V2executor_aclnnGroupedMatmulAddV2.py进行端到端验证JSON 用例文件定义了输入输出规格executor 中给出了基于 PyTorch 的参考实现torch.matmul 分组切分 累加可作为理解算子语义最直接的“标杆代码”。结语GroupedMatmulAdd 是 CANN ops-transformer 中面向 MoE 分组推理与训练场景的高效算子它以 K 轴分组支持各专家权重 K 维不等以 FLOAT32 累加输出保证精度并通过两段式 aclnn 接口屏蔽了构图与 workspace 管理细节。掌握本文中的公式语义、groupList 两种编码、数据类型组合与示例代码即可在 Ascend 950 / Atlas A2 / A3 系列产品上快速接入该算子若需了解更细的 tiling 策略与 kernel 实现可直接深入 op_host/op_tiling 与 op_kernel 目录继续阅读源码。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考