CANN ops-transformer 通信算子解析：aclnnMoeDistributeCombineTeardown 两段式接口详解

CANN ops-transformer 通信算子解析aclnnMoeDistributeCombineTeardown 两段式接口详解【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformeraclnnMoeDistributeCombineTeardown是 CANN ops-transformer 中 MoEMixture of Experts分布式专家并行EP通信链路里的收尾算子负责对aclnnMoeDistributeCombineSetup收集回来的跨卡 token 数据做通信状态确认并按专家权重做加权求和最终恢复出本卡的 batch 输出。读完本文你将掌握该算子的功能定位、两段式调用方式、全部入参约束与取值规则并能够结合仓库源码写出可在 Ascend 950DT 上运行的直调示例。一、功能定位MoE 分布式 Combine 链路的最后一步MoE 分布式推理/训练中token 会按专家路由结果被分发Dispatch到不同的 EP 卡上执行专家计算之后再把结果收集Combine回原卡。在 CANN ops-transformer 中这一过程由四个接口配套完成aclnnMoeDistributeDispatchSetup分发前处理生成 dispatch 所需的通信命令与扩展索引aclnnMoeDistributeDispatchTeardown分发收尾aclnnMoeDistributeCombineSetup收集前处理生成 combine 所需的通信命令commCmdInfo与量化扩展数据quantExpandXaclnnMoeDistributeCombineTeardown收集收尾确认通信状态并对数据做加权求和。本接口功能接收aclnnMoeDistributeCombineSetup发来的数据并整合乘权重再相加。其计算公式为$$ ataOut AllToAllV(expandX) $$$$ xOut Sum(expertScales * ataOut expertScales * sharedExpertX) $$按MoeDistributeDispatchSetup和MoeDistributeDispatchTeardown算子收集数据的路径原路返还本算子只做通信状态确认和数据整理对aclnnMoeDistributeCombineSetup接口发送的数据做通信状态确认并按照专家的 scale 权重做加权求和。注意该接口必须与aclnnMoeDistributeDispatchSetup、aclnnMoeDistributeDispatchTeardown及aclnnMoeDistributeCombineSetup配套使用。从算子定义源码 op_host/moe_distribute_combine_teardown_def.cpp 可以看到该算子在信息库中注册了ascend950对应 Ascend 950DT的 AICore 配置并声明MC2().HcclGroup(group_ep)——即以 EP 通信域group_ep为通信组这正是其参与跨卡集合通信的直接证据。二、两段式接口与函数原型每个算子分为两段式接口必须先调用aclnnMoeDistributeCombineTeardownGetWorkspaceSize接口获取入参并根据计算流程计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnMoeDistributeCombineTeardown接口执行计算。第一段接口原型aclnnStatus aclnnMoeDistributeCombineTeardownGetWorkspaceSize( const aclTensor* expandX, const aclTensor* quantExpandX, const aclTensor* expertIds, const aclTensor* expandIdx, const aclTensor* expertScales, const aclTensor* commCmdInfo, const aclTensor* xActiveMaskOptional, const aclTensor* sharedExpertXOptional, const char* groupEp, int64_t epWorldSize, int64_t epRankId, int64_t moeExpertNum, int64_t expertShardType, int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, int64_t commQuantMode, int64_t commType, const char* commAlg, aclTensor* xOut, uint64_t* workspaceSize, aclOpExecutor** executor)第二段接口原型aclnnStatus aclnnMoeDistributeCombineTeardown( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)调用流程为先用第一段接口完成入参校验并计算出 workspace 大小、生成executor随后在 Device 侧按workspaceSize申请内存示例中通过aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST)申请最后调用第二段接口在指定 stream 上执行计算。在接口实现源码 op_api/aclnn_moe_distribute_combine_teardown.cpp 中可以观察到几个实现细节第一段接口会先校验当前 NPU 架构必须为Ops::Base::DAV_3510即 Ascend 950DT 的昇腾 AI Core 架构否则直接返回ACLNN_ERR_PARAM_INVALID随后执行CheckParams对 6 个必选输入、groupEp与xOut做空指针检查并对groupEp名称长度HCCL_GROUP_NAME_MAX做越界校验第二段接口会通过NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE)将通信服务类型设置为 MTE 通路这与文档中当前仅支持 commType2URMA 通路的约束相互印证。三、aclnnMoeDistributeCombineTeardownGetWorkspaceSize 参数详解3.1 Tensor 类参数下表完整列出第一段接口的 Tensor 参数要求参数名输入/输出描述使用说明数据类型数据格式维度非连续TensorexpandXaclTensor*输入根据expertIds进行扩展过的token特征不支持空Tensor。FLOAT16、BFLOAT16ND(A, H)√quantExpandXaclTensor*输入对应aclnnMoeDistributeCombineSetup中的quantExpandX输出不支持空Tensor。INT8ND(A, tokenMsgSize)√expertIdsaclTensor*输入每个token的topK个专家索引不支持空Tensor。INT32ND(Bs, K)√expandIdxaclTensor*输入对应aclnnMoeDistributeDispatchSetup中的expandIdx输出不支持空Tensor。INT32ND(BS * K, )√expertScalesaclTensor*输入每个token的topK个专家的权重不支持空Tensor。FLOAT32ND(BS, K)√commCmdInfoaclTensor*输入aclnnMoeDistributeCombineSetup的输出通信的cmd信息不支持空Tensor。INT32ND((A epWorldSize) * 16, )√xActiveMaskOptionalaclTensor*可选输入表示token是否参与通信不支持空Tensor。可选择传入有效数据或填空指针传空指针时表示所有token都参与通信。BOOLND(BS, )√sharedExpertXOptionalaclTensor*可选输入表示共享专家计算后的Token不支持空Tensor。要求是一个2D或3D的Tensor当Tensor为2D时shape为(BS,H)当Tensor为3D时前两位的乘积需等于BS第三维需等于H。数据类型需跟expandX保持一致。可选择传入有效数据或填空指针传入有效数据时sharedExpertNum需为0。FLOAT16、BFLOAT16ND(BS, H)或(a, b, H)√xOutaclTensor*输出表示处理后的token不支持空Tensor。数据类型与expandX保持一致。FLOAT16、BFLOAT16ND(BS, H)√workspaceSize输出返回需要在Device侧申请的workspace大小。-----executor输出返回op执行器包含了算子计算流程。-----这些输入/输出的类型约束与算子定义文件 op_host/moe_distribute_combine_teardown_def.cpp 完全对应expand_x/shared_expert_x/x支持DT_FLOAT16与DT_BF16quant_expand_x为DT_INT8expert_ids/expand_idx/comm_cmd_info为DT_INT32expert_scales为DT_FLOATfloat32x_active_mask为DT_BOOL格式均为FORMAT_ND且所有输入均声明了AutoContiguous()。在 shape 推断源码 op_host/moe_distribute_combine_teardown_infershape.cpp 中可以看到输出xOut的 shape 推导逻辑其第 0 维取自expertIds的第 0 维即 BS第 1 维取自expandX的第 1 维即 H输出数据类型与expandX保持一致——这印证了输出 (BS, H)、数据类型与 expandX 一致的约定。3.2 标量与字符串参数参数名输入描述使用说明groupEp输入EP通信域名称字符串长度范围为[1, 128)epWorldSize输入EP通信域size取值支持[2, 384]epRankId输入EP域本卡Id取值范围[0, epWorldSize)。同一个EP通信域中各卡的epRankId不重复。moeExpertNum输入MoE专家数量取值范围(0, 512]。满足moeExpertNum % (epWorldSize - sharedExpertRankNum) 0。expertShardType输入共享专家卡分布类型当前仅支持传0表示共享专家卡排在MoE专家卡前面。sharedExpertNum输入共享专家数量当前取值范围[0, 4]。0表示无共享专家。当前仅支持无共享专家。sharedExpertRankNum输入共享专家卡数量当前取值范围[0, epWorldSize / 2]。globalBs输入EP域全局的batch size大小当每个rank的Bs数一致场景下globalBs Bs * epWorldSize或globalBs 0当每个rank的Bs数不一致场景下globalBs maxBs * epWorldSize其中maxBs表示单卡Bs最大值。commQuantMode输入通信量化类型取值范围[0, 2]0表示通信时不进行量化当前仅支持0。commType输入通信方案选择取值范围[0, 2]0表示AICPU-SDMA方案1表示CCU方案2表示URMA方案当前仅支持2。commAlg输入通信算法选择仅支持传入空指针或空字符串图模式原型 op_graph/moe_distribute_combine_teardown_proto.h 给出了这些属性的默认值expert_shard_type默认 0、shared_expert_num默认 1、shared_expert_rank_num默认 0、global_bs默认 0、comm_quant_mode默认 0、comm_type默认 0、comm_alg默认空字符串其中group_ep、ep_world_size、ep_rank_id、moe_expert_num为必选属性。3.3 平台相关的参数限制Ascend 950DT不支持共享专家场景。epWorldSize 当前取值仅支持 2、8。moeExpertNum 表示 MoE 专家数量当前仅能传入 32。expertShardType 当前仅支持传 0表示共享专家卡排在 MoE 专家卡前面。sharedExpertNum 表示共享专家数量当前不支持共享专家仅能传入 0。sharedExpertRankNum 表示共享专家卡数当前不支持共享专家仅能传入 0。commQuantMode 当前仅支持传入 0表示不进行量化。commType 取值范围 [0, 2]当前仅支持 2表示 URMA 通路。commAlg 当前版本不支持传空指针即可。Atlas A3 训练系列产品/Atlas A3 推理系列产品限制项与 Ascend 950DT 完全一致epWorldSize 仅支持 2、8moeExpertNum 仅能传入 32不支持共享专家commType 仅支持 2commAlg 传空指针。四、aclnnMoeDistributeCombineTeardown 参数说明参数名输入/输出描述workspace输入在Device侧申请的workspace内存地址。workspaceSize输入在Device侧申请的workspace大小由第一段接口aclnnMoeDistributeCombineTeardownGetWorkspaceSize获取。executor输入op执行器包含了算子计算流程。stream输入指定执行任务的stream流。五、返回值与错误码两段接口均返回aclnnStatus状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入和输出的必选参数Tensor是空指针。ACLNN_ERR_PARAM_INVALID161002输入和输出的数据类型不在支持的范围内。ACLNN_ERR_INNER_TILING_ERROR561002输入和输出的shape不在支持的范围内。ACLNN_ERR_INNER_TILING_ERROR561002参数的取值不在支持的范围内。结合接口实现代码可以看到第一段接口在正式进入 tiling 之前就会执行CheckParams校验并返回ACLNN_ERR_PARAM_NULLPTR或ACLNN_ERR_PARAM_INVALID而 shape、参数取值相关的校验则在 host 侧 tiling 阶段完成对应ACLNN_ERR_INNER_TILING_ERROR561002——tiling 基类 op_host/op_tiling/moe_distribute_combine_teardown_tiling_base.h 中定义了CheckAttrs、CheckTensorDim、CheckTensorShapeRelation、CheckTensorDataType、CheckTensorFormat、CheckHcclBuffsize等一系列校验步骤最终由 op_host/op_tiling/moe_distribute_combine_teardown_tiling.cpp 注册的MoeDistributeCombineTeardownTilingFunc驱动执行。六、约束说明6.1 确定性计算aclnnMoeDistributeCombineTeardown 默认确定性实现。6.2 配套使用约束aclnnMoeDistributeDispatchSetup、aclnnMoeDistributeDispatchTeardown、aclnnMoeDistributeCombineSetup、aclnnMoeDistributeCombineTeardown四个接口必须配套使用。调用接口过程中使用的groupEp、epWorldSize、moeExpertNum、expertShardType、sharedExpertNum、sharedExpertRankNum、globalBs、commQuantMode、commType、commAlg参数取值所有卡需保持一致其中groupEp、epWorldSize、expertShardType、sharedExpertNum、sharedExpertRankNum、globalBs、commQuantMode、commType、commAlg参数取值在网络中不同层中也需保持一致且与 dispatch/combine 各阶段接口的对应参数保持一致。6.3 shape 符号说明参数说明中的 shape 符号含义如下A本卡需要分发的最大 token 数量取值范围如下对于共享专家要满足 A BS * epWorldSize * sharedExpertNum / sharedExpertRankNum当 globalBs 非 0 时要满足 A globalBs * sharedExpertNum / sharedExpertRankNum。对于 MoE 专家当 globalBs 为 0 时要满足 A BS * epWorldSize * min(localExpertNum, K)当 globalBs 非 0 时要满足 A globalBs * min(localExpertNum, K)。Hhidden size 隐藏层大小。取值为 [1024, 8196]。当前仅支持 4096、7168。HS通信时的数据大小Align512( Align32(H) Align8(H) / 8 * sizeof(float) )其中 Align512(x) ((x 512 - 1) / 512) * 512Align32(x) ((x 32 - 1) / 32) * 32Align8(x) ((x 8 - 1) / 8) * 8。BSbatch sequence size即本卡最终输出的 token 数量。取值范围为 0 BS ≤ 512。当前仅支持 8、16、256。K选取 topK 个专家取值范围为 0 K ≤ 16 同时满足 0 K ≤ moeExpertNum。当前仅支持 6、8。localExpertNum本卡专家数量。对于共享专家卡localExpertNum 1对于 MoE 专家卡localExpertNum moeExpertNum / (epWorldSize - sharedExpertRankNum)。moeExpertNum 当前仅支持 32。tokenMsgSize每个 token 在数据通信时的维度信息计算公式为 Align512(Align32(H) Align8(H) / 8 * sizeof(float))其中 AlignN(x) ((x N - 1) / N * N)。当前不支持共享专家。sharedExpertNum 和 sharedExpertRankNum 当前仅支持 0。6.4 HCCL_BUFFSIZE 环境变量调用本接口前需检查HCCL_BUFFSIZE环境变量取值是否合理该环境变量表示单个通信域占用内存大小单位 MB不配置时默认为 200MB。Ascend 950DT要求 2 且满足 4 * (localExpertNum * maxBs * epWorldSize * Align512(Align32(2 * H) 44) (K sharedExpertNum) * maxBs * Align512(2 * H))localExpertNum 需使用 MoE 专家卡的本卡专家数。Atlas A3 训练系列产品/Atlas A3 推理系列产品要求 2 且满足 2 * (localExpertNum * maxBs * epWorldSize * Align512(Align32(2 * H) 44) (K sharedExpertNum) * maxBs * Align512(2 * H))localExpertNum 需使用 MoE 专家卡的本卡专家数。Align512(x) ((x 512 - 1) / 512) * 512Align32(x) ((x 32 - 1) / 32) * 32。6.5 通信域与通信方式约束一个模型中的aclnnMoeDistributeDispatchSetup、aclnnMoeDistributeDispatchTeardown、aclnnMoeDistributeCombineSetup、aclnnMoeDistributeCombineTeardown接口仅支持相同 EP 通信域且该通信域中不允许有其他算子。Ascend 950DT仅支持 URMA 通信与 commType 仅支持 2 的约束一致。七、调用示例Ascend 950DT 上aclnnMoeDistributeCombineTeardown必须在aclnnMoeDistributeCombineSetup调用之后调用。仓库提供了可直接运行的完整示例examples/test_aclnn_moe_distribute_combine_teardown.cpp配套的单算子直调测试位于 tests/ut/op_api/test_aclnn_moe_distribute_combine_teardown.cpp其测试用例覆盖了不同 shape、可选输入传入/置空等参数组合并在SetUpTestCase中显式将平台架构设置为Ops::Base::DAV_3510。7.1 场景参数配置示例中的典型场景参数与文档约束完全吻合int64_t bs 16; // BS文档要求 0 BS ≤ 512当前支持 8、16、256 int64_t h 4096; // H文档要求当前仅支持 4096、7168 int64_t k 6; // K文档要求当前仅支持 6、8 int64_t expertSharedType 0; // expertShardType仅支持 0 int64_t sharedExpertNum 0; // 当前不支持共享专家仅支持 0 int64_t sharedExpertRankNum 0; int64_t moeExpertNum 32; // 当前仅支持 32 int64_t commQuantMode 0; // 不进行量化 int64_t epWorldSize DEV_NUM; // 2 卡 EP 域 int64_t commType 2; // URMA 通路 int64_t globalBS bs * epWorldSize; // 各 rank Bs 一致场景globalBs Bs * epWorldSize7.2 shape 推导与 tensor 创建int64_t localExpertNum; int64_t localToken; if (args.epRankId sharedExpertRankNum) { localExpertNum 1; localToken globalBS / sharedExpertRankNum; } else { localExpertNum moeExpertNum / (epWorldSize - sharedExpertRankNum); localToken globalBS * (localExpertNum k ? localExpertNum : k); } std::vectorint64_t expandXShape{tpWorldSize * localToken, h}; // (A, H) std::vectorint64_t expertIdsShape{bs, k}; // (BS, K) std::vectorint64_t expertScalesShape{bs, k}; // (BS, K) std::vectorint64_t expandIdxShape{bs * k}; // (BS * K,) std::vectorint64_t quantExpandXOutShape{tpWorldSize * localToken, AlignN(AlignN(h, 32) AlignN(h, 8) / 8 * sizeof(float), 512)}; // (A, tokenMsgSize) std::vectorint64_t commCmdInfoOutShapeforcombine{(localToken epWorldSize) * 16}; // ((A epWorldSize) * 16,) std::vectorint64_t xOutShape{bs, h}; // (BS, H)其中AlignN(x, n) (x n - 1) / n * n与文档中tokenMsgSize Align512(Align32(H) Align8(H) / 8 * sizeof(float))的定义一致。7.3 两段式调用teardown 阶段// 1) 第一段计算 workspace 并获取 executor ret aclnnMoeDistributeCombineTeardownGetWorkspaceSize( expandX, quantExpandXOut, expertIds, expandIdx, expertScales, commCmdInfoOutforCombine, nullptr, nullptr, // xActiveMaskOptional、sharedExpertXOptional 传空指针 hcomEpName, epWorldSize, args.epRankId, moeExpertNum, expertSharedType, sharedExpertNum, sharedExpertRankNum, globalBS, commQuantMode, commType, nullptr, // commAlg 传空指针 xOut, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, ...); // 2) 按需申请 workspace if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, ...); } // 3) 第二段执行计算 ret aclnnMoeDistributeCombineTeardown(workspaceAddr, workspaceSize, executor, args.combineteardownstream); CHECK_RET(ret ACL_SUCCESS, ...); // 4) 等待 stream 同步后释放资源 ret aclrtSynchronizeStreamWithTimeout(args.combineteardownstream, timeOut);完整示例还演示了多进程/多线程的 EP 多卡组织方式main中先aclInit对每个 rank 执行aclrtSetDevice、aclrtCreateContext、aclrtCreateStream通过HcclCommInitAll初始化 HCCL 通信域再为每个 rank 启动一个线程执行LaunchOneProcess每个进程内通过HcclGetCommName从HcclComm句柄拿到通信域名称hcomEpName作为groupEp传入算子。运行结束后依次销毁 tensor、释放 Device 内存、销毁 stream/context 并aclrtResetDevice。八、源码实现纵览从 aclnn 接口到 Ascend C Kernel除了上述接口层代码本算子目录结构完整各层实现均可在仓库中定位接口层op_api/aclnn_moe_distribute_combine_teardown.cpp 与 op_api/aclnn_moe_distribute_combine_teardown.h负责入参校验与两段式入口算子信息库op_host/moe_distribute_combine_teardown_def.cpp声明输入/输出/属性及 AICore 配置shape 与 dtype 推断op_host/moe_distribute_combine_teardown_infershape.cpphost 侧 tilingop_host/op_tiling/ 下按架构分为arch22与arch35两套实现其中arch35对应 Ascend 950DT 的 DAV_3510 架构共享 moe_distribute_combine_teardown_tiling_base.cpp 中的公共校验与 workspace 设置逻辑设备侧 kernelop_kernel/moe_distribute_combine_teardown.cpp为 AIV-only 的 Ascend C 算子按__NPU_ARCH__ 3510条件编译加载 arch35 实现tiling key 定义见 op_kernel/moe_distribute_combine_teardown_tiling_key.h当前模板参数仅区分TILINGKEY_TP是否带 TP 维度图模式原型op_graph/moe_distribute_combine_teardown_proto.h供图编译Graph 模式使用配置与测试host 侧算子二进制配置位于 op_host/config/ascend950/单元测试覆盖 aclnn 接口tests/ut/op_api/、tilingtests/ut/op_host/与 kerneltests/ut/op_kernel/各层。如需进一步理解本算子的前序环节可参考 aclnnMoeDistributeCombineSetup 接口文档 及其调用示例若希望了解本仓库的整体贡献流程与算子开发规范可阅读 CONTRIBUTING.md。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考