CANN ops-transformer 算子解析：aclnnMoeTokenUnpermuteWithRoutingMap 的 MoE Token 逆置换与权重累加实现-北京尧图网络

算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载导读aclnnMoeTokenUnpermuteWithRoutingMap是 CANN ops-transformer 大模型算子库中用于 MoEMixture of Experts路由结果回写的关键算子。在 MoE 网络中Token 会先被aclnnMoeTokenPermuteWithRoutingMap按路由映射打散Permute到各专家槽位本算子则负责将专家计算完成后的permutedTokens依据sortedIndices记录的映射关系逆置换Unpermute并累加回原始的unpermutedTokens同时支持按probs路由权重对 token 进行加权。读完本文你将掌握该算子的计算公式、两段式 aclnn 接口原型、全部参数约束、返回码语义以及一份可直接参考编译运行的完整调用示例。产品支持情况根据 aclnnMoeTokenUnpermuteWithRoutingMap 接口文档与算子目录下的 README该算子在不同产品上的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持Kirin X90 处理器系列产品支持不支持 BFLOAT16Kirin 9030 处理器系列产品支持不支持 BFLOAT16从算子定义源码 op_host/moe_token_unpermute_with_routing_map_def.cpp 可以看到AICore().AddConfig分别注册了ascend910b、ascend910_93、ascend950、kirinx90、kirin9030等硬件配置与上表的产品支持情况一一对应其中 Kirin 系列的配置额外声明了PrecisionReduceFlag(true)且输入数据类型仅支持FLOAT16、FLOAT这正是 README 中Kirin 系列不支持 BFLOAT16的源码依据。功能说明与计算公式接口功能aclnnMoeTokenUnpermuteWithRoutingMap对经过aclnnMoeTokenPermuteWithRoutingMap处理的permutedTokens执行逆置换将其累加回原unpermutedTokens根据sortedIndices存储的下标取出permutedTokens中存储的输入数据如果存在probs数据permutedTokens会先与probs相乘最后进行累加求和输出计算结果。该算子与正向的aclnnMoeTokenPermuteWithRoutingMap形成打散—计算—还原的完整闭环是 MoE 网络中每个专家并行计算完成后将结果按 token 维度聚合回原始布局的关键一步。基础量定义计算公式中涉及的基础量定义如下$$ topK_num permutedTokens.size(0) // routingMapOptional.size(0) $$其中topK_num表示每个 token 预留的最大专家槽位数。当paddedMode为false时每个 token 实际选择的专家数可以小于等于topK_num未使用的槽位在sortedIndices中以-1表示计算时跳过该槽位。$$ numExperts probs.size(1) $$$$ numTokens probs.size(0) $$$$ capacity sortedIndices.size(0) // numExperts $$其中capacity表示在paddedMode为true时每个专家能够处理的 token 个数。从 tiling 源码 op_host/moe_token_unpermute_with_routing_map_tiling.cpp 中可以看到numOutTokens取自permutedTokens第 0 维totalLength取自sortedIndices第 0 维tokensNum即 numTokens与numExperts在存在probs时取自probs的 shape这些定义与文档公式完全一致。四种计算分支根据probs是否为 None 以及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:\quad 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:\quad unpermutedTokens[i//topK_num] permutedTokens[sortedIndices[i]] $$源码实现路径印证四种计算分支在算子内核入口 op_kernel/moe_token_unpermute_with_routing_map.cpp 中通过 tilingKey 进行分发TILING_KEY_IS(1000)走KernelMoeTokenUnpermuteWithRoutingMapPadpaddedMode 为 true 的路径对应源码 op_kernel/moe_token_unpermute_with_routing_map_pad.h其 Init 中直接以sortedIndices、probs为入参TILING_KEY_IS(1)haveProbstrue即 probs 不为 None 的非 padding 路径TILING_KEY_IS(0)haveProbsfalse即 probs 为 None 的非 padding 路径对应源码 op_kernel/moe_token_unpermute_with_routing_map_not_pad.h。非 padding 路径会先通过KernelMaskedSelectV3执行maskedSelect得到permute_probs对应公式中的probs.T.maskedSelect(routingMap.T)再由主 kernel 完成乘权与累加。在 aclnn 接口层 op_host/op_api/aclnn_moe_token_unpermute_with_routing_map.cpp 中paddedMode 为 true 时通过l0op::Sort对sortedIndices排序得到permuteTokenId与outIndex随后调用l0op::InplaceIndexAddWithSortedAscend 910B/910_93/DAV_3510 架构且满足 dtype 条件时或l0op::InplaceIndexAddAiCore完成排序后按索引累加paddedMode 为 false 时则由l0op::MoeTokenUnpermuteWithRoutingMap内核直接产出unpermutedTokens。这种排序索引累加的组合从底层印证了公式13中sort index_add的语义。两段式接口与函数原型与其他 aclnn 算子一致本算子采用两段式接口调用方式必须先调用aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnMoeTokenUnpermuteWithRoutingMap接口执行计算。第一段接口获取 workspace 大小与执行器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);两段式接口的实现在 op_host/op_api/aclnn_moe_token_unpermute_with_routing_map.cpp 中第一段接口完成参数校验、构图CREATE_EXECUTOR并通过uniqueExecutor-GetWorkspaceSize()返回 workspace 大小第二段接口则通过CommonOpExecutorRun完成实际计算下发。另外tiling 阶段在 op_host/moe_token_unpermute_with_routing_map_tiling.cpp 中申请了16 * 1024 * 1024 * 2字节的系统 workspace用于内核计算中的中间数据交换。aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize 参数说明第一段接口的完整参数说明如下表参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorpermutedTokensaclTensor*输入表示输入 tokenShape 中的 capacity 表示每个专家能够处理的 token 个数BFLOAT16、FLOAT16、FLOATNDpaddedMode 为 false(tokens_num * topK_num, hidden_size)paddedMode 为 true(experts_num * capacity, hidden_size)√sortedIndicesaclTensor*输入表示输入输出梯度的映射关系paddedMode 为 false 时要求索引取值范围 [0, tokens_num * topK_num - 1]paddedMode 为 true 时索引取值范围 [0, tokens_num - 1]INT32NDpaddedMode 为 false(tokens_num * topK_num)paddedMode 为 true(experts_num * capacity)√routingMapOptionalaclTensor*输入计算公式中的 routingMapOptional代表对应位置的 Token 是否被对应专家处理当输入 probsOptional 为空指针时不需要此输入应该传入空指针。当数据类型为 INT8取值支持 0、1当数据类型为 bool取值支持 true、falseINT8、BOOLND(tokens_num, experts_num)√probsOptionalaclTensor*输入计算公式中的 probsOptional代表对应位置的 Token 被对应专家处理后的结果在最终结果中的权重数据类型与 permutedTokens 相同或者当 permutedTokens 是 BFLOAT16 时 probsOptional 支持 FLOATBFLOAT16、FLOAT16、FLOATND与 routingMapOptional 一致√paddedModebool输入表示填充模式是否开启true 表示开启 paddedModefalse 表示关闭 paddedMode----restoreShapeOptionalaclIntArray*输入表示 unpermutedTokens 的 shapesize 大小为 2INT64---unpermutedTokensaclTensor*输出正向输出结果计算公式中的 unpermutedTokens-BFLOAT16、FLOAT16、FLOATND(tokens_num, hidden_size)√outIndexaclTensor*输出表示输出的索引值计算公式中的 outIndex当 paddedMode 为 false 时outIndex 表示槽位索引取值范围 [0, tokens_num * topK_num - 1]sortedIndices 中的 -1 仅表示无效输入槽位不改变 outIndex 的槽位索引范围当 paddedMode 为 true 时索引取值范围 [0, experts_num * capacity - 1]INT32NDpaddedMode 为 false(tokens_num * topK_num)paddedMode 为 true(experts_num * capacity)√permuteTokenIdaclTensor*输出计算公式中的 permuteTokenId索引取值范围 [0, tokens_num - 1]INT32NDpaddedMode 为 false(tokens_num * topK_num)paddedMode 为 true(experts_num * capacity)√permuteProbsaclTensor*输出计算公式中的 permuteProbs表示输出经过排序后的 probs与 probsOptional 相同BFLOAT16、FLOAT16、FLOATND1√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----关于非连续 Tensor列接口在内部会通过l0op::Contiguous将各输入统一转为连续内存再参与计算见 op_host/op_api/aclnn_moe_token_unpermute_with_routing_map.cpp 中l0op::Contiguous(permutedTokens, ...)等调用因此支持非连续张量作为输入输出侧则通过l0op::ViewCopy将计算结果写回用户指定的可能非连续的输出张量。数据类型与 shape 的源码校验接口第一段会对参数进行严格的 dtype 与 shape 校验对应 op_host/op_api/aclnn_moe_token_unpermute_with_routing_map.cpp 中的CheckDtypeValid与CheckShapeValidpermutedTokens、unpermutedTokens必须是二维TRANSPOSE_SHAPE_SIZE 2且 dtype 属于 {FLOAT16, FLOAT, BF16}sortedIndices、outIndex、permuteTokenId必须是一维INDEX_SHAPE 1且 dtype 为 INT32routingMapOptional的 dtype 支持 INT8、UINT8、BOOL存在probsOptional时其 dtype 必须与permutedTokens一致唯一例外是permutedTokens为 BF16 时probsOptional允许为 FLOAT即混合精度场景接口内通过l0op::Cast将 FP32 中间结果转回 BF16 输出对应源码中的isMixed逻辑probsOptional与routingMapOptional同时存在时两者的 shape 需保持一致格式方面各输入若为私有格式会输出 warning建议使用 ND 格式。返回值与错误码接口返回aclnnStatus状态码具体含义参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001必选输入或必选输出的 Tensor 是空指针ACLNN_ERR_PARAM_INVALID161002输入或输出的数据类型或 shape 不在支持的范围内ACLNN_ERR_INNER_NULLPTR561103topK_num 512ACLNN_ERR_INNER_NULLPTR561103probsOptional 的 shape 不在支持的范围内上述校验逻辑同样出现在 tiling 阶段的InputParamCheck中当topK 512时直接返回ge::GRAPH_FAILED并打印topK can not larger than 512见 op_host/moe_token_unpermute_with_routing_map_tiling.cpp。此外CheckNotNull还规定probsOptional与restoreShapeOptional不能同时为空指针。另外源码中还对空 Tensor 场景做了专门处理当sortedIndices为空或paddedMode false且permutedTokens为空时第一段接口会直接构造ZerosLike的零值图将各输出置零后即返回无需下发实际计算。aclnnMoeTokenUnpermuteWithRoutingMap 参数说明第二段接口参数相对简单均由第一段接口产生或指定参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值同样为aclnnStatus具体参见 aclnn 返回码。约束说明使用该算子时需要遵守以下约束确定性计算aclnnMoeTokenUnpermuteWithRoutingMap默认采用非确定性实现支持通过aclrtCtxSetSysParamOpt开启确定性。从 op_host/op_api/aclnn_moe_token_unpermute_with_routing_map.cpp 可以看到接口会通过aclrtGetSysParamOpt(ACL_OPT_DETERMINISTIC, ...)查询确定性开关在 Ascend 910B/910_93 上开启确定性时会切换到InplaceIndexAddWithSorted的确定性实现路径。topK_num 512paddedMode 为 false 时每个 token 最多预留 topK_num 个专家槽位routingMap 中每行为 1或 true的个数小于等于 topK_numsortedIndices 中允许使用 -1 表示无效槽位。以下场景后续版本会拦截如果提示 warning建议整改paddedMode 为 true且 topK_num experts_numtiling 阶段会打印topK is larger than experts_num, which may cause unexpected behaviorpaddedMode 为 true且 capacity tokens_numroutingMap 的数据类型或 shape 不符合要求输入 tensor 的数据格式不为 ND。Kirin 系列限制Kirin X90 / Kirin 9030 处理器系列产品不支持 BFLOAT16 输入参见 README 与算子定义中config_kirin仅注册DT_FLOAT16、DT_FLOAT的声明。调用示例以下为完整调用示例代码仅供参考具体编译和执行过程请参考编译与运行样例。该示例的独立可编译版本位于 examples/test_aclnn_moe_token_unpermute_with_routing_map.cppUT 测试版本位于 tests/ut/op_host/op_api/test_aclnn_moe_token_unpermute_with_routing_map.cpp。#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); // 调用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; } int main() { // 1. 固定写法device/stream初始化参考acl对外接口列表 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // check根据自己的需要处理 CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口定义构造 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; //in 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}; // 创建self aclTensor 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); // 创建out aclTensor 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需要修改为具体的API uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnMoeTokenUnpermuteWithRoutingMap第一段接口 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); // 根据第一段接口计算出的workspaceSize申请device内存 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侧需要根据具体API的接口定义修改 auto unpermutedTokensSize GetShapeSize(unpermutedTokensShape); std::vectorfloat unpermutedTokensData(unpermutedTokensSize, 0); ret aclrtMemcpy(unpermutedTokensData.data(), unpermutedTokensData.size() * sizeof(unpermutedTokensData[0]), 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和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(permutedTokens); aclDestroyTensor(sortedIndices); aclDestroyTensor(routingMapOptional); aclDestroyTensor(probs); aclDestroyTensor(unpermutedTokens); aclDestroyTensor(outIndex); aclDestroyTensor(permuteTokenId); aclDestroyTensor(permuteProbs); // 7. 释放device资源需要根据具体API的接口定义修改 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; }示例要点解读对照上表参数说明示例中的关键点如下示例采用paddedMode true此时permutedTokens的 shape 应为(experts_num * capacity, hidden_size)示例中取(2, 2)对应experts_num1、capacity2或experts_num2、capacity1的等价情形routingMapOptional使用charINT8构造取值全为 1表示每个 token 均被对应专家处理restoreShapeOptional通过aclCreateIntArray创建值为{2, 2}即unpermutedTokens的目标 shape 为(tokens_num2, hidden_size2)permuteProbs输出 shape 为{2}一维与参数表中维度(shape) 为 1的说明一致workspace 只有在workspaceSize 0时才需要aclrtMalloc申请使用完毕后必须aclrtFree释放。测试与验证仓库为该算子提供了完整的多层测试覆盖可用于验证调用与结果正确性ST 用例位于 tests/st/aclnnMoeTokenUnpermuteWithRoutingMap/executor_aclnnMoeTokenUnpermuteWithRoutingMap.py其用例模式枚举了padded、non_padded、non_padded_minus_one三种场景CASE_MODES其中non_padded_minus_one专门覆盖sortedIndices中含-1无效槽位的分支对应的数据生成配置见 atk_aclnnMoeTokenUnpermuteWithRoutingMap.jsonUT 用例tiling 单测见 tests/ut/op_host/test_moe_token_unpermute_with_routing_map_tiling.cppaclnn 接口单测见 tests/ut/op_host/op_api/test_aclnn_moe_token_unpermute_with_routing_map.cppkernel 单测见 tests/ut/op_kernel/test_moe_token_unpermute_with_routing_map.cpp。总结aclnnMoeTokenUnpermuteWithRoutingMap是 ops-transformer 中 MoE 路由数据回写链路的核心算子。它通过sortedIndices与routingMap精确还原每个 token 与其专家槽位的映射关系支持paddedMode开关以适配按专家填充容量与按 token 预留 topK 槽位两种主流 MoE 布局并支持可选的probs权重乘加与 BF16FLOAT 混合精度计算。在使用时重点关注topK_num 512的硬约束、probsOptional与restoreShapeOptional不能同时为空的入参约束以及 paddedMode 下topK_num experts_num、capacity tokens_num等会被后续版本拦截的告警场景。结合本文的两段式接口说明与完整调用示例即可在支持的 Ascend A2/A3/950 及 Kirin 系列产品上正确接入该算子。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐MoeTokenUnpermuteWithRoutingMap 算子全解析CANN ops-transformer 中 MoE Token 反置换与加权累加实现MoeTokenUnpermuteWithRoutingMap 算子全解析CANN ops transformer 中 MoE Token 反置换与加权累加实算子库人工智能深度学习AscendMoeTokenUnpermute 算子深度指南CANN ops-transformer 中的 MoE Token 去重排与加权聚合实现MoeTokenUnpermute 算子深度指南CANN ops transformer 中的 MoE Token 去重排与加权聚合实现导读 MoeToke算子库人工智能深度学习AscendCANN ops-transformer MoeTokenUnpermute 算子深度解析MoE Token 反置换与加权归约的 NPU 实现及 aclnn 调用实战CANN ops transformer MoeTokenUnpermute 算子深度解析MoE Token 反置换与加权归约的 NPU 实现及 aclnn算子库人工智能深度学习Ascend上一篇解锁Redash数据潜能从插件开发到社区共建全指南下一篇Instant Weather通知系统WorkManager NotificationHelper实现智能提醒创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

