CANN ops-transformer 通算融合算子 aclnnMatmulAllReduceV3MatMul 与 AllReduce 深度融合原理与调用实战【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本文围绕 CANN ops-transformer 仓库中mc2/matmul_all_reduce模块的aclnnMatmulAllReduceV3接口系统讲解其产品支持情况、融合计算公式、两段式 API 函数原型、全部入参与属性约束、通信引擎 commMode 选型以及完整的多卡调用示例并结合仓库源码与单元测试深入剖析其入参校验、通信引擎传递与确定性计算等实现细节帮助开发者正确、高效地在 NPU 上使用该通算融合算子。算子定位为什么需要 MatMul AllReduce 融合在 Transformer 类大模型训练与推理中MatMul矩阵乘与AllReduce全量规约通信是出现频率最高的两类算子。传统实现将两者作为独立算子先后执行先由计算核完成MatMul将中间结果落盘到 Global Memory再启动通信引擎把各卡结果做AllReduce归约。这带来了两次额外的访存开销与一次算子调度开销在大规模并行如专家并行、张量并行场景下通信与计算的重叠效率直接决定集群扩展性。aclnnMatmulAllReduceV3将两者融合为单一算子在计算核完成矩阵乘与逐元素加法的同时把结果通过通信引擎直接完成跨卡归约减少中间张量搬运与算子下发次数。该接口兼容aclnnMatmulAllReduceV2的全部功能并新增通信引擎参数commMode允许调用方在 AICPU 与 CCU 两种通信引擎之间自选不同芯片支持情况不同从而在不同数据量、不同组网规模下获得更优的通信性能。产品支持情况aclnnMatmulAllReduceV3在 CANN ops-transformer 当前仓库中的产品支持矩阵如下产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持版本前提使用该接口时请确保驱动固件包和 CANN 包均为配套的 8.0.RC2 版本或配套的更高版本否则会引发报错如 Bus Error 等。该信息与仓库 READMEmc2/matmul_all_reduce/README.md中的产品支持矩阵一致且 README 进一步说明A2 系列支持 1/2/4/8 卡950 系列支持 1/2/4/8/16/32/64 卡且仅支持 HCCS 链路 all mesh 组网。功能说明与计算公式接口功能完成 MatMul 计算与 AllReduce 通信融合兼容aclnnMatmulAllReduceV2支持的功能新增通信引擎参数commMode支持自选通信引擎。计算公式$$ output AllReduce(x1 x2 bias x3) $$即先计算x1与x2的矩阵乘叠加一维偏置bias与同 shape 的x3再在指定通信域group上按reduceOp当前仅支持 sum做一次全量归约得到最终输出output。需要说明的是bias、x3 在aclnnMatmulAllReduceV3中均可选输入可为空与仓库 README.md 中非量化场景情形2的公式保持一致。两段式接口与函数原型与 CANN AscendCL 其它算子一样aclnnMatmulAllReduceV3采用两段式接口先调用aclnnMatmulAllReduceV3GetWorkspaceSize完成入参校验获取计算所需 workspace 大小以及包含算子计算流程的执行器executor再调用aclnnMatmulAllReduceV3真正执行计算。两段式接口的函数原型如下aclnnStatus aclnnMatmulAllReduceV3GetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, const aclTensor *bias, const aclTensor *x3, const char *group, const char *reduceOp, const char *commMode, int64_t commTurn, int64_t streamMode, const aclTensor *output, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnMatmulAllReduceV3( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)该原型定义可在仓库头文件 mc2/matmul_all_reduce/op_api/aclnn_matmul_all_reduce_v3.h 中查到头文件中domain aclnn_ops_infer标注了该接口面向推理场景的域归属。aclnnMatmulAllReduceV3GetWorkspaceSize 参数详解第一阶段接口共 12 个参数完整说明如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1输入MatMul 计算的左矩阵即计算公式中的 x1当前版本仅支持二维或三维输入支持不转置场景FLOAT16、BFLOAT16ND2-3×x2输入MatMul 计算的右矩阵即计算公式中的 x2当前版本仅支持二维输入支持转置/不转置场景ND 格式下支持最后两轴转置情况下的非连续 tensor其他非连续 tensor 不支持FLOAT16、BFLOAT16ND2×x3输入MatMul 计算后的 add 计算即计算公式中的 x3shape 与 MatMul 计算后的 shape 一致FLOAT16、BFLOAT16ND2√bias输入对应计算公式中的 bias 偏移当前版本仅支持一维输入FLOAT16、BFLOAT16ND0-1√group输入通信域名称通过 Hccl 提供的接口extern HcclResult HcclGetCommName(HcclComm comm, char* commName);获取其中 commName 即为 groupString---reduceOp输入reduce 操作类型当前版本仅支持输入 sumString---commMode输入通信引擎参数用于指定通信引擎根据芯片型号有不同的输入限制详见下文约束说明String---commTurn输入通信数据切分数即总数据量/单次通信量当前版本仅支持输入 0INT64---streamMode输入AscendCL 流模式的枚举当前版本仅支持枚举值 1INT64---output输出MatMul 计算与 AllReduce 通信的结果即计算公式中的 outputoutput 的维数与 x1 一致FLOAT16、BFLOAT16ND2-3√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----几个需要特别留意的要点x2 的转置与非连续接口支持 x2 转置由内部IsTransposeLastTwoDims依据 stride 判断且仅在最后两轴转置场景下才允许 x2 非连续其余非连续 x2 一律返回参数非法。源码实现见 mc2/matmul_all_reduce/op_api/matmul_all_reduce_util.cpp 中MatmulAllReduceCheckValidContiguous。x1 不支持转置op_api 实现中transposeX1固定为false见 mc2/matmul_all_reduce/op_api/aclnn_matmul_all_reduce_v3.cpp与文档约束一致。bias 的维度文档标注 shape 为 0-1 维实际仅支持一维输入且一维长度须等于 output 最后一维即 n0 维标量场景对应空 bias 的退化情形。aclnnMatmulAllReduceV3 执行接口参数详解第二阶段接口共 4 个参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnMatmulAllReduceV3GetWorkspaceSize获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值与入参校验错误码两段式接口均返回aclnnStatus状态码完整错误码说明参见 aclnn 返回码。第一阶段接口会完成入参校验以下场景将报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 x1、x2 或 output 是空指针ACLNN_ERR_PARAM_INVALID161002x1、x2、bias、x3 或 output 的数据类型不在支持范围之内ACLNN_ERR_PARAM_INVALID161002reduceOp、streamMode 不在合法范围内ACLNN_ERR_PARAM_INVALID161002x1、x2、bias、x3 或 output 的 shape 不符合约束要求从源码看第一段接口的校验顺序非常清晰mc2/matmul_all_reduce/op_api/matmul_all_reduce_util.cpp 中MatmulAllReduceCheckParams空指针检查MatmulAllReduceCheckNotNull(x1, x2, output)数据类型检查MatmulAllReduceCheckDtypeValid要求 x1、x2、x3、bias、output 的类型都在 FLOAT16/BFLOAT16 支持列表内且 x1、x2、x3、bias、output 彼此数据类型一致属性检查MatmulAllReduceCheckAttrreduceOp必须等于sumstrcmp与REDUCE_OP_SUM比对streamMode必须等于NUM_ACL_STOP_ON_FAILURE即 1格式检查非 310P 场景下 x2 必须为 ND 格式shape 检查MatmulAllReduceCheckShapex2 必须 2 维x1 维数在 2~3 之间x1 最后一维k必须等于 x2 第一维koutput shape 等于 x1 前序维拼接 x2 最后一维bias 必须一维且长度为 nx3 与 output shape 一致连续性检查A2ASCEND910B场景下x2 非连续时必须为转置场景否则报ACLNN_ERR_PARAM_INVALID。此外aclnnMatmulAllReduceV3在入参前会额外校验commMode合法性IsCommModeValid非法值直接返回ACLNN_ERR_PARAM_INVALID。约束说明通信引擎 commMode 支持度commMode是本接口相对 V2 的核心新增参数用于自选通信引擎Atlas A2 训练/推理系列产品目前不支持指定通信引擎commMode仅允许输入ai_cpu即使用 AICPU 通信引擎Ascend 950PR / Ascend 950DT通信引擎支持 AICPU 与 CCUcommMode允许输入ai_cpu或ccu。其中 CCU 仅支持单机 UB 域内互联AICPU 可支持跨机 UB 域内互联同一条通信链路内只能选择同一种通信引擎。两种通信引擎的特点如下AICPU不占用计算核通信效率高但通信静态开销较大对小数据量通信场景不友好适用于大数据高带宽场景CCU能够减少访存带宽与计算核占用但受限于片上资源支持的通信域数量有限适用于高带宽、低时延的通信场景。从源码可以看到该参数在运行时如何生效mc2/matmul_all_reduce/op_api/aclnn_matmul_all_reduce_v3.cppGetWorkspaceSize阶段将字符串commMode映射为枚举值ai_cpu→Mc2Comm::COMM_MODE_AICPUccu→Mc2Comm::COMM_MODE_CCU通过NnopbaseSetUserHandle挂到 executor 上若既不是ai_cpu也不是ccu直接返回ACLNN_ERR_PARAM_INVALID执行阶段aclnnMatmulAllReduceV3依据当前 NPU 架构决定 HCCL server 类型——在 950源码中DAV_3510上按 handle 中保存的 commMode 设置NNOPBASE_HCCL_SERVER_TYPE_CCU或NNOPBASE_HCCL_SERVER_TYPE_AICPU在 A2 上则固定使用 AICPU。单元测试 CSVmc2/matmul_all_reduce/tests/ut/op_api/test_aclnn_matmul_all_reduce_v3.csv也验证了这一约束A2 上传入ccu返回PARAM_INVALID950 上ccu与ai_cpu均 SUCCESS而非法字符串invalid返回PARAM_INVALID。确定性计算Atlas A2 训练/推理系列产品aclnnMatmulAllReduceV3默认非确定性实现支持通过配置HCCL_DETERMINISTIC环境变量为true开启确定性计算Ascend 950PR / Ascend 950DTaclnnMatmulAllReduceV3默认确定性实现。确定性计算的通用背景可参考仓库文档 确定性计算说明。shape、数据类型与组网约束增量场景不使能 MC2全量场景使能 MC2输入 x1 可为二维或三维其 shape 为(b, s, k)或(m, k)x2 必须为二维shape 为(k, n)轴满足 mm 算子入参要求k 轴相等bias 若非空其 shape 为(n)b*s、m、k、n的值均不得超过 2147483647INT32_MAX当输入 x1 的 shape 为(b, s, k)时输出 output 的 shape 为(b, s, n)当输入 x1 的 shape 为(m, k)时输出 output 的 shape 为(m, n)x1、x2、bias 计算输入的数据类型要和 output 计算输出的数据类型一致传入的 x1、x2、x3 或 output 不为空指针仅支持 HCCS 链路 all mesh 组网Atlas A2 训练/推理系列产品支持 1、2、4、8 卡Ascend 950PR / Ascend 950DT支持 1、2、4、8、16、32、64 卡Atlas A2 训练/推理系列产品一个模型中的通算融合 MC2 算子仅支持相同通信域。空 tensor 支持度支持 k 为 0 的场景此时输出为bias x3支持bs/m/n为 0此时传入的输出也应该是空 tensor此场景不进入 kernel 计算直接返回。上述空 tensor 行为在测试用例中同样有覆盖CSV 中empty_K、empty_M、empty_N三个用例在 A2 上均返回 SUCCESS。完整调用示例以下为文档提供的完整多卡调用示例展示了从设备初始化、HCCL 通信域创建、tensor 创建、两段式接口调用到资源释放的完整流程。示例基于 FP16、x1(32,64) x2(64,128) bias(128) x3(32,128)的 2 卡场景。编译与执行过程请参考 编译与运行样例。说明示例代码调用了部分 HCCL 集合通信库接口HcclGetCommName、HcclCommInitAll、HcclCommDestroy这些接口来自 HCCL APIC头文件hccl/hccl.h随 CANN 环境一并提供。#include iostream #include vector #include thread #include hccl/hccl.h #include aclnn/opdev/fp16_t.h #include aclnnop/aclnn_matmul_all_reduce_v3.h int ndev 2; #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i: shape) { shapeSize * i; } return shapeSize; } templatetypename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void **deviceAddr, aclDataType dataType, aclTensor **tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } struct Args { uint32_t rankId; HcclComm hcclComm; aclrtStream stream; aclrtContext context; }; int launchOneThreadMatmulAllReduce(Args args) { int ret; ret aclrtSetCurrentContext(args.context); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetCurrentContext failed. ERROR: %d\n, ret); return ret); char hcom_name[128]; ret HcclGetCommName(args.hcclComm, hcom_name); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT([ERROR] HcclGetCommName failed. ret %d \n, ret); return -1); LOG_PRINT([INFO] rank %d hcom: %s stream: %p, context : %p\n, args.rankId, hcom_name, args.stream, args.context); std::vectorint64_t x1Shape {32, 64}; std::vectorint64_t x2Shape {64, 128}; std::vectorint64_t biasShape {128}; std::vectorint64_t x3Shape {32, 128}; std::vectorint64_t outShape {32, 128}; void *x1DeviceAddr nullptr; void *x2DeviceAddr nullptr; void *biasDeviceAddr nullptr; void *x3DeviceAddr nullptr; void *outDeviceAddr nullptr; aclTensor *x1 nullptr; aclTensor *x2 nullptr; aclTensor *bias nullptr; aclTensor *x3 nullptr; aclTensor *out nullptr; int64_t commTurn 0; int64_t streamMode 1; uint64_t workspaceSize 0; aclOpExecutor *executor; void *workspaceAddr nullptr; long long x1ShapeSize GetShapeSize(x1Shape); long long x2ShapeSize GetShapeSize(x2Shape); long long biasShapeSize GetShapeSize(biasShape); long long x3ShapeSize GetShapeSize(x3Shape); long long outShapeSize GetShapeSize(outShape); std::vectorop::fp16_t x1HostData(x1ShapeSize, 1); std::vectorop::fp16_t x2HostData(x2ShapeSize, 1); std::vectorop::fp16_t biasHostData(biasShapeSize, 1); std::vectorop::fp16_t x3HostData(x3ShapeSize, 1); std::vectorop::fp16_t outHostData(outShapeSize, 0); // 创建tensor ret CreateAclTensor(x1HostData, x1Shape, x1DeviceAddr, aclDataType::ACL_FLOAT16, x1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(x2HostData, x2Shape, x2DeviceAddr, aclDataType::ACL_FLOAT16, x2); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(biasHostData, biasShape, biasDeviceAddr, aclDataType::ACL_FLOAT16, bias); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(x3HostData, x3Shape, x3DeviceAddr, aclDataType::ACL_FLOAT16, x3); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT16, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 调用第一段接口 ret aclnnMatmulAllReduceV3GetWorkspaceSize(x1, x2, bias, x3, hcom_name, sum, , commTurn, streamMode, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMatmulAllReduceV3GetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用第二段接口 ret aclnnMatmulAllReduceV3(workspaceAddr, workspaceSize, executor, args.stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMatmulAllReduceV3 failed. ERROR: %d\n, ret); return ret); //固定写法同步等待任务执行结束 ret aclrtSynchronizeStreamWithTimeout(args.stream, 10000); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); LOG_PRINT(device%d aclnnMatmulAllReduceV3 execute success \n, args.rankId); // 释放device资源需要根据具体API的接口定义修改 if (x1 ! nullptr) { aclDestroyTensor(x1); } if (x2 ! nullptr) { aclDestroyTensor(x2); } if (bias ! nullptr) { aclDestroyTensor(bias); } if (x3 ! nullptr) { aclDestroyTensor(x3); } if (out ! nullptr) { aclDestroyTensor(out); } if (x1DeviceAddr ! nullptr) { aclrtFree(x1DeviceAddr); } if (x2DeviceAddr ! nullptr) { aclrtFree(x2DeviceAddr); } if (biasDeviceAddr ! nullptr) { aclrtFree(biasDeviceAddr); } if (x3DeviceAddr ! nullptr) { aclrtFree(x3DeviceAddr); } if (outDeviceAddr ! nullptr) { aclrtFree(outDeviceAddr); } if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(args.stream); HcclCommDestroy(args.hcclComm); aclrtDestroyContext(args.context); aclrtResetDevice(args.rankId); return 0; } int main(int argc, char *argv[]) { int ret; int32_t devices[ndev]; for (int i 0; i ndev; i) { devices[i] i; } HcclComm comms[128]; ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); // 初始化集合通信域 for (int i 0; i ndev; i) { ret aclrtSetDevice(devices[i]); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); } ret HcclCommInitAll(ndev, devices, comms); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(HcclCommInitAll failed. ERROR: %d\n, ret); return ret); Args args[ndev]; aclrtStream stream[ndev]; aclrtContext context[ndev]; for (uint32_t rankId 0; rankId ndev; rankId) { ret aclrtSetDevice(rankId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateContext(context[rankId], rankId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateContext failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream[rankId]); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); } // 启动多线程 std::vectorstd::unique_ptrstd::thread threads(ndev); for (uint32_t rankId 0; rankId ndev; rankId) { args[rankId].rankId rankId; args[rankId].hcclComm comms[rankId]; args[rankId].stream stream[rankId]; args[rankId].context context[rankId]; threads[rankId].reset(new(std::nothrow) std::thread(launchOneThreadMatmulAllReduce, std::ref(args[rankId]))); } for (uint32_t rankId 0; rankId ndev; rankId) { threads[rankId]-join(); } aclFinalize(); return 0; }调用流程的关键节点拆解如下通信域准备aclInit后通过aclrtSetDevice绑定多卡用HcclCommInitAll一次性创建多卡的 HCCL 通信域每个 rank 的线程通过HcclGetCommName拿到通信域名字符串作为group参数tensor 创建CreateAclTensor内部先aclrtMalloc申请 Device 内存、aclrtMemcpy拷入 host 数据再按连续布局计算 strides 并调用aclCreateTensor创建aclTensor第一段接口aclnnMatmulAllReduceV3GetWorkspaceSize完成校验并返回workspaceSize与executor若workspaceSize 0则aclrtMalloc申请 workspace第二段接口aclnnMatmulAllReduceV3提交任务到指定 stream随后用aclrtSynchronizeStreamWithTimeout同步等待执行完成资源回收依次销毁aclTensor、释放 Device 内存与 workspace、销毁 stream、销毁 HCCL 通信域与 context最后aclrtResetDevice与aclFinalize。示例中第一段接口的commMode传入了空字符串表示由系统按当前芯片选择默认通信引擎A2 固定为ai_cpu950 默认引擎实际开发中可按上文约束说明显式传入ai_cpu或ccu。从源码看实现原理op_api 层参数校验与 commMode 传递aclnnMatmulAllReduceV3的 op_api 实现位于 mc2/matmul_all_reduce/op_api/aclnn_matmul_all_reduce_v3.cpp。核心逻辑分三步入参校验aclnnMatmulAllReduceV3GetWorkspaceSize先调用MatmulAllReduceCheckParams完成空指针、dtype、attrreduceOp/streamMode、格式、shape、连续性检查再调用IsCommModeValid校验 commMode 字符串合法性下调内核接口内部将 x1 固定为不转置transposeX1 false通过IsTransposeLastTwoDims(x2)自动识别 x2 是否转置然后调用aclnnInnerMatmulAllReduceGetWorkspaceSize下发 IR该内核接口同时承载量化分支因此 V3 在此通过NnopbaseDisableOptionalInput将 4~9 号量化相关可选输入scale/offset/dequant/pertoken/commQuantScale 等全部禁用commMode 传递commMode字符串被映射为枚举后经NnopbaseSetUserHandle挂到 executor 上执行阶段aclnnMatmulAllReduceV3再按 950/非 950 分别设置NNOPBASE_HCCL_SERVER_TYPE_CCU或NNOPBASE_HCCL_SERVER_TYPE_AICPU最终调用aclnnInnerMatmulAllReduce真正下发执行。公共的校验工具函数集中在 mc2/matmul_all_reduce/op_api/matmul_all_reduce_util.h 与 mc2/matmul_all_reduce/op_api/matmul_all_reduce_util.cpp 中包括MatmulAllReduceCheckParams、MatmulAllReduceCheckNotNull、MatmulAllReduceCheckDtypeValid、MatmulAllReduceCheckAttr、MatmulAllReduceCheckShape、MatmulAllReduceCheckValidEmptyTensor、MatmulAllReduceCheckValidContiguous等V2/V3 及量化系列接口共用。算子定义层动态特性配置算子的 IR 定义在 mc2/matmul_all_reduce/op_host/matmul_all_reduce_def.cpp。MatmulAllReduceOpDef 为输出y配置了 950 全系列的 FLOAT16/BFLOAT16/FLOAT 支持列表与 ND 格式并为 AICore 配置了动态编译、动态格式、动态 rank 与动态 shape 支持意味着该算子支持动态 shape 下运行时重新做 tiling 切分。值得说明的是该 OpDef 同时承载量化分支README 中的情形 3~7因此 dtype/format 列表比 V3 文档限定的 FLOAT16/BFLOAT16 更宽V3 接口本身仍以文档约束为准。从测试用例验证行为单元测试 mc2/matmul_all_reduce/tests/ut/op_api/test_aclnn_matmul_all_reduce_v3.cpp 通过 CSV 参数化驱动测试用例数据见 mc2/matmul_all_reduce/tests/ut/op_api/test_aclnn_matmul_all_reduce_v3.csv覆盖了合法场景2D/3D x1、BF16、empty_K/empty_M/empty_N空 tensor、950 上的ccu与ai_cpu通信引擎非法场景x1 四维、x1 空、x2 三维、x3 三维、reduceOp 为max、streamMode 为 0/-1、dtype 不匹配、k/n 轴不匹配、bias/x3 与 output 不匹配、x2 为 FRACTAL_NZ 格式、A2 上传ccu、非法 commMode 字符串等均返回ACLNN_ERR_PARAM_INVALID额外用例转置 x2非连续但满足最后两轴转置SUCCESS非转置非连续 x2 返回PARAM_INVALIDexecutor为空时第二段接口返回ACLNN_ERR_INNER。接口家族与版本演进aclnnMatmulAllReduceV3属于matmul_all_reduce算子家族该模块 docs 目录下共存有非量化、量化、权重量化三个方向的接口文档可交叉参考非量化aclnnMatmulAllReduce、aclnnMatmulAllReduceV2、aclnnMatmulAllReduceV3本文量化INT8/FP8 等aclnnQuantMatmulAllReduce至aclnnQuantMatmulAllReduceV5权重量化aclnnWeightQuantMatmulAllReduce、aclnnWeightQuantMatmulAllReduceV2。V3 相对 V2 的核心差异在于新增commMode通信引擎参数允许 950 系列在 AICPU 与 CCU 之间自选为高带宽、低时延场景提供了更细粒度的通信调优手段计算语义上两者一致均实现output AllReduce(x1 x2 bias x3)。相关文档入口mc2/matmul_all_reduce/docs/aclnnMatmulAllReduceV2.md 与 mc2/matmul_all_reduce/README.md。总结aclnnMatmulAllReduceV3是 CANN ops-transformer 中面向 Transformer 大模型全量并行训练/推理场景的典型通算融合算子通过将矩阵乘、偏置加法、x3累加与跨卡 AllReduce 合入一次下发有效压缩中间访存与调度开销。使用该接口时需重点把握两段式调用流程、FLOAT16/BFLOAT16 与 ND 格式约束、sum/1/0 三个属性的取值限制、x2 转置场景下的非连续支持、以及commMode随芯片的差异化选型A2 仅ai_cpu950 可选ai_cpu/ccu。结合仓库源码与单元测试开发者可以据此快速定位参数校验失败原因并在多卡组网上正确搭建通算融合的计算流程。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-transformer 通算融合算子 aclnnMatmulAllReduceV3：MatMul 与 AllReduce 深度融合原理与调用实战
CANN ops-transformer 通算融合算子 aclnnMatmulAllReduceV3MatMul 与 AllReduce 深度融合原理与调用实战【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本文围绕 CANN ops-transformer 仓库中mc2/matmul_all_reduce模块的aclnnMatmulAllReduceV3接口系统讲解其产品支持情况、融合计算公式、两段式 API 函数原型、全部入参与属性约束、通信引擎 commMode 选型以及完整的多卡调用示例并结合仓库源码与单元测试深入剖析其入参校验、通信引擎传递与确定性计算等实现细节帮助开发者正确、高效地在 NPU 上使用该通算融合算子。算子定位为什么需要 MatMul AllReduce 融合在 Transformer 类大模型训练与推理中MatMul矩阵乘与AllReduce全量规约通信是出现频率最高的两类算子。传统实现将两者作为独立算子先后执行先由计算核完成MatMul将中间结果落盘到 Global Memory再启动通信引擎把各卡结果做AllReduce归约。这带来了两次额外的访存开销与一次算子调度开销在大规模并行如专家并行、张量并行场景下通信与计算的重叠效率直接决定集群扩展性。aclnnMatmulAllReduceV3将两者融合为单一算子在计算核完成矩阵乘与逐元素加法的同时把结果通过通信引擎直接完成跨卡归约减少中间张量搬运与算子下发次数。该接口兼容aclnnMatmulAllReduceV2的全部功能并新增通信引擎参数commMode允许调用方在 AICPU 与 CCU 两种通信引擎之间自选不同芯片支持情况不同从而在不同数据量、不同组网规模下获得更优的通信性能。产品支持情况aclnnMatmulAllReduceV3在 CANN ops-transformer 当前仓库中的产品支持矩阵如下产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持版本前提使用该接口时请确保驱动固件包和 CANN 包均为配套的 8.0.RC2 版本或配套的更高版本否则会引发报错如 Bus Error 等。该信息与仓库 READMEmc2/matmul_all_reduce/README.md中的产品支持矩阵一致且 README 进一步说明A2 系列支持 1/2/4/8 卡950 系列支持 1/2/4/8/16/32/64 卡且仅支持 HCCS 链路 all mesh 组网。功能说明与计算公式接口功能完成 MatMul 计算与 AllReduce 通信融合兼容aclnnMatmulAllReduceV2支持的功能新增通信引擎参数commMode支持自选通信引擎。计算公式$$ output AllReduce(x1 x2 bias x3) $$即先计算x1与x2的矩阵乘叠加一维偏置bias与同 shape 的x3再在指定通信域group上按reduceOp当前仅支持 sum做一次全量归约得到最终输出output。需要说明的是bias、x3 在aclnnMatmulAllReduceV3中均可选输入可为空与仓库 README.md 中非量化场景情形2的公式保持一致。两段式接口与函数原型与 CANN AscendCL 其它算子一样aclnnMatmulAllReduceV3采用两段式接口先调用aclnnMatmulAllReduceV3GetWorkspaceSize完成入参校验获取计算所需 workspace 大小以及包含算子计算流程的执行器executor再调用aclnnMatmulAllReduceV3真正执行计算。两段式接口的函数原型如下aclnnStatus aclnnMatmulAllReduceV3GetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, const aclTensor *bias, const aclTensor *x3, const char *group, const char *reduceOp, const char *commMode, int64_t commTurn, int64_t streamMode, const aclTensor *output, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnMatmulAllReduceV3( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)该原型定义可在仓库头文件 mc2/matmul_all_reduce/op_api/aclnn_matmul_all_reduce_v3.h 中查到头文件中domain aclnn_ops_infer标注了该接口面向推理场景的域归属。aclnnMatmulAllReduceV3GetWorkspaceSize 参数详解第一阶段接口共 12 个参数完整说明如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1输入MatMul 计算的左矩阵即计算公式中的 x1当前版本仅支持二维或三维输入支持不转置场景FLOAT16、BFLOAT16ND2-3×x2输入MatMul 计算的右矩阵即计算公式中的 x2当前版本仅支持二维输入支持转置/不转置场景ND 格式下支持最后两轴转置情况下的非连续 tensor其他非连续 tensor 不支持FLOAT16、BFLOAT16ND2×x3输入MatMul 计算后的 add 计算即计算公式中的 x3shape 与 MatMul 计算后的 shape 一致FLOAT16、BFLOAT16ND2√bias输入对应计算公式中的 bias 偏移当前版本仅支持一维输入FLOAT16、BFLOAT16ND0-1√group输入通信域名称通过 Hccl 提供的接口extern HcclResult HcclGetCommName(HcclComm comm, char* commName);获取其中 commName 即为 groupString---reduceOp输入reduce 操作类型当前版本仅支持输入 sumString---commMode输入通信引擎参数用于指定通信引擎根据芯片型号有不同的输入限制详见下文约束说明String---commTurn输入通信数据切分数即总数据量/单次通信量当前版本仅支持输入 0INT64---streamMode输入AscendCL 流模式的枚举当前版本仅支持枚举值 1INT64---output输出MatMul 计算与 AllReduce 通信的结果即计算公式中的 outputoutput 的维数与 x1 一致FLOAT16、BFLOAT16ND2-3√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----几个需要特别留意的要点x2 的转置与非连续接口支持 x2 转置由内部IsTransposeLastTwoDims依据 stride 判断且仅在最后两轴转置场景下才允许 x2 非连续其余非连续 x2 一律返回参数非法。源码实现见 mc2/matmul_all_reduce/op_api/matmul_all_reduce_util.cpp 中MatmulAllReduceCheckValidContiguous。x1 不支持转置op_api 实现中transposeX1固定为false见 mc2/matmul_all_reduce/op_api/aclnn_matmul_all_reduce_v3.cpp与文档约束一致。bias 的维度文档标注 shape 为 0-1 维实际仅支持一维输入且一维长度须等于 output 最后一维即 n0 维标量场景对应空 bias 的退化情形。aclnnMatmulAllReduceV3 执行接口参数详解第二阶段接口共 4 个参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnMatmulAllReduceV3GetWorkspaceSize获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值与入参校验错误码两段式接口均返回aclnnStatus状态码完整错误码说明参见 aclnn 返回码。第一阶段接口会完成入参校验以下场景将报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 x1、x2 或 output 是空指针ACLNN_ERR_PARAM_INVALID161002x1、x2、bias、x3 或 output 的数据类型不在支持范围之内ACLNN_ERR_PARAM_INVALID161002reduceOp、streamMode 不在合法范围内ACLNN_ERR_PARAM_INVALID161002x1、x2、bias、x3 或 output 的 shape 不符合约束要求从源码看第一段接口的校验顺序非常清晰mc2/matmul_all_reduce/op_api/matmul_all_reduce_util.cpp 中MatmulAllReduceCheckParams空指针检查MatmulAllReduceCheckNotNull(x1, x2, output)数据类型检查MatmulAllReduceCheckDtypeValid要求 x1、x2、x3、bias、output 的类型都在 FLOAT16/BFLOAT16 支持列表内且 x1、x2、x3、bias、output 彼此数据类型一致属性检查MatmulAllReduceCheckAttrreduceOp必须等于sumstrcmp与REDUCE_OP_SUM比对streamMode必须等于NUM_ACL_STOP_ON_FAILURE即 1格式检查非 310P 场景下 x2 必须为 ND 格式shape 检查MatmulAllReduceCheckShapex2 必须 2 维x1 维数在 2~3 之间x1 最后一维k必须等于 x2 第一维koutput shape 等于 x1 前序维拼接 x2 最后一维bias 必须一维且长度为 nx3 与 output shape 一致连续性检查A2ASCEND910B场景下x2 非连续时必须为转置场景否则报ACLNN_ERR_PARAM_INVALID。此外aclnnMatmulAllReduceV3在入参前会额外校验commMode合法性IsCommModeValid非法值直接返回ACLNN_ERR_PARAM_INVALID。约束说明通信引擎 commMode 支持度commMode是本接口相对 V2 的核心新增参数用于自选通信引擎Atlas A2 训练/推理系列产品目前不支持指定通信引擎commMode仅允许输入ai_cpu即使用 AICPU 通信引擎Ascend 950PR / Ascend 950DT通信引擎支持 AICPU 与 CCUcommMode允许输入ai_cpu或ccu。其中 CCU 仅支持单机 UB 域内互联AICPU 可支持跨机 UB 域内互联同一条通信链路内只能选择同一种通信引擎。两种通信引擎的特点如下AICPU不占用计算核通信效率高但通信静态开销较大对小数据量通信场景不友好适用于大数据高带宽场景CCU能够减少访存带宽与计算核占用但受限于片上资源支持的通信域数量有限适用于高带宽、低时延的通信场景。从源码可以看到该参数在运行时如何生效mc2/matmul_all_reduce/op_api/aclnn_matmul_all_reduce_v3.cppGetWorkspaceSize阶段将字符串commMode映射为枚举值ai_cpu→Mc2Comm::COMM_MODE_AICPUccu→Mc2Comm::COMM_MODE_CCU通过NnopbaseSetUserHandle挂到 executor 上若既不是ai_cpu也不是ccu直接返回ACLNN_ERR_PARAM_INVALID执行阶段aclnnMatmulAllReduceV3依据当前 NPU 架构决定 HCCL server 类型——在 950源码中DAV_3510上按 handle 中保存的 commMode 设置NNOPBASE_HCCL_SERVER_TYPE_CCU或NNOPBASE_HCCL_SERVER_TYPE_AICPU在 A2 上则固定使用 AICPU。单元测试 CSVmc2/matmul_all_reduce/tests/ut/op_api/test_aclnn_matmul_all_reduce_v3.csv也验证了这一约束A2 上传入ccu返回PARAM_INVALID950 上ccu与ai_cpu均 SUCCESS而非法字符串invalid返回PARAM_INVALID。确定性计算Atlas A2 训练/推理系列产品aclnnMatmulAllReduceV3默认非确定性实现支持通过配置HCCL_DETERMINISTIC环境变量为true开启确定性计算Ascend 950PR / Ascend 950DTaclnnMatmulAllReduceV3默认确定性实现。确定性计算的通用背景可参考仓库文档 确定性计算说明。shape、数据类型与组网约束增量场景不使能 MC2全量场景使能 MC2输入 x1 可为二维或三维其 shape 为(b, s, k)或(m, k)x2 必须为二维shape 为(k, n)轴满足 mm 算子入参要求k 轴相等bias 若非空其 shape 为(n)b*s、m、k、n的值均不得超过 2147483647INT32_MAX当输入 x1 的 shape 为(b, s, k)时输出 output 的 shape 为(b, s, n)当输入 x1 的 shape 为(m, k)时输出 output 的 shape 为(m, n)x1、x2、bias 计算输入的数据类型要和 output 计算输出的数据类型一致传入的 x1、x2、x3 或 output 不为空指针仅支持 HCCS 链路 all mesh 组网Atlas A2 训练/推理系列产品支持 1、2、4、8 卡Ascend 950PR / Ascend 950DT支持 1、2、4、8、16、32、64 卡Atlas A2 训练/推理系列产品一个模型中的通算融合 MC2 算子仅支持相同通信域。空 tensor 支持度支持 k 为 0 的场景此时输出为bias x3支持bs/m/n为 0此时传入的输出也应该是空 tensor此场景不进入 kernel 计算直接返回。上述空 tensor 行为在测试用例中同样有覆盖CSV 中empty_K、empty_M、empty_N三个用例在 A2 上均返回 SUCCESS。完整调用示例以下为文档提供的完整多卡调用示例展示了从设备初始化、HCCL 通信域创建、tensor 创建、两段式接口调用到资源释放的完整流程。示例基于 FP16、x1(32,64) x2(64,128) bias(128) x3(32,128)的 2 卡场景。编译与执行过程请参考 编译与运行样例。说明示例代码调用了部分 HCCL 集合通信库接口HcclGetCommName、HcclCommInitAll、HcclCommDestroy这些接口来自 HCCL APIC头文件hccl/hccl.h随 CANN 环境一并提供。#include iostream #include vector #include thread #include hccl/hccl.h #include aclnn/opdev/fp16_t.h #include aclnnop/aclnn_matmul_all_reduce_v3.h int ndev 2; #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i: shape) { shapeSize * i; } return shapeSize; } templatetypename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void **deviceAddr, aclDataType dataType, aclTensor **tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } struct Args { uint32_t rankId; HcclComm hcclComm; aclrtStream stream; aclrtContext context; }; int launchOneThreadMatmulAllReduce(Args args) { int ret; ret aclrtSetCurrentContext(args.context); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetCurrentContext failed. ERROR: %d\n, ret); return ret); char hcom_name[128]; ret HcclGetCommName(args.hcclComm, hcom_name); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT([ERROR] HcclGetCommName failed. ret %d \n, ret); return -1); LOG_PRINT([INFO] rank %d hcom: %s stream: %p, context : %p\n, args.rankId, hcom_name, args.stream, args.context); std::vectorint64_t x1Shape {32, 64}; std::vectorint64_t x2Shape {64, 128}; std::vectorint64_t biasShape {128}; std::vectorint64_t x3Shape {32, 128}; std::vectorint64_t outShape {32, 128}; void *x1DeviceAddr nullptr; void *x2DeviceAddr nullptr; void *biasDeviceAddr nullptr; void *x3DeviceAddr nullptr; void *outDeviceAddr nullptr; aclTensor *x1 nullptr; aclTensor *x2 nullptr; aclTensor *bias nullptr; aclTensor *x3 nullptr; aclTensor *out nullptr; int64_t commTurn 0; int64_t streamMode 1; uint64_t workspaceSize 0; aclOpExecutor *executor; void *workspaceAddr nullptr; long long x1ShapeSize GetShapeSize(x1Shape); long long x2ShapeSize GetShapeSize(x2Shape); long long biasShapeSize GetShapeSize(biasShape); long long x3ShapeSize GetShapeSize(x3Shape); long long outShapeSize GetShapeSize(outShape); std::vectorop::fp16_t x1HostData(x1ShapeSize, 1); std::vectorop::fp16_t x2HostData(x2ShapeSize, 1); std::vectorop::fp16_t biasHostData(biasShapeSize, 1); std::vectorop::fp16_t x3HostData(x3ShapeSize, 1); std::vectorop::fp16_t outHostData(outShapeSize, 0); // 创建tensor ret CreateAclTensor(x1HostData, x1Shape, x1DeviceAddr, aclDataType::ACL_FLOAT16, x1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(x2HostData, x2Shape, x2DeviceAddr, aclDataType::ACL_FLOAT16, x2); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(biasHostData, biasShape, biasDeviceAddr, aclDataType::ACL_FLOAT16, bias); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(x3HostData, x3Shape, x3DeviceAddr, aclDataType::ACL_FLOAT16, x3); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT16, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 调用第一段接口 ret aclnnMatmulAllReduceV3GetWorkspaceSize(x1, x2, bias, x3, hcom_name, sum, , commTurn, streamMode, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMatmulAllReduceV3GetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用第二段接口 ret aclnnMatmulAllReduceV3(workspaceAddr, workspaceSize, executor, args.stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMatmulAllReduceV3 failed. ERROR: %d\n, ret); return ret); //固定写法同步等待任务执行结束 ret aclrtSynchronizeStreamWithTimeout(args.stream, 10000); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); LOG_PRINT(device%d aclnnMatmulAllReduceV3 execute success \n, args.rankId); // 释放device资源需要根据具体API的接口定义修改 if (x1 ! nullptr) { aclDestroyTensor(x1); } if (x2 ! nullptr) { aclDestroyTensor(x2); } if (bias ! nullptr) { aclDestroyTensor(bias); } if (x3 ! nullptr) { aclDestroyTensor(x3); } if (out ! nullptr) { aclDestroyTensor(out); } if (x1DeviceAddr ! nullptr) { aclrtFree(x1DeviceAddr); } if (x2DeviceAddr ! nullptr) { aclrtFree(x2DeviceAddr); } if (biasDeviceAddr ! nullptr) { aclrtFree(biasDeviceAddr); } if (x3DeviceAddr ! nullptr) { aclrtFree(x3DeviceAddr); } if (outDeviceAddr ! nullptr) { aclrtFree(outDeviceAddr); } if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(args.stream); HcclCommDestroy(args.hcclComm); aclrtDestroyContext(args.context); aclrtResetDevice(args.rankId); return 0; } int main(int argc, char *argv[]) { int ret; int32_t devices[ndev]; for (int i 0; i ndev; i) { devices[i] i; } HcclComm comms[128]; ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); // 初始化集合通信域 for (int i 0; i ndev; i) { ret aclrtSetDevice(devices[i]); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); } ret HcclCommInitAll(ndev, devices, comms); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(HcclCommInitAll failed. ERROR: %d\n, ret); return ret); Args args[ndev]; aclrtStream stream[ndev]; aclrtContext context[ndev]; for (uint32_t rankId 0; rankId ndev; rankId) { ret aclrtSetDevice(rankId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateContext(context[rankId], rankId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateContext failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream[rankId]); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); } // 启动多线程 std::vectorstd::unique_ptrstd::thread threads(ndev); for (uint32_t rankId 0; rankId ndev; rankId) { args[rankId].rankId rankId; args[rankId].hcclComm comms[rankId]; args[rankId].stream stream[rankId]; args[rankId].context context[rankId]; threads[rankId].reset(new(std::nothrow) std::thread(launchOneThreadMatmulAllReduce, std::ref(args[rankId]))); } for (uint32_t rankId 0; rankId ndev; rankId) { threads[rankId]-join(); } aclFinalize(); return 0; }调用流程的关键节点拆解如下通信域准备aclInit后通过aclrtSetDevice绑定多卡用HcclCommInitAll一次性创建多卡的 HCCL 通信域每个 rank 的线程通过HcclGetCommName拿到通信域名字符串作为group参数tensor 创建CreateAclTensor内部先aclrtMalloc申请 Device 内存、aclrtMemcpy拷入 host 数据再按连续布局计算 strides 并调用aclCreateTensor创建aclTensor第一段接口aclnnMatmulAllReduceV3GetWorkspaceSize完成校验并返回workspaceSize与executor若workspaceSize 0则aclrtMalloc申请 workspace第二段接口aclnnMatmulAllReduceV3提交任务到指定 stream随后用aclrtSynchronizeStreamWithTimeout同步等待执行完成资源回收依次销毁aclTensor、释放 Device 内存与 workspace、销毁 stream、销毁 HCCL 通信域与 context最后aclrtResetDevice与aclFinalize。示例中第一段接口的commMode传入了空字符串表示由系统按当前芯片选择默认通信引擎A2 固定为ai_cpu950 默认引擎实际开发中可按上文约束说明显式传入ai_cpu或ccu。从源码看实现原理op_api 层参数校验与 commMode 传递aclnnMatmulAllReduceV3的 op_api 实现位于 mc2/matmul_all_reduce/op_api/aclnn_matmul_all_reduce_v3.cpp。核心逻辑分三步入参校验aclnnMatmulAllReduceV3GetWorkspaceSize先调用MatmulAllReduceCheckParams完成空指针、dtype、attrreduceOp/streamMode、格式、shape、连续性检查再调用IsCommModeValid校验 commMode 字符串合法性下调内核接口内部将 x1 固定为不转置transposeX1 false通过IsTransposeLastTwoDims(x2)自动识别 x2 是否转置然后调用aclnnInnerMatmulAllReduceGetWorkspaceSize下发 IR该内核接口同时承载量化分支因此 V3 在此通过NnopbaseDisableOptionalInput将 4~9 号量化相关可选输入scale/offset/dequant/pertoken/commQuantScale 等全部禁用commMode 传递commMode字符串被映射为枚举后经NnopbaseSetUserHandle挂到 executor 上执行阶段aclnnMatmulAllReduceV3再按 950/非 950 分别设置NNOPBASE_HCCL_SERVER_TYPE_CCU或NNOPBASE_HCCL_SERVER_TYPE_AICPU最终调用aclnnInnerMatmulAllReduce真正下发执行。公共的校验工具函数集中在 mc2/matmul_all_reduce/op_api/matmul_all_reduce_util.h 与 mc2/matmul_all_reduce/op_api/matmul_all_reduce_util.cpp 中包括MatmulAllReduceCheckParams、MatmulAllReduceCheckNotNull、MatmulAllReduceCheckDtypeValid、MatmulAllReduceCheckAttr、MatmulAllReduceCheckShape、MatmulAllReduceCheckValidEmptyTensor、MatmulAllReduceCheckValidContiguous等V2/V3 及量化系列接口共用。算子定义层动态特性配置算子的 IR 定义在 mc2/matmul_all_reduce/op_host/matmul_all_reduce_def.cpp。MatmulAllReduceOpDef 为输出y配置了 950 全系列的 FLOAT16/BFLOAT16/FLOAT 支持列表与 ND 格式并为 AICore 配置了动态编译、动态格式、动态 rank 与动态 shape 支持意味着该算子支持动态 shape 下运行时重新做 tiling 切分。值得说明的是该 OpDef 同时承载量化分支README 中的情形 3~7因此 dtype/format 列表比 V3 文档限定的 FLOAT16/BFLOAT16 更宽V3 接口本身仍以文档约束为准。从测试用例验证行为单元测试 mc2/matmul_all_reduce/tests/ut/op_api/test_aclnn_matmul_all_reduce_v3.cpp 通过 CSV 参数化驱动测试用例数据见 mc2/matmul_all_reduce/tests/ut/op_api/test_aclnn_matmul_all_reduce_v3.csv覆盖了合法场景2D/3D x1、BF16、empty_K/empty_M/empty_N空 tensor、950 上的ccu与ai_cpu通信引擎非法场景x1 四维、x1 空、x2 三维、x3 三维、reduceOp 为max、streamMode 为 0/-1、dtype 不匹配、k/n 轴不匹配、bias/x3 与 output 不匹配、x2 为 FRACTAL_NZ 格式、A2 上传ccu、非法 commMode 字符串等均返回ACLNN_ERR_PARAM_INVALID额外用例转置 x2非连续但满足最后两轴转置SUCCESS非转置非连续 x2 返回PARAM_INVALIDexecutor为空时第二段接口返回ACLNN_ERR_INNER。接口家族与版本演进aclnnMatmulAllReduceV3属于matmul_all_reduce算子家族该模块 docs 目录下共存有非量化、量化、权重量化三个方向的接口文档可交叉参考非量化aclnnMatmulAllReduce、aclnnMatmulAllReduceV2、aclnnMatmulAllReduceV3本文量化INT8/FP8 等aclnnQuantMatmulAllReduce至aclnnQuantMatmulAllReduceV5权重量化aclnnWeightQuantMatmulAllReduce、aclnnWeightQuantMatmulAllReduceV2。V3 相对 V2 的核心差异在于新增commMode通信引擎参数允许 950 系列在 AICPU 与 CCU 之间自选为高带宽、低时延场景提供了更细粒度的通信调优手段计算语义上两者一致均实现output AllReduce(x1 x2 bias x3)。相关文档入口mc2/matmul_all_reduce/docs/aclnnMatmulAllReduceV2.md 与 mc2/matmul_all_reduce/README.md。总结aclnnMatmulAllReduceV3是 CANN ops-transformer 中面向 Transformer 大模型全量并行训练/推理场景的典型通算融合算子通过将矩阵乘、偏置加法、x3累加与跨卡 AllReduce 合入一次下发有效压缩中间访存与调度开销。使用该接口时需重点把握两段式调用流程、FLOAT16/BFLOAT16 与 ND 格式约束、sum/1/0 三个属性的取值限制、x2 转置场景下的非连续支持、以及commMode随芯片的差异化选型A2 仅ai_cpu950 可选ai_cpu/ccu。结合仓库源码与单元测试开发者可以据此快速定位参数校验失败原因并在多卡组网上正确搭建通算融合的计算流程。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考