算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载本篇文章围绕 CANN ops-transformer 仓库中 MoEMixture of Experts专家并行EP场景下的aclnnMoeTokenUnpermuteWithEpGrad算子展开完整讲解其功能定位、数学公式、两段式 aclnn 接口原型、全部入参与出参约束、错误码语义并结合moe/moe_token_unpermute_with_ep_grad目录下的算子注册、infershape、tiling 与 AICore Kernel 源码剖析其底层实现原理。读完本文你将掌握在 Atlas A2/A3 及 Ascend 950 系列产品上通过 aclnn 两段式接口正确调用该算子并能够依据源码理解其数据流与核间/核内切分策略。一、算子背景与功能定位MoeTokenUnpermuteWithEpGrad是aclnnMoeTokenUnpermuteWithEp的反向传播算子用于 MoE 模型中Token 逆置换这一环节的梯度回传。在 MoE 专家并行Expert ParallelEP的训练框架中前向过程通常包含permute置换根据路由结果将 token 按专家分组重排得到permutedTokens专家计算每个专家处理属于自己的 token 子集unpermute逆置换将专家计算结果按照sortedIndices重新排列回原始 token 顺序得到unpermutedTokens。在反向传播阶段unpermutedTokens会收到来自上游的梯度unpermutedTokensGrad需要将其沿逆置换路径回传给permutedTokens得到permutedTokensGradOut若前向使用了路由概率probs还需将梯度沿probs回传得到probsGradOut。这正是本算子要完成的计算。与同类算子不同的是本算子额外引入了rangeOptional属性用于表达 EP 切分后的有效 token 范围每个设备只负责[rangeOptional[0], rangeOptional[1])范围内的 token 子集从而避免在反向传播时跨设备处理不属于自己的 token实现与 moe_token_unpermute_with_ep_grad/README.md 描述一致的前后向对齐。二、产品支持情况根据 aclnnMoeTokenUnpermuteWithEpGrad.md 中的产品适配声明各产品支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持该声明与算子注册代码 moe_token_unpermute_with_ep_grad_def.cpp 中的 AICore 配置一致注册了ascend910b、ascend910_93对应 Atlas A2/A3 系列并为ascend950单独配置了DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)等动态能力开关。三、功能与计算公式算子接口功能为aclnnMoeTokenUnpermuteWithEp 的反向传播即同时计算permutedTokens与probs两个输入的梯度。根据probs是否传入分为两条计算路径其中下标 $i \in {0, 1, 2, ..., num_tokens - 1}$。3.1 probs 非 None 的计算路径当传入probsOptional时计算分三步完成第一步构建 unpermutedTokens按范围裁剪的逆置换当rangeOptional[0] sortedIndices[i] rangeOptional[1]时$$ unpermutedTokens[i] permutedTokensOptional[sortedIndices[i]-rangeOptional[0]] $$否则索引落在 EP 范围之外视为截断$$ unpermutedTokens[i] 0 $$第二步计算 probs 的梯度沿 topK 维做求和归约$$ unpermutedTokens unpermutedTokens.reshape(-1, topkNum, hiddenSize) $$$$ unpermutedTokens unpermutedTokensGrad.unsqueeze(1) \times unpermutedTokens $$$$ probsGrad \sum_{k0}^{topkNum}(unpermutedTokens_{i,j,k}) $$即将unpermutedTokensGrad与逆置换结果逐元素相乘后对 topK 维度做求和得到 shape 为(tokens_num, topK)的probsGradOut。第三步计算 permutedTokens 的梯度当rangeOptional[0] sortedIndices[i] rangeOptional[1]时$$ permutedTokensGradOut[sortedIndices[i]] ((unpermutedTokensGrad.unsqueeze(1) \times probs.unsqueeze(-1)).reshape(-1, hiddenSize))[i] $$即先计算unpermutedTokensGrad与probs的外积再按sortedIndices逆散射回permutedTokensGradOut的对应位置。3.2 probs 为 None 的计算路径当不传入probsOptional时只计算permutedTokens的梯度退化为纯数据搬移式散射当rangeOptional[0] sortedIndices[i] rangeOptional[1]时$$ permutedTokensGradOut[sortedIndices[i]-rangeOptional[0]] unpermutedTokensGrad[i] $$即把unpermutedTokensGrad[i]按sortedIndices[i]直接写入permutedTokensGradOut不涉及任何乘加运算。该路径在 Kernel 侧对应 moe_token_unpermute_grad_prob_none_with_ep.h 中MoeTokenUnpermuteGradProbNone的实现从源码看其核心循环仅包含 rowIdMap 搬运、unpermuted_tokens_grad 分块搬入与 permuted_tokens_grad 散射搬出无需向量计算单元参与乘加。四、两段式接口与函数原型每个 aclnn 算子遵循两段式接口规范必须先调用第一段接口aclnnMoeTokenUnpermuteWithEpGradGetWorkspaceSize完成入参校验、获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用第二段接口aclnnMoeTokenUnpermuteWithEpGrad真正执行计算。第一段接口原型aclnnStatus aclnnMoeTokenUnpermuteWithEpGradGetWorkspaceSize( const aclTensor *unpermutedTokensGrad, const aclTensor *sortedIndices, const aclTensor *permutedTokensOptional, const aclTensor *probsOptional, bool paddedMode, const aclIntArray *restoreShapeOptional, const aclIntArray *rangeOptional, int64_t topkNum, const aclTensor *permutedTokensGradOut, const aclTensor *probsGradOut, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口原型aclnnStatus aclnnMoeTokenUnpermuteWithEpGrad( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)其中第二段接口的workspace、workspaceSize、executor均来自第一段接口的输出stream指定任务执行所在的 Device 流。五、第一段接口参数详解aclnnMoeTokenUnpermuteWithEpGradGetWorkspaceSize的完整参数说明如下表参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorunpermutedTokensGradaclTensor输入公式中的 unpermutedTokensGrad正向输出 unpermutedTokens 的梯度tokens_num 代表 token 个数hidden_size 代表 token 的维度大小BFLOAT16、FLOAT16、FLOAT32ND(tokens_num, hidden_size)√sortedIndicesaclTensor输入公式中的 sortedIndices索引取值范围 [0, tokens_num * topK - 1]设置 rangeOptional 时仅处理 range 范围内元素INT32ND(tokens_num * topK)√permutedTokensOptionalaclTensor输入公式中的 permutedTokensOptional要求 shape 与 sortedIndices 和 unpermutedTokensGrad 匹配与 unpermutedTokensGrad 一致ND(tokens_num * topK, hidden_size)√probsOptionalaclTensor输入公式中的 probsOptional当 probsOptional 传入时以 probsOptional 第 2 维作为有效 topK当 probsOptional 不传入时以 topkNum 作为有效 topK默认 1BFLOAT16、FLOAT16、FLOAT32ND(tokens_num, topK)√paddedModebool输入true 表示开启 paddedModefalse 表示关闭 paddedModepaddedMode 表示填充模式目前仅支持 false----restoreShapeOptionalaclIntArray输入公式中的 restoreShapeOptional当 paddedMode 为 true 后生效否则不会对其进行操作当前仅支持 nullptr----rangeOptionalaclIntArray输入公式中的 rangeOptionalep 切分的有效范围支持 nullptr 或长度为 2 的 aclIntArray。为 nullptr 或取值为 {-1, -1} 时不生效生效时要求归一化后的起始位置不大于结束位置----topkNumint64_t输入公式中的 topkNum每个 token 被选中的专家个数probsOptional 不传入时作为有效 topKprobsOptional 传入时调用侧应保证 topkNum 与 probsOptional 第 2 维语义一致。有效 topK 取值范围为 [1, 512]----permutedTokensGradOutaclTensor输出permutedTokens 的梯度公式中的 permutedTokensGradOut-与 unpermutedTokensGrad 一致ND(tokens_num * topK, hidden_size)×probsGradOutaclTensor输出公式中的 probsGradOutprobs 的梯度probsOptional 传入时 shape 与 probsOptional 一致与 probsOptional 一致ND(tokens_num, topK)×workspaceSizeuint64_t输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor输出返回 op 执行器包含了算子计算流程-----补充说明源自源码可选输入的语义在 moe_token_unpermute_with_ep_grad_def.cpp 中permuted_tokens与probs均以ParamType(OPTIONAL)注册且所有输入均声明了AutoContiguous()框架会在必要时自动将非连续输入转为连续内存后再下发 Kernel。属性默认值padded_mode默认falserestore_shape默认{1, 1}range默认{-1, -1}不生效topk_num默认1。这与接口层paddedMode 仅支持 false、restoreShapeOptional 仅支持 nullptr、rangeOptional 默认不生效的约束相互印证。形状推导规则moe_token_unpermute_with_ep_grad_infershape.cpp 规定输入必须满足permutedTokens与unpermutedTokensGrad为 2 维、sortedIndices为 1 维输出permutedTokensGradOut的 shape 为(sortedIndices第0维, hidden_size)probsGradOut的 shape 直接拷贝probs的 shape当probs为 nullptr 时probsGradOut置空。混精推导infershape 中probsGradOut的数据类型优先取probs的数据类型支持 probs 与 permutedTokens 类型不同的混精场景否则与permutedTokens一致。六、返回值与错误码接口返回aclnnStatus状态码具体含义参见 aclnn返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001必选输入或必选输出的 Tensor 是空指针ACLNN_ERR_PARAM_INVALID161002输入和输出的数据类型或数据格式不在支持的范围内ACLNN_ERR_INNER_TILING_ERROR561002topkNum 512ACLNN_ERR_INNER_TILING_ERROR561002输入和输出的 shape 不符合要求ACLNN_ERR_INNER_TILING_ERROR561002rangeOptional[1] rangeOptional[0]从源码看上述校验与 tiling 阶段的一致性检查吻合moe_token_unpermute_with_ep_grad_tiling.cpp 中定义了MAX_TOPK 512、EP_RANGE_SIZE 2等常量并在切分前对 topK 与 range 合法性进行判定。第二段接口aclnnMoeTokenUnpermuteWithEpGrad的参数只有 4 个均为输入参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnMoeTokenUnpermuteWithEpGradGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream七、约束说明调用本算子必须遵守以下约束确定性aclnnMoeTokenUnpermuteWithEpGrad默认为确定性实现同一输入多次调用结果可复现。Shape 匹配不支持 broadcast各输入、输出 shape 需显式匹配。permutedTokensOptional传入时其第 0 维需等于sortedIndices第 0 维、第 1 维需等于unpermutedTokensGrad第 1 维hidden_sizeprobsOptional传入时其第 0 维需等于unpermutedTokensGrad第 0 维tokens_numpermutedTokensGradOut第 0 维需等于sortedIndices第 0 维、第 1 维需等于 hidden_sizeprobsOptional传入时probsGradOut需与probsOptional一致。有效 topK 约束probsOptional传入时有效 topK 为probsOptional第 2 维不传入时有效 topK 为topkNum默认 1有效 topK 取值范围为 [1, 512]。rangeOptional 约束支持 nullptr不生效全程处理或长度为 2 的 aclIntArray。传入时支持负值按 tokens_num * topK 归一化会被裁剪到 [0, tokens_num * topK] 范围内表示左闭右开区间 [start, end)要求归一化后 start end且仅该范围内 sortedIndices 元素参与写入。空 Tensor 约束当 tokens_num、hidden_size 或有效 topK 为 0 时为空 Tensor 场景。Ascend 950PR/Ascend 950DT 支持空 Tensor第一段接口快速返回workspaceSize 为 0不进行有效计算Atlas A3 训练/推理系列产品与 Atlas A2 训练/推理系列产品不支持空 Tensor第一段接口返回错误。八、源码实现剖析8.1 Kernel 入口与分支调度Kernel 入口位于 moe_token_unpermute_with_ep_grad.cpp通过TILING_KEY区分四种组合0padded_mode False不存在 probs1padded_mode False存在 probs10padded_mode True不存在 probs11padded_mode True存在 probs。分支0/10走MoeTokenUnpermuteGradProbNone纯散射路径分支1/11在开启DTYPE_PROBS编译宏后走MoeTokenUnpermuteGradProbNotNone乘加归约路径。结合约束paddedMode 仅支持 false实际可用分支为0与1其余分支为未来扩展预留。8.2 有 probs 路径的计算流水从 moe_token_unpermute_grad_prob_not_none_with_ep.h 的实现可以看到典型的 AICore 流水设计双缓冲inQueueUnpermuted、inQueuePermutedTokens、inQueueProb、outQueuePermutedTokensGrad、outQueueProbGrad均以DOUBLE_BUFFER值为 2定义于 moe_token_unpermute_grad_base_with_ep.h申请空间并通过 ping-pong 标志位与 event 同步EVENT_ID0~EVENT_ID7交替使用两块 buffer实现 MTE2 搬运与 V 向量计算的流水重叠。索引驻留寄存器/UB为避免频繁GetValue访问 Global Memory单次循环最多将 512 个sortedIndices值与 512 个probs值缓存在本地数组indicesArray、probsArray上限INDICES_PROBS_MAX_RESERVE_NUM 512。probs 梯度归约probsGrad需要对 hiddenSize 维做逐元素相乘后求和。源码针对不同规模选择不同归约策略moe_token_unpermute_grad_base_with_ep.h 中的ReduceSumFunchiddensizeLen 4096 时先用二分加法累加到 4096再经多级BlockReduceSum与WholeReduceSum归约到 1 个标量 512、 64 时按同样思路递减小于 64 时直接WholeReduceSum。范围截断搬入与散射均检查inputGradOffset start inputGradOffset end inputGradOffset numOutTokens越界位置写 0搬入或跳过写出散射与公式中 rangeOptional 边界外置 0 的语义一致。8.3 Tiling 切分策略Tiling 阶段定义于 moe_token_unpermute_with_ep_grad_tiling.cpp切分信息结构体见 moe_token_unpermute_with_ep_grad_tiling.h核间切分按可用 AIV 核数将tokensNum均分前formerCoreNum个核每核多处理 1 个 tokentokenNumEachCore tokenNumTailCore 1其余tailCoreNum个核处理tokenNumTailCore个 token每核对应的sortedIndices量为token数 * topKrowIdMapEachCore/rowIdMapTailCore。核内切分无 probs 路径CoreSplitInfoProbIsNone依据 UB 总大小与 hiddenSize 对齐值估算单次可搬入的 indices 数量hiddenSize 过大时按hiddenSizeLoopTimes切块循环处理并保证单次搬入量最小 32B 对齐。核内切分有 probs 路径CoreSplitInfoProbIsNotNone优先令indicesReserveNum取 topK 的倍数且不超过INDICES_RESERVE_MAX_NUM 256当单次载入空间不足时切换为切分 hiddenSize 的策略并保证indicesNumPerLoop是 topK 的倍数、topK 是permutedTokenNumPerLoop的倍数以满足循环边界约束。8.4 算子注册与配置算子注册于 moe_token_unpermute_with_ep_grad_def.cpp输入数据类型白名单为 {BF16, FP16, FP32}sortedIndices固定 INT32格式全部为 ND。ascend910b与ascend910_93使用普通 AICore 配置ascend950配置了动态 shape、动态 rank 支持与静态编译标志。平台二进制配置位于 ascend910b 配置目录 与 ascend910_93 配置目录含binary.json与simplified_key.ini。九、调用示例可编译参考以下为 test_aclnn_moe_token_unpermute_with_ep_grad.cpp 的完整示例逻辑编译与运行步骤可参考编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_moe_token_unpermute_with_ep_grad.h #include iostream #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; } void PrintOutResult(std::vectorint64_t shape, void **deviceAddr) { auto size GetShapeSize(shape); std::vectorfloat resultData(size, 0); auto ret aclrtMemcpy( resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET( ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ); for (int64_t i 0; i size; i) { LOG_PRINT(mean result[%ld] is: %f\n, i, resultData[i]); } } 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 API手册 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t permutedTokensShape {3, 2}; std::vectorint64_t unpermutedTokensGradShape {1, 2}; std::vectorint64_t probsShape {1, 3}; std::vectorint64_t sortedIndicesShape {3}; std::vectorint64_t permutedTokensGradShape {3, 2}; std::vectorint64_t probsGradShape {1, 3}; void* permutedTokensDeviceAddr nullptr; void* unpermutedTokensGradDeviceAddr nullptr; void* probsDeviceAddr nullptr; void* sortedIndicesDeviceAddr nullptr; void* permutedTokensGradDeviceAddr nullptr; void* probsGradDeviceAddr nullptr; aclTensor* permutedTokens nullptr; aclTensor* unpermutedTokensGrad nullptr; aclTensor* probs nullptr; aclTensor* sortedIndices nullptr; bool paddedMode false; aclTensor *permutedTokensGrad nullptr; aclTensor *probsGrad nullptr; std::vectorfloat permutedTokensHostData {1, 1, 1, 1, 1, 1}; std::vectorfloat unpermutedTokensGradHostData {1, 1}; std::vectorfloat probsHostData {1, 1, 1}; std::vectorint sortedIndicesHostData {0, 1, 2}; std::vectorfloat permutedTokensGradHostData {0, 0, 0, 0, 0, 0}; std::vectorfloat probsGradHostData {0, 0, 0}; ret CreateAclTensor(permutedTokensHostData, permutedTokensShape, permutedTokensDeviceAddr, aclDataType::ACL_FLOAT, permutedTokens); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(unpermutedTokensGradHostData, unpermutedTokensGradShape, unpermutedTokensGradDeviceAddr, aclDataType::ACL_FLOAT, unpermutedTokensGrad); 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(sortedIndicesHostData, sortedIndicesShape, sortedIndicesDeviceAddr, aclDataType::ACL_INT32, sortedIndices); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(permutedTokensGradHostData, permutedTokensGradShape, permutedTokensGradDeviceAddr, aclDataType::ACL_FLOAT, permutedTokensGrad); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(probsGradHostData, probsGradShape, probsGradDeviceAddr, aclDataType::ACL_FLOAT, probsGrad); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的Api名称 uint64_t workspaceSize 0; aclOpExecutor *executor; // 调用aclnnMoeTokenUnpermuteWithEpGrad第一段接口 ret aclnnMoeTokenUnpermuteWithEpGradGetWorkspaceSize(unpermutedTokensGrad, sortedIndices,permutedTokens, probs, paddedMode, nullptr, nullptr, 3, permutedTokensGrad, probsGrad, workspaceSize, executor); CHECK_RET( ret ACL_SUCCESS, LOG_PRINT(aclnnMoeTokenUnpermuteWithEpGradGetWorkspaceSize 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); } // 调用aclnnMoeTokenUnpermuteWithEpGrad第二段接口 ret aclnnMoeTokenUnpermuteWithEpGrad(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMoeTokenUnpermuteWithEpGrad 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的接口定义修改 PrintOutResult(permutedTokensGradShape, permutedTokensGradDeviceAddr); PrintOutResult(probsGradShape, probsGradDeviceAddr); // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(permutedTokens); aclDestroyTensor(unpermutedTokensGrad); aclDestroyTensor(sortedIndices); aclDestroyTensor(probs); aclDestroyTensor(permutedTokensGrad); aclDestroyTensor(probsGrad); // 7. 释放device资源 aclrtFree(permutedTokensDeviceAddr); aclrtFree(unpermutedTokensGradDeviceAddr); aclrtFree(probsDeviceAddr); aclrtFree(sortedIndicesDeviceAddr); aclrtFree(permutedTokensGradDeviceAddr); aclrtFree(probsGradDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中的 shape 组合tokens_num1、topK3、hidden_size2对应公式验证逻辑sortedIndices {0, 1, 2}表明三个 token 位置的专家索引均落在有效范围内probs {1, 1, 1}时permutedTokensGradOut各位置应等于unpermutedTokensGrad * probs的外积展开值probsGradOut应等于unpermutedTokensGrad各元素与对应permutedTokens之积在 topK 维的累加。读者可据此手工推算并核对打印结果。十、测试与验证仓库为该算子提供了多级测试覆盖UT 测试test_moeTokenUnpermuteWithEpGrad_infershape.cpp验证 shape 与 dtype 推导规则覆盖正常 shape、未知 rank[-2]、probs 为空等分支test_moe_token_unpermute_with_ep_grad_tiling.cpp验证核间/核内切分参数的合理性包括 topK 上限、hiddenSize 切分、indices 预留数量对齐等test_moe_token_unpermute_with_ep_grad.cpp在 AICore 侧对 Kernel 做数值验证覆盖无 probs / 有 probs 两条路径及 EP 范围截断场景。ST 测试目录 tests/st/aclnnMoeTokenUnpermuteWithEpGrad 下提供了 ATK 用例描述文件atk_aclnnMoeTokenUnpermuteWithEpGrad.json与执行器脚本executor_aclnnMoeTokenUnpermuteWithEpGrad.py可在真实 NPU 环境上对 aclnn 接口做端到端验证。十一、总结aclnnMoeTokenUnpermuteWithEpGrad是 CANN ops-transformer 中面向 MoE 专家并行训练的反向传播算子通过sortedIndices与rangeOptional完成 EP 切分范围内的梯度逆散射并借助probs可选输入实现了permutedTokensGradOut与probsGradOut的双路梯度输出。其 aclnn 两段式接口、参数约束与错误码语义完整记录在 aclnnMoeTokenUnpermuteWithEpGrad.md 中若需进一步理解正向算子或查看同族 MoE 路由算子如MoeTokenUnpermute、MoeTokenPermute系列可参考 moe 目录 下的其余算子文档与源码。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer 算子解析aclnnMoeTokenPermuteGrad 接口详解与 MoE Token Permute 反向传播实现CANN ops transformer 算子解析aclnnMoeTokenPermuteGrad 接口详解与 MoE Token Permute 反向传播实算子库人工智能深度学习AscendCANN ops-transformer MoeTokenUnpermuteGrad 算子详解MoE Token Unpermute 反向传播的 aclnn 接口原理与实战CANN ops transformer MoeTokenUnpermuteGrad 算子详解MoE Token Unpermute 反向传播的 aclnn算子库人工智能深度学习AscendMoeTokenUnpermuteWithRoutingMap 算子全解析CANN ops-transformer 中 MoE Token 反置换与加权累加实现MoeTokenUnpermuteWithRoutingMap 算子全解析CANN ops transformer 中 MoE Token 反置换与加权累加实算子库人工智能深度学习Ascend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考