CANN ops-transformer 算子解析：aclnnMoeTokenUnpermuteWithRoutingMap 的 MoE Token 逆置换与权重累加实现

相关推荐

深圳GEO优化服务商哪家靠谱？五个硬指标与避坑指南

SCI投稿避坑指南：从Submitted到Under Review的完整链路

逆向必学：PE文件结构核心字段与加壳脱壳实战解析

更多资讯

视频会议系统操作手册：从结构设计到doc格式落地全攻略

Raycast 2.0 深度体验：AI 启动器重构与高效工作流配置指南

react-admin 生态全景指南：官方包、Enterprise Edition 与第三方扩展组件导航

Apache SkyWalking 后端存储配置指南：使用 H2 内存数据库作为默认存储

alt-tab-macos 设置侧边栏邮箱地址智能换行：EmailLineWrap 断行算法与实现解析

微带天线仿真设计全流程解析：从理论计算到工程实践

今日更新

OneUptime 自定义探针（Custom Probe）部署实战：私网监控、代理配置与断连排障全指南

大众TL52625前端框架材料要求详解：从性能测试到落地执行

TiXL 浮点运算算子库 Lib.numbers.float 完全指南：44 个算子的参数详解、源码原理与实战串联

本周精选

深入解析Transformer多头注意力机制与工程优化

OpenClaw 的 Skills 跑学习任务，模型通道改到 TaoToken 通道行不行？

ChatGPT报错Oops, an error occurred! 全链路排查指南

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？