MoeTokenUnpermuteWithRoutingMap 算子全解析CANN ops-transformer 中 MoE Token 反置换与加权累加实现【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本文是 CANN ops-transformer 开源算子库中MoeTokenUnpermuteWithRoutingMap算子的深度技术指南。该算子是 MoEMixture of Experts稀疏路由流水线的收尾环节负责将专家计算完成后打散的permutedTokens按路由映射还原回原始 token 顺序并支持按权重probs加权累加。读完本文你将掌握该算子的功能语义、四种计算分支paddedMode × 有无 probs、全部输入输出参数的 shape/类型约束、两段式 aclnn 接口调用方法以及其 host 侧 tiling 与 kernel 侧实现的源码级原理。算子定位MoE 路由链路中的反置换Unpermute环节在 MoE 大模型推理/训练中token 并不会被所有专家处理而是通过路由Routing机制为每个 token 挑选 topK 个专家。为了提升访存效率CANN 算子库采用先置换、再计算、后还原的流水线设计Permute置换由MoeTokenPermuteWithRoutingMap类算子依据routingMap将原始 token 重排为按专家分组的紧凑布局permutedTokens使每个专家可以连续读取自己负责的 token专家计算各个专家并行处理permutedTokens中属于自己的那部分数据Unpermute反置换MoeTokenUnpermuteWithRoutingMap根据sortedIndices记录的下标把处理结果累加回unpermutedTokens中原始 token 所在的位置完成结果还原。如果存在probstoken 被各专家处理结果的加权权重算子会先将permutedTokens与对应的probs相乘再做累加求和等效于对多个专家的输出做加权平均。在 MoE 网络如 Mixtral 类架构中这一加权累加正是多个专家贡献合并的标准语义。产品支持情况该算子在不同硬件平台上的支持情况如下表所示产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√注意Kirin X90 / Kirin 9030 处理器系列产品不支持 BFLOAT16permutedTokens、probs 等浮点输入仅支持 FLOAT16 与 FLOAT32。这一平台差异在算子定义的源码中也有对应体现Ascend 侧配置如ascend910b、ascend910_93、ascend950注册了 BF16 数据类型组合而kirinx90、kirin9030的配置只声明了 FLOAT16 与 FLOAT 两种类型参见 moe_token_unpermute_with_routing_map_def.cpp。功能说明与计算公式算子功能可概括为对经过aclnnMoeTokenPermuteWithRoutingMap处理的permutedTokens依据sortedIndices中存储的下标取回输入数据若存在probs则将permutedTokens与probs相乘最后累加求和并输出结果。计算前需要先明确以下派生量的定义topK_num每个 token 预留的最大专家槽位数topK_num permutedTokens.size(0) // routingMapOptional.size(0)。当paddedMode为false时每个 token 实际选择的专家数可以小于等于topK_num未使用的槽位在sortedIndices中以-1表示计算时跳过该槽位。numExperts专家数量numExperts probs.size(1)。numTokenstoken 数量numTokens probs.size(0)。capacitypaddedMode 下每个专家处理的固定 token 数capacity sortedIndices.size(0) // numExperts。根据probs是否存在与paddedMode的取值计算分为四种分支1probs 不为 NonepaddedMode 为 true 时permuteProbs[i // capacity, sortedIndices[i]] probs[i] permutedTokens permutedTokens * permuteProbs unpermutedTokens zeros(restoreShape, dtypepermutedTokens.dtype, devicepermutedTokens.device) permuteTokenId, outIndex sortedIndices.sort(dim-1) unpermutedTokens[permuteTokenId[i]] permutedTokens[outIndex[i]]2probs 不为 NonepaddedMode 为 false 时T 为转置操作permuteProbs probs.T.maskedSelect(routingMap.T) permutedTokens permutedTokens * permuteProbs unpermutedTokens zeros(restoreShape, dtypepermutedTokens.dtype, devicepermutedTokens.device) if sortedIndices[i] 0: unpermutedTokens[i // topK_num] permutedTokens[sortedIndices[i]] * permuteProbs[i]3probs 为 NonepaddedMode 为 true 时permuteTokenId, outIndex sortedIndices.sort(dim-1) unpermutedTokens[permuteTokenId[i]] permutedTokens[outIndex[i]]4probs 为 NonepaddedMode 为 false 时if sortedIndices[i] 0: unpermutedTokens[i // topK_num] permutedTokens[sortedIndices[i]]可以看到两条语义主线paddedMode true 走专家容量对齐布局每个专家固定处理capacity个 token累加目标位置通过sortedIndices.sort得到的permuteTokenId定位paddedMode false 走token 槽位对齐布局累加目标位置直接由i // topK_num推导同时用sortedIndices[i] 0跳过无效槽位。参数说明算子的输入、输出与属性参数如下表所示参数名输入/输出/属性描述数据类型数据格式permutedTokens输入待计算输入对应公式中的permutedTokensFLOAT16、BFLOAT16、FLOAT32NDsortedIndices输入对应公式中的sortedIndices。paddedMode为 false 时非负值表示permutedTokens中的有效行号-1 表示该槽位没有对应专家输出计算时跳过INT32NDroutingMapOptional可选输入代表对应位置的 Token 是否被对应专家处理对应公式中的routingMapOptionalINT8、BOOLNDprobsOptional可选输入代表对应位置的 Token 被对应专家处理后的结果在最终结果中的权重对应公式中的probsFLOAT16、BFLOAT16、FLOAT32NDpaddedMode属性true 表示开启 paddedModefalse 表示关闭 paddedMode。开启 paddedMode 时每个专家固定处理 capacity 个 token。关闭 paddedMode 时每个 token 预留 topK_num 个专家槽位实际被小于等于 topK_num 个专家处理未使用槽位在 sortedIndices 中以 -1 表示BOOL-restoreShapeOptional属性代表 unpermutedTokens 的 shapeaclIntArray*-unpermutedTokens输出对应公式中的unpermutedTokensFLOAT16、BFLOAT16、FLOAT32NDoutIndex输出对应公式中的outIndexINT32NDpermuteTokenId输出对应公式中的permuteTokenIdINT32NDpermuteProbs输出表示输出经过排序后的 probs对应公式中的permutedProbsFLOAT16、BFLOAT16、FLOAT32ND补充说明来自接口文档 aclnnMoeTokenUnpermuteWithRoutingMap.md 与算子定义源码permutedTokens的 shape 依模式而定paddedMode 为 false 时为(tokens_num * topK_num, hidden_size)paddedMode 为 true 时为(experts_num * capacity, hidden_size)其中 capacity 表示每个专家能够处理的 token 个数。sortedIndices的索引范围paddedMode 为 false 时要求取值在[0, tokens_num * topK_num - 1]paddedMode 为 true 时在[0, tokens_num - 1]。routingMapOptional的 shape 为(tokens_num, experts_num)当输入probsOptional为空指针时该输入也应传入空指针。INT8 类型取值支持 0、1BOOL 类型支持 true、false。probsOptional的 shape 与routingMapOptional一致数据类型需与permutedTokens相同当permutedTokens为 BFLOAT16 时probsOptional额外支持 FLOAT。restoreShapeOptional是一个aclIntArray*size 大小为 2即(tokens_num, hidden_size)用于指定unpermutedTokens的输出 shape。outIndex在 paddedMode 为 false 时表示槽位索引取值范围[0, tokens_num * topK_num - 1]sortedIndices中的 -1 仅表示无效输入槽位不改变 outIndex 的槽位索引范围paddedMode 为 true 时取值范围[0, experts_num * capacity - 1]。permuteTokenId索引取值范围为[0, tokens_num - 1]。输出permuteProbs的 shape 为 1 维长度与sortedIndices一致paddedMode 为 false 时为tokens_num * topK_numpaddedMode 为 true 时为experts_num * capacity。这些约束在算子定义文件 moe_token_unpermute_with_routing_map_def.cpp 中通过OpDef注册表可查证permuted_tokens与unpermuted_tokens声明为 REQUIRED 且支持 BF16/FLOAT16/FLOAT 的 ND 格式routing_map与probs为 OPTIONALout_index、permute_token_id为 REQUIRED 的 INT32两个属性drop_and_pad对应paddedMode默认 false与restore_shape默认空数组也被显式注册。约束说明topK_num ≤ 512。paddedMode 为 false 时每个 token 最多预留topK_num个专家槽位routingMap中每行为 1或 true的个数须小于等于topK_numsortedIndices中允许使用 -1 表示无效槽位。确定性计算aclnnMoeTokenUnpermuteWithRoutingMap默认是非确定性实现可通过aclrtCtxSetSysParamOpt开启确定性。以下场景在后续版本会被拦截若当前提示 warning建议整改paddedMode 为 true且topK_num experts_numpaddedMode 为 true且capacity tokens_numroutingMap 的数据类型或 shape 不符合要求输入 tensor 的数据格式不为 ND。从源码看topK_num 512的拦截在 tiling 阶段由InputParamCheck完成见 moe_token_unpermute_with_routing_map_tiling.cpp 中的OP_CHECK_IF(topK 512, ...)而capacity tokens_num与topK experts_num会在 tiling 时打印 warning 日志OP_LOGW与文档中的约束描述一一对应。调用说明两段式 aclnn 接口与 CANN 算子库中绝大多数 aclnn 接口一致该算子采用两段式调用方式详见 两段式接口说明先调用aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize获取计算所需的 workspace 大小与执行器再调用aclnnMoeTokenUnpermuteWithRoutingMap执行计算。完整样例代码位于 test_aclnn_moe_token_unpermute_with_routing_map.cpp接口头文件位于 aclnn_moe_token_unpermute_with_routing_map.h。第一段接口原型aclnnStatus aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize( const aclTensor *permutedTokens, const aclTensor *sortedIndices, const aclTensor *routingMapOptional, const aclTensor *probsOptional, bool paddedMode, const aclIntArray *restoreShapeOptional, aclTensor *unpermutedTokens, aclTensor *outIndex, aclTensor *permuteTokenId, aclTensor *permuteProbs, uint64_t *workspaceSize, aclOpExecutor **executor);第二段接口原型aclnnStatus aclnnMoeTokenUnpermuteWithRoutingMap( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream);第一段接口各参数的 shape、类型与使用说明汇总自接口文档参数输入/输出描述使用说明数据类型数据格式permutedTokens输入表示输入 tokenShape 中的 capacity 表示每个专家能够处理的 token 个数BFLOAT16、FLOAT16、FLOATNDsortedIndices输入表示输入输出梯度的映射关系paddedMode 为 false 时索引范围[0, tokens_num*topK_num-1]为 true 时[0, tokens_num-1]INT32NDroutingMapOptional输入代表对应位置 Token 是否被对应专家处理probsOptional 为空指针时此输入也传空指针INT8 取值 0/1bool 取值 true/falseINT8、BOOLNDprobsOptional输入对应 Token 被专家处理结果在最终结果中的权重数据类型与 permutedTokens 相同permutedTokens 为 BFLOAT16 时支持 FLOATBFLOAT16、FLOAT16、FLOATNDpaddedMode输入表示填充模式是否开启true/false--restoreShapeOptional输入表示 unpermutedTokens 的 shapesize 大小为 2INT64-unpermutedTokens输出正向输出结果-BFLOAT16、FLOAT16、FLOATNDoutIndex输出输出的索引值paddedMode 为 false 时表示槽位索引INT32NDpermuteTokenId输出公式中的 permuteTokenId索引范围[0, tokens_num-1]INT32NDpermuteProbs输出输出经过排序后的 probs与 probsOptional 相同BFLOAT16、FLOAT16、FLOATNDworkspaceSize输出需要在 Device 侧申请的 workspace 大小---executor输出op 执行器包含算子计算流程---返回值与错误码接口返回aclnnStatus状态码具体参见 aclnn 返回码说明。第一段接口完成入参校验以下场景会报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001必选输入或必选输出的 Tensor 是空指针ACLNN_ERR_PARAM_INVALID161002输入或输出的数据类型或 shape 不在支持的范围内ACLNN_ERR_INNER_NULLPTR561103topK_num 512或 probsOptional 的 shape 不在支持的范围内第二段接口的workspace、workspaceSize、executor、stream分别指定 Device 侧 workspace 内存地址、其大小由第一段接口获取、算子执行器以及执行任务所在 Stream。完整调用示例与运行流程下面给出一个 paddedMode true、带 probs 的最小可运行示例节选自接口文档与仓库 example完整的 paddedMode false / 带 -1 槽位等三个回归示例见 test_aclnn_moe_token_unpermute_with_routing_map.cpp。示例中 shape 设为{2, 2}数据全 1最终unpermutedTokens应为 2×2 的全 2 矩阵。#include acl/acl.h #include aclnnop/aclnn_moe_token_unpermute_with_routing_map.h #include iostream #include vector #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 shape_size 1; for (auto i : shape) { shape_size * i; } return shape_size; } int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); 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); 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); 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]; } *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 std::vectorint64_t permutedTokensShape {2, 2}; std::vectorint64_t sortedIndicesShape {2}; std::vectorint64_t routingMapOptionalShape {2, 2}; std::vectorint64_t probsShape {2, 2}; std::vectorint64_t unpermutedTokensShape {2, 2}; std::vectorint64_t outIndexShape {2}; std::vectorint64_t permuteTokenIdShape {2}; std::vectorint64_t permuteProbsShape {2}; void* permutedTokensDeviceAddr nullptr; void* sortedIndicesDeviceAddr nullptr; void* routingMapOptionalDeviceAddr nullptr; void* probsDeviceAddr nullptr; void* unpermutedTokensDeviceAddr nullptr; void* outIndexDeviceAddr nullptr; void* permuteTokenIdDeviceAddr nullptr; void* permuteProbsDeviceAddr nullptr; aclTensor* permutedTokens nullptr; aclTensor* sortedIndices nullptr; aclTensor* routingMapOptional nullptr; aclTensor* probs nullptr; aclTensor* unpermutedTokens nullptr; aclTensor* outIndex nullptr; aclTensor* permuteTokenId nullptr; aclTensor* permuteProbs nullptr; bool paddedMode true; std::vectorint64_t restoreShapeOptionalData {2, 2}; aclIntArray* restoreShapeOptional aclCreateIntArray(restoreShapeOptionalData.data(), restoreShapeOptionalData.size()); std::vectorfloat permutedTokensHostData {1.0, 1.0, 1.0, 1.0}; std::vectorint sortedIndicesHostData {1, 1}; std::vectorchar routingMapOptionalHostData {1, 1, 1, 1}; std::vectorfloat probsHostData {1, 1, 1, 1}; std::vectorfloat unpermutedTokensHostData {0, 0, 0, 0}; std::vectorint outIndexHostData {0, 0}; std::vectorint permuteTokenIdHostData {0, 0}; std::vectorfloat permuteProbsHostData {0, 0}; ret CreateAclTensor(permutedTokensHostData, permutedTokensShape, permutedTokensDeviceAddr, aclDataType::ACL_FLOAT, permutedTokens); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(sortedIndicesHostData, sortedIndicesShape, sortedIndicesDeviceAddr, aclDataType::ACL_INT32, sortedIndices); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(routingMapOptionalHostData, routingMapOptionalShape, routingMapOptionalDeviceAddr, aclDataType::ACL_INT8, routingMapOptional); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(probsHostData, probsShape, probsDeviceAddr, aclDataType::ACL_FLOAT, probs); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(unpermutedTokensHostData, unpermutedTokensShape, unpermutedTokensDeviceAddr, aclDataType::ACL_FLOAT, unpermutedTokens); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(outIndexHostData, outIndexShape, outIndexDeviceAddr, aclDataType::ACL_INT32, outIndex); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(permuteTokenIdHostData, permuteTokenIdShape, permuteTokenIdDeviceAddr, aclDataType::ACL_INT32, permuteTokenId); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(permuteProbsHostData, permuteProbsShape, permuteProbsDeviceAddr, aclDataType::ACL_FLOAT, permuteProbs); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 两段式调用 CANN 算子库 API uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize( permutedTokens, sortedIndices, routingMapOptional, probs, paddedMode, restoreShapeOptional, unpermutedTokens, outIndex, permuteTokenId, permuteProbs, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); void* workspaceAddr nullptr; 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 aclnnMoeTokenUnpermuteWithRoutingMap(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMoeTokenUnpermuteWithRoutingMap failed. ERROR: %d\n, ret); return ret); // 4. 同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 将 device 侧结果拷回 host 并打印 auto unpermutedTokensSize GetShapeSize(unpermutedTokensShape); std::vectorfloat unpermutedTokensData(unpermutedTokensSize, 0); ret aclrtMemcpy(unpermutedTokensData.data(), unpermutedTokensData.size() * sizeof(float), unpermutedTokensDeviceAddr, unpermutedTokensSize * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i unpermutedTokensSize; i) { LOG_PRINT(unpermutedTokensData[%ld] is: %f\n, i, unpermutedTokensData[i]); } // 6. 释放 aclTensor / aclIntArray aclDestroyTensor(permutedTokens); aclDestroyTensor(sortedIndices); aclDestroyTensor(routingMapOptional); aclDestroyTensor(probs); aclDestroyTensor(unpermutedTokens); aclDestroyTensor(outIndex); aclDestroyTensor(permuteTokenId); aclDestroyTensor(permuteProbs); aclDestroyIntArray(restoreShapeOptional); // 7. 释放 device 资源 aclrtFree(permutedTokensDeviceAddr); aclrtFree(sortedIndicesDeviceAddr); aclrtFree(routingMapOptionalDeviceAddr); aclrtFree(probsDeviceAddr); aclrtFree(unpermutedTokensDeviceAddr); aclrtFree(outIndexDeviceAddr); aclrtFree(permuteTokenIdDeviceAddr); aclrtFree(permuteProbsDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }调用步骤可归纳为初始化 ACL 资源 → 构造输入输出 aclTensor → 第一段接口取 workspace 大小并创建执行器 → 申请 workspace → 第二段接口执行 → 同步 Stream → 拷回结果 → 释放资源。编译与运行样例的完整流程可参考 编译与运行样例说明。源码级实现剖析算子定义OpDef 注册在 moe_token_unpermute_with_routing_map_def.cpp 中算子通过OP_ADD(MoeTokenUnpermuteWithRoutingMap)完成注册声明了 4 个输入、4 个输出和 2 个属性并针对不同平台配置了差异化的 AICore 配置Ascend 平台ascend910b、ascend910_93、ascend950启用DynamicCompileStaticFlag、DynamicRankSupportFlag、DynamicShapeSupportFlag即支持动态 shape 与动态 rankKirin 平台kirinx90、kirin9030额外启用DynamicFormatFlag、PrecisionReduceFlag但数据类型收窄为 FLOAT16/FLOAT 两档且关闭了NeedCheckSupportFlag。这一注册信息与 README 中Kirin 系列不支持 BFLOAT16的说明互相印证。Tiling 策略按 paddedMode 分流两条切分路径Tiling 逻辑集中在 moe_token_unpermute_with_routing_map_tiling.cpp入口TilingMoeTokenUnpermuteWithRoutingMap读取属性padded_mode后分派到两条独立路径paddedMode falseTilingCompute从 permutedTokens 第 0 维得到numOutTokens、第 1 维得到hiddenSize从 sortedIndices 得到totalLength若存在 probs则tokensNum、numExperts取自 probs 的 shapetopK permutedTokens.size(0) / tokensNum。随后执行SetCoreNum按 token 数取核数、TilingHiddenSize按 UB 容量对 hidden 维切分512 字节对齐、SetBufferNum在 2~4 之间权衡 buffer 数、TilingToken按剩余 UB 空间切分 token 维等一系列切分。同时通过SetTilingKey设置 tilingKey 第 0 位0 表示 probs 为 None1 表示 probs 非 None。存在 probs 时还会SetScheduleMode(1)并预留至少 5120 字节给 sorted_indices 和 probs 的暂存空间。paddedMode trueTilingMoeTokenUnpermuteWithRoutingMapPad以numExperts * capacity为行总数按 AIV 核数划分 front_core / tail_core 两类核并结合 UB 容量把每核要处理的 token 再切成多轮循环loop_time_each_front_core、num_tokens_front_core_each_loop等字段tilingKey 固定为 1000。若capacity tokens_num会打印 warning与文档中的约束提示一致。tiling 数据结构定义在 moe_token_unpermute_with_routing_map_tiling.hMoeTokenUnpermuteWithRoutingMapTilingData记录hidden_size、top_k、num_out_tokens、hidden/token 两维的切分信息length/num/remain、used_core_num、buffer_num并内嵌MaskedSelectMTUTilingData用于 maskedSelect 子流程与MoeTokenUnpermuteWithRoutingMapPadTilingData用于 pad 模式两个结构体。Kernel 实现三种模板分支Kernel 入口在 moe_token_unpermute_with_routing_map.cpp根据 tilingKey 分派tilingKey 1000paddedMode true调用KernelMoeTokenUnpermuteWithRoutingMapPad实现见 moe_token_unpermute_with_routing_map_pad.h。该实现按 front/tail 核划分 blockOffset每个核读取 sortedIndices 定位 probs 中对应元素probsGM.GetValue(sortedIndicesGM.GetValue(posIdx) * num_experts posIdx / capacity)用DataCopyPad写回 unpermutedTokens。tilingKey 1probs 存在先执行KernelMaskedSelectV3完成permuteProbs probs.T.maskedSelect(routingMap.T)对应 masked_select.h再调用KernelMoeTokenUnpermuteWithRoutingMapT1, T2, T3, true做带权累加。tilingKey 0probs 为 None调用KernelMoeTokenUnpermuteWithRoutingMapT1, T2, T3, false做纯累加。非 pad 内核 moe_token_unpermute_with_routing_map_not_pad.h 的关键逻辑IsValidInputIndex检查in_token_index 0 in_token_index num_out_tokens这正是跳过 sortedIndices 中 -1 无效槽位语义的落点CalPartOutToken对每个 token 的 topK 个槽位逐个处理先取索引与权重判断validIndex与prob_value ! 0决定是否拷贝输入然后CalFirstToken/CalToken在 Vector 计算单元上完成Muls乘 prob与Add累加——注意非 FLOAT 输入会先Cast到 float 中间缓冲计算最后再Cast回原类型CAST_RINT输出这是 tiling 中为 FP16/BF16 预留 cast 空间的直接原因hidden 维超出单次 UB 容纳能力时按hidden_splited_num/hidden_splited_remain分块处理非 32 字节对齐的片段走DataCopyPad。测试验证仓库为该算子提供了多层测试UThost 侧单测tests/ut/op_host/op_api/test_aclnn_moe_token_unpermute_with_routing_map.cpp 覆盖有 probs paddedModefalse有 probs paddedModetrue等典型组合用TensorDesc(...).ValueRange(...)构造随机数据并通过OP_API_UT宏断言输出tiling 单测位于 test_moe_token_unpermute_with_routing_map_tiling.cpp。Kernel 单测tests/ut/op_kernel/test_moe_token_unpermute_with_routing_map.cpp 验证核函数数值正确性。ST系统级测试tests/st/aclnnMoeTokenUnpermuteWithRoutingMap/ 下提供 ATK 测试配置atk_aclnnMoeTokenUnpermuteWithRoutingMap.json与执行脚本executor_aclnnMoeTokenUnpermuteWithRoutingMap.py。可运行回归示例examples/test_aclnn_moe_token_unpermute_with_routing_map.cpp 内置三个场景paddedModetrue 基础示例、paddedModefalse 无 -1 槽位回归期望 unpermutedTokens 为{11,22,33,44,...}的加权累加结果、paddedModefalse 带 -1 槽位且 probs 按槽位对齐的示例验证无效槽位权重为 0 时被跳过并用 1e-3 容差自动断言结果。总结MoeTokenUnpermuteWithRoutingMap是 CANN ops-transformer 中 MoE 路由链路的关键收尾算子其核心价值在于通过sortedIndices与routingMap把专家并行计算后的离散结果高效还原为原始 token 布局并内建 probs 加权累加语义。掌握 paddedMode 两种布局模式专家容量对齐 vs token 槽位对齐、-1无效槽位的跳过规则、topK_num ≤ 512等约束以及两段式 aclnn 接口的参数契约即可在 Ascend A2/A3/950 及 Kirin 系列平台上正确集成与调优该算子。如需进一步深入可结合本文引用的算子定义、tiling 与 kernel 源码以及三层测试用例继续研读。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考