CANN ops-transformer MoeFusedTopk 算子详解分组 TopK 专家选择的 aclnn 接口与 NPU 实现【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformerMoeFusedTopk 是 CANN ops-transformer 仓库中面向 MoEMixture-of-Experts混合专家模型的融合算子它把Sigmoid 激活 → 分组 TopK 排序 → 前 k 个专家选取 → 归一化 → 专家映射整条链路融合为一次 NPU 计算。本文以 moe_fused_topk 官方接口文档 为主线结合 算子定义、aclnn 两段式接口实现、Tiling 实现 与 NPU kernel 实现系统讲解其功能语义、接口参数、约束条件、调用方式与底层实现原理。读完本文你将能够理解分组 TopK 专家选择的完整计算流程正确配置全部 8 个属性参数并通过 aclnnMoeFusedTopkGetWorkspaceSize / aclnnMoeFusedTopk 两段式接口在 NPU 上完成一次完整调用。一、功能概述MoE 中分组选专家为什么要做一次融合在 MoE 架构中每个 token 需要从若干专家中选择最合适的 topK 个专家进行前向计算。经典实现往往拆成 Sigmoid、Add、TopK、ReduceSum、Sort、Gather 等多个独立算子串行执行多次访存与多次 kernel 启动带来明显开销。MoeFusedTopk 将这些步骤融合为一个算子接口功能MoE 计算中对输入 x 做 Sigmoid 计算对计算结果分组进行排序最后根据分组排序的结果选取前 k 个专家见 接口文档功能说明。它的特殊之处在于分组语义专家被划分为 groupNum 个组先在每个组内取 topN 个专家并求和以组内 topN 分数之和作为组排名依据选出分数最高的 groupTopk 个组随后只在被选中的组所覆盖的专家范围内做第二次 TopK得到最终的 topK 个专家及其分数。这种先粗筛组、再细筛专家的两级筛选策略可以显著缩小最终 TopK 的搜索空间同时保证不同组例如不同机器/不同算子分组之间不会出现专家选择失衡。1.1 产品支持情况MoeFusedTopk 并非所有昇腾产品都支持。根据接口文档与 算子 README 的产品支持表产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×在算子定义层面moe_fused_topk_def.cpp 中通过AICore().AddConfig(ascend910b)、AddConfig(ascend910_93)注册了 910b/A3 平台配置并为 ascend950 配置了动态编译DynamicCompileStaticFlag、动态 RankDynamicRankSupportFlag与动态 ShapeDynamicShapeSupportFlag支持同时指定使用moe_fused_topk_apt的 APTAdaptive Parameter Tilingkernel 文件。各平台下的二进制 kernel 配置可查看 ascend910b 配置目录 与 ascend950 配置目录。二、计算流程与公式拆解算子核心计算链路在接口文档中有明确的公式定义与 ST 测试的 golden 实现 完全对应逐段说明如下。第 1 步Sigmoid 激活对输入 x每个 token 对应各个专家的分数做 Sigmoid$$ sigmoidRes sigmoid(x) $$第 2 步叠加偏置 addNum将逐专家的偏置项 addNum 加到 Sigmoid 结果上得到用于排序的分数$$ normOut sigmoidRes addNum $$第 3 步分组内 TopN 组间排序按 groupNum 对 normOut 分组每组取组内分数最大的 topN 个专家求和以该和作为组的得分对全部组做 TopK取出分数最高的 groupTopk 个组$$ groupOut,\ groupId TopK(ReduceSum(TopK(Split(normOut,\ groupCount),\ ktopN,\ dim-1),\ dim-1),\ kgroupTopk) $$第 4 步组内二次 TopK 得到专家索引根据第 3 步得到的 groupId 取出对应组覆盖的专家分数做第二次 TopK得到分数与专家索引$$ normY,\ indices TopK(normOut[groupId,\ :],\ ktopK) $$第 5 步Gather 得到输出分数注意最终输出的分数 y 取自Sigmoid 之后、加偏置之前的sigmoidRes而不是normOut偏置只参与排序不进入输出$$ y gather(sigmoidRes,\ indices) $$第 6 步可选归一化当isNorm为 true 时对 y 按输入的 scale 做归一化$$ y \frac{y}{ReduceSum(y,\ dim-1)} \times scale $$第 7 步可选物理专家到逻辑专家映射当enableExpertMapping为 true 时将 indices 中的物理专家按 mappingNum 与 mappingTable 映射到逻辑专家得到最终的 indices 输出。这一能力用于支持专家并行场景下物理专家与逻辑专家编号不一致的部署方式。上述流程的每一环都能在 kernel 源码中找到对应实现ActivateAndAddSigmoidAddmoe_fused_topk_kernel.h、GroupReduceSumInternelImpl组内 TopKReduceSum、GroupTopkImpl组排序掩码广播、GatherSigmoidImplSortExtractGather、NormImplReduceSumMuls、CopyOut专家映射更新完整流水见Process()函数。三、两段式 aclnn 接口与函数原型MoeFusedTopk 遵循 CANN aclnn 两段式接口规范详见 两段式接口说明必须先调用第一段aclnnMoeFusedTopkGetWorkspaceSize完成参数校验、构图并获取 workspace 大小与执行器再调用第二段aclnnMoeFusedTopk真正下发计算。3.1 第一段接口原型aclnnMoeFusedTopkGetWorkspaceSize( const aclTensor* x, const aclTensor* addNum, const aclTensor* mappingNum, const aclTensor* mappingTable, uint32_t groupNum, uint32_t groupTopk, uint32_t topN, uint32_t topK, uint32_t activateType, bool isNorm, float scale, bool enableExpertMapping, aclTensor* y, aclTensor* indices, uint64_t* workspaceSize, aclOpExecutor** executor)3.2 第二段接口原型aclnnStatus aclnnMoeFusedTopk( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)从 aclnn_moe_fused_topk.cpp 的实现可以看到第一段接口的完整职责先执行CheckParams空指针/数据类型/Shape 校验随后通过l0op::Contiguous将非连续 Tensor 转为连续 Tensor对应参数表中非连续 Tensor 支持 √调用l0op::MoeFusedTopk构图再用ViewCopy将算子中间输出拷贝到用户指定的 y 与 indices最后通过executor-GetWorkspaceSize()返回 workspace 大小。第二段接口则直接调用CommonOpExecutorRun执行。四、参数说明第一段接口4.1 输入参数参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 Tensorx输入每个 token 对应各个专家的分数-FLOAT16、BFLOAT16、FLOAT32ND(numToken, expertNum)√addNum输入与输入 x 进行计算的偏置值-与 x 一致ND(expertNum)√mappingNum输入每个物理专家被实际映射到的逻辑专家数量enableExpertMapping 为 false 时不启用INT32ND(expertNum)√mappingTable输入每个物理专家/逻辑专家映射表enableExpertMapping 为 false 时不启用maxMappingNum 小于等于 128INT32ND(expertNum, maxMappingNum)√groupNum输入分组数量必须大于 0-----groupTopk输入被选择的组的数量必须大于 0-----topN输入组内选取的用于求和的专家数量必须大于 0-----topK输入最终选取的专家数量必须大于 0-----activateType输入激活类型当前只支持 0ACTIVATION_SIGMOID-----isNorm输入是否对输出进行归一化-----scale输入归一化后的系数乘-----enableExpertMapping输入是否开启物理专家到逻辑专家的映射-----4.2 输出参数参数名输入/输出描述数据类型数据格式维度(shape)非连续 Tensory输出Device 侧的 aclTensor输出每个 token 的 topK 分数FLOAT32ND(numToken, topK)√indices输出Device 侧的 aclTensortopK 个专家与 token 的映射关系INT32ND(numToken, topK)√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小----executor输出返回 op 执行器包含了算子计算流程----关于数据类型有两个容易忽略的细节x 与 addNum 数据类型必须一致。在 aclnn_moe_fused_topk.cpp 中通过OP_CHECK(x-GetDataType() addNum-GetDataType())强制校验同时算子定义 moe_fused_topk_def.cpp 中 x 与 add_num 均声明为{DT_FLOAT, DT_FLOAT16, DT_BF16}。y 恒为 FLOAT32、indices 恒为 INT32与输入精度无关。这一设计在 算子定义 与 ascend950 二进制配置FP32/FP16/BF16 三套 bin 的输出均固定为 float32 与 int32中保持一致。同时AscendC kernel 在 InitTilingData 中把内部计算统一提升为 float 精度即输入为 FP16/BF16 时先 Cast 成 float 参与全部排序计算避免低精度累积误差。4.3 属性参数语义对应算子 Attr上表中的 8 个标量参数在算子定义中对应 8 个 Attrmoe_fused_topk_def.cppAttrC 类型默认值含义group_numInt1分组数量必须大于 0group_topkInt1被选择的组数量必须大于 0top_nInt1组内用于求和的专家数量必须大于 0top_kInt1最终选取的专家数量必须大于 0activate_typeInt1激活类型当前仅支持 0ACTIVATION_SIGMOIDis_normBooltrue是否对输出归一化scaleFloat1.0归一化后的系数乘enable_expert_mappingBoolfalse是否开启物理专家到逻辑专家映射Tiling 阶段会从 TilingContext 中按固定索引读取这 8 个 Attrmoe_fused_topk_tiling.cpp其中enableExpertMapping还会决定 tilingKey置 1 表示启用映射分支见GetTilingKeykernel 侧则通过模板参数uint32_t enableExpertMapping在编译期展开映射逻辑moe_fused_topk_kernel.h。五、返回值与错误码两段接口均返回aclnnStatus状态码具体枚举含义参见 aclnn 返回码说明。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入 x 或者 addNum 为空指针ACLNN_ERR_PARAM_NULLPTR161001输出 y 或者 indices 为空指针ACLNN_ERR_PARAM_NULLPTR161001当 enableExpertMapping 为 true 时输入 mappingNum 或者 mappingTable 为空指针ACLNN_ERR_PARAM_INVALID161002输入或者输出的数据类型或数据格式不在支持的范围内ACLNN_ERR_PARAM_INVALID161002输入的参数不满足约束ACLNN_ERR_PARAM_INVALID161002输入输出的 Shape 不满足约束这些错误场景与 aclnn_moe_fused_topk.cpp 中的三段校验函数一一对应CheckNotNull空指针、CheckDtypeValid数据类型、CheckShapeShape 与参数约束。其中CheckShape内部还逐条实现了下一章的约束校验逻辑是约束说明的源码级实现。六、约束说明确定性计算aclnnMoeFusedTopk 默认确定性实现详见 确定性计算说明。x 和 addNum 数据类型必须一致。expertNum 必须为 groupNum 的整数倍。groupTopk 小于等于 groupNum。maxMappingNummappingTable 第二维小于等于 128。topK 小于等于 expertNum。topN 小于等于 expertNum / groupNum即组内专家数。expertNum 小于等于 1024。groupNum 小于等于 256。以上约束均在 CheckShape 中逐条实现常量GROUP_NUM_MAX256U、EXPERT_NUM_MAX1024U、MAX_MAPPING_NUM_MAX128U定义在文件头部。此外当输入 x 为空 Tensor 时接口直接返回workspaceSize0而不执行计算aclnn_moe_fused_topk.cpp属于空 Tensor 特判逻辑。七、调用示例完整的可编译示例位于 examples/test_aclnn_moe_fused_topk.cpp与接口文档中的示例代码一致。其调用流程也是所有 aclnn 算子的标准模板分为以下步骤步骤 1初始化 ACL 环境。aclInit→aclrtSetDevice→aclrtCreateStream示例中的Init函数。步骤 2构造输入与输出 aclTensor。示例选用的参数int64_t num_token 16; int64_t expert_num 32; int64_t max_mapping_num 16; uint32_t groupNum 2; // 分成 2 组每组 16 个专家 uint32_t groupTopk 2; // 2 组全部选中 uint32_t topN 2; // 组内取前 2 个专家分数求和作为组得分 uint32_t topK 4; // 最终为每个 token 选出 4 个专家 uint32_t activateType 0; // ACTIVATION_SIGMOID bool isNorm false; float scale 1.0; bool enableExpertMapping true; // 开启专家映射对应 shapestd::vectorint64_t xShape {num_token, expert_num}; // (16, 32) std::vectorint64_t addNumShape {expert_num}; // (32) std::vectorint64_t mappingNumShape {expert_num}; // (32) std::vectorint64_t mappingTableShape {expert_num, max_mapping_num}; // (32, 16) std::vectorint64_t yShape {num_token, topK}; // (16, 4) std::vectorint64_t indicesShape {num_token, topK}; // (16, 4)创建 Tensor 时示例中的CreateAclTensor模板函数完成aclrtMalloc申请 Device 内存 →aclrtMemcpy拷贝 Host 数据 → 计算连续 stride →aclCreateTensor创建 ND 格式的 aclTensor。注意输入 x 用ACL_FLOATFLOAT32、mappingNum/mappingTable 用ACL_INT32与前面参数表一一对应y 为ACL_FLOAT、indices 为ACL_INT32。步骤 3两段式调用。先调用第一段接口获取 workspaceSize 与 executor再按需申请 workspace 内存并调用第二段接口uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段接口参数校验 构图 计算 workspace 大小 ret aclnnMoeFusedTopkGetWorkspaceSize(x, addNum, mappingNum, mappingTable, groupNum, groupTopk, topN, topK, activateType, isNorm, scale, enableExpertMapping, y, indices, workspaceSize, executor); CHECK_FREE_RET(ret ACL_SUCCESS, ...); // 根据 workspaceSize 申请 device 内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_FREE_RET(ret ACL_SUCCESS, ...); } // 第二段接口执行计算 ret aclnnMoeFusedTopk(workspaceAddr, workspaceSize, executor, stream); CHECK_FREE_RET(ret ACL_SUCCESS, ...);步骤 4同步等待。调用aclrtSynchronizeStream(stream)等待任务执行结束。步骤 5取回结果。通过aclrtMemcpy将 y及 indices从 Device 拷贝回 Host并打印验证。最后通过aclrtDestroyStream、aclrtResetDevice、aclFinalize完成清理示例的Finalize函数。关于 workspace 的规模可以结合 Tiling 源码理解SetKernelTiling中usedWorkspaceSize SYS_WORKSPACESIZE usedCoreNum_ * workspacePerCore_其中SYS_WORKSPACESIZE为固定 16 MB 系统开销workspacePerCore_ secondDimSize_ * sizeof(float)是每个核用于暂存sigmoidRes addNum即 normOut的中间缓冲见 moe_fused_topk_tiling.cpp。整个 workspace 占用与 expertNumsecondDimSize和实际使用的核数成正比这也是第一段接口必须返回精确 workspaceSize 的原因。具体的编译与运行方法请参考 编译与运行样例对应单测见 tests/ut/op_host/op_api/test_aclnn_moe_fused_topk.cppST 测试的 ATK 封装见 executor_aclnnMoeFusedTopk.py其 CPU 侧 golden 实现可作为理解计算语义的参考先 Sigmoid、再 Add、reshape 分组后 topksum、按组得分排序并掩码、二次 sortgather、可选 norm 与专家映射。八、底层实现原理Tiling 与 Kernel 流水8.1 Host 侧 Tiling核间并行与 UB 切分Tiling 逻辑位于 moe_fused_topk_tiling.cpp核心决策包括按 token 维度并行切核GetUsedCore以 firstDimSizenumToken与coreNum的关系决定每个核处理的 batch 数batchPerCore与尾批tailBatch保证各核负载均衡当 numToken 小于等于核数时每核仅处理 1 个 token。UBUnified Buffer切分SplitUb按 32 字节对齐计算 tilingData、xIn、addNum、yOut、indicesOut、sigmoidBuf、sigmoidAdd、sorted、topkValue、assist、mappingNum、tempBuf 等各缓冲区的需求总量超出可用 UB 时报错返回。其中排序缓冲区按SORT_UNIT32对齐浮点缓冲区按 8 元素块对齐这对应了 AscendC TopK/Sort 指令对数据布局的硬件对齐要求。TopK 临时空间计算GetTopKTiling与GetTmpBuffSize分别通过AscendC::TopKTilingFunc、GetTopKMaxMinTmpSize、GetSigmoidMaxMinTmpSize、GetSortTmpSize、GetBroadCastMaxMinTmpSize等库函数精确计算 Sigmoid、TopK、Sort、BroadCast 各自需要的最大/最小临时缓冲取最大值作为topkMaxValue_/topkMinValue_写入 tiling 数据。8.2 Device 侧 Kernel五段流水Kernel 主流水定义在 moe_fused_topk_kernel.h 的Process()中与公式一一对应CopyInAddNum将 addNum 读入 VECCALC 缓冲非 float 输入先 Cast 为 float对每个 batchCopyInX读入 x →ActivateAndAdd完成Sigmoid(x)与Add(sigmoidRes, addNum)CopyToWorkspace/CopyFromWorkspace将 normOut 暂存到 Global workspace再按分组 reshape 回读未选中的补齐位置填充-3.4e38负无穷对应NEGATIVE_MIN_VAULE_FP32常量为后续排序做准备GroupReduceSumInternelImpl调用 AscendCTopKcalT, false, false, false, TOPK_NORMAL对每个组取 topN再逐个ReduceSum得到组得分GroupTopkImpl用Sort对组得分降序排序将选中的 groupTopk 个组置 1、其余置 0再通过BroadCast将掩码广播到组内所有专家并Mul回 normOut未选中组被置为负无穷GatherSigmoidImpl对掩码后的分数做SortExtract取出前 topK 个索引索引乘 4FLOAT_BYTES得到字节偏移后用Gather从sigmoidBuf未加偏置的 Sigmoid 结果取数得到 yNormImpl可选ReduceSum(y)→ 取倒数 →Muls逐元素乘再乘scaleCopyOut可选专家映射对每个选中的专家按mappingNum[expertId]计算冗余偏移(batchOffset_ loop) % expertMappingNum再通过mappingTable[expertId * tableDim_ offset]查表将物理专家索引替换为逻辑专家索引最后将 y 与 indices 写回 Global Memory。从该实现可以看到算子的几个关键工程决策全程在 float 精度下计算以保证排序稳定性用 workspace 中转实现分组这一非连续内存视图用负无穷填充实现未选中组被排除的语义Sigmoid 结果与排序结果分离存放保证输出 y 不受偏置影响。九、小结MoeFusedTopk 是 CANN ops-transformer 仓库中面向 MoE 场景的典型融合算子将激活、分组排序、两级 TopK、归一化、专家映射收敛为单算子通过两段式 aclnn 接口对外提供确定性计算能力。理解它的关键是把握三条主线计算语义偏置只参与排序、输出取自 Sigmoid 结果、两级筛选顺序、参数契约x/addNum 同类型、y 恒为 FLOAT32、8 个属性默认值及其约束以及实现路径Tiling 的核间并行与 workspace 规划、Kernel 的五段流水与负无穷掩码技巧。如需进一步深入可结合 接口文档、算子 README、ST golden 实现 与 调用示例 对照阅读。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考