MoeTokenUnpermuteGrad 算子全解析aclnnMoeTokenUnpermuteGrad 两段式接口、梯度公式与 NPU 实现【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本篇技术指南聚焦 CANN ops-transformer 仓库中的MoeTokenUnpermuteGrad算子及其对外 aclnn 接口aclnnMoeTokenUnpermuteGrad。该算子是 MoEMixture of Experts架构中MoeTokenUnpermute正向算子将按专家重排的 token 恢复到原始顺序的反向传播算子用于 MoE 专家路由的反向训练场景。读完本文你将掌握该接口的函数原型、两段式调用流程、每个参数的约束与取值范围、梯度计算公式的数学含义以及其在 op_host 校验、tiling 切分、AscendC kernel 计算上的源码级实现细节并可直接套用仓库提供的完整 C 调用示例进行开发与调试。算子定位与背景在大模型 MoE 训练中token 通常会被路由dispatch到不同专家正向计算完成后再通过 unpermute 操作将 token 还原到原始顺序继续后续网络层计算。MoeTokenUnpermuteGrad正是这一还原过程的梯度算子它接收unpermutedTokensGrad正向 unpermute 输出对 loss 的梯度结合路由阶段记录的sortedIndices映射关系与可选的路由权重probs计算出需要回传的permutedTokensGradpermuted tokens 的梯度以及probsGrad路由权重的梯度。该算子位于仓库的 moe/moe_token_unpermute_grad 目录与其正向算子MoeTokenUnpermute、以及 moe/moe_token_unpermute_with_ep_grad带 expert parallel 的梯度版本形成算子族。从源码结构看该目录遵循 CANN 算子仓的标准四段式布局op_host算子定义、infershape、tiling、op_kernelAscendC 核函数、op_apiaclnn 接口封装、docs/examples/tests文档、样例与测试。功能说明与梯度计算公式接口功能为aclnnMoeTokenUnpermute的反向传播。根据 aclnnMoeTokenUnpermuteGrad.md 与 README.md计算公式按probs是否传入分为两种情况。probs 非 None 时的梯度计算正向过程可描述为unpermutedTokens[i] permutedTokens[sortedIndices[i]]即按照sortedIndices将 permuted token 重新排列得到形状为(tokens_num, topK_num, hidden_size)的unpermutedTokensreshape 后再与路由权重probs相乘。对应反向传播公式如下$$ unpermutedTokens unpermutedTokensGrad.unsqueeze(1) * unpermutedTokens $$$$ probsGrad \sum_{k0}^{K}(unpermutedTokens_{i,j,k}) $$$$ permutedTokensGrad[sortedIndices[i]] ((unpermutedTokensGrad.unsqueeze(1) * probs.unsqueeze(-1)).reshape(-1, hiddenSize))[i] $$即probsGrad通过对unpermutedTokensGrad与permutedTokens即正向 unpermute 的输入 token逐元素乘积累加沿 hidden 维做 reduce sum得到反映了路由权重对 loss 的敏感度permutedTokensGrad由unpermutedTokensGrad扩展出 topK 维与probs扩展出 hidden 维逐元素相乘再按sortedIndices索引 scatter 回 permuted token 的原始位置。probs 为 None 时的梯度计算当不传入路由权重时梯度计算退化为纯索引搬移$$ permutedTokensGrad[sortedIndices[i]] unpermutedTokensGrad[i] $$即仅按照sortedIndices将梯度 scatter 回对应位置此时topK_num 1。此时probsGrad输出不会被写入有效值。两段式接口与函数原型与 CANN 其他 aclnn 算子一致本算子采用两段式接口设计先调用aclnnMoeTokenUnpermuteGradGetWorkspaceSize完成入参校验、计算 workspace 大小并创建执行器再调用aclnnMoeTokenUnpermuteGrad真正下发计算。aclnnStatus aclnnMoeTokenUnpermuteGradGetWorkspaceSize( const aclTensor *permutedTokensOptional, const aclTensor *unpermutedTokensGrad, const aclTensor *sortedIndices, const aclTensor *probsOptional, bool paddedMode, const aclIntArray *restoreShapeOptional, aclTensor *permutedTokensGradOut, aclTensor *probsGradOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnMoeTokenUnpermuteGrad( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnMoeTokenUnpermuteGradGetWorkspaceSize 参数详解参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorpermutedTokensOptionalaclTensor输入表示输入 token正向 permute 后的 token对应公式中的permutedTokens-BFLOAT16、FLOAT16、FLOAT32ND(tokens_num * topK_num, hidden_size)√unpermutedTokensGradaclTensor输入unpermutedTokens 的梯度-与 permutedTokens 一致ND(tokens_num, hidden_size)√sortedIndicesaclTensor输入表示输入输出梯度的映射关系对应sortedIndices取值范围为 [0, tokens_num * topK_num - 1]且索引值不重复INT32ND(tokens_num * topK_num)√probsOptionalaclTensor输入表示 token 选择指定专家的权重当 probsOptional 不为空时topK_num 等于 probsOptional 第 2 维当 probsOptional 为空时topK_num1BFLOAT16、FLOAT16、FLOAT32ND(tokens_num, topK_num)√paddedModebool输入true 表示开启 paddedModefalse 表示关闭 paddedMode目前仅支持 false----restoreShapeOptionalaclIntArray输入当 paddedMode 为 true 后生效否则不会对其进行操作。当 paddedMode 为 true 时此为 unpermutedTokens 的 shape当前仅支持 nullptraclIntArray*---permutedTokensGradOutaclTensor输出permutedTokens 的梯度-与 permutedTokens 一致ND(tokens_num * topK_num, hidden_size)×probsGradOutaclTensor输出probs 的梯度-与 probsOptional 一致ND(tokens_num, topK_num)×workspaceSizeuint64_t输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor输出返回 op 执行器包含了算子计算流程-----其中“非连续 Tensor”列标注 √ 表示允许传入非连续非 contiguous张量。从 aclnn_moe_token_unpermute_grad.cpp 的源码可见接口内部会对所有输入统一执行l0op::Contiguous转为连续张量并通过ViewCopy将计算结果拷贝回非连续输出视图见该文件第 152-157 行这解释了为何输入支持非连续而输出标注为 ×。返回值与错误码aclnnStatus返回状态码具体参见aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入或输出的 Tensor 是空指针ACLNN_ERR_PARAM_INVALID161002输入和输出的数据类型不在支持的范围内在源码中这两类校验分别对应MoeTokenUnpermuteGradCheck::CheckNotNull空指针检查与CheckDtypeValid数据类型检查。CheckDtypeValid内部通过MOE_GRAD_DTYPE_SUPPORT_LIST_X仅允许 FLOAT16、BF16、FLOAT与MOE_GRAD_DTYPE_SUPPORT_LIST_ROW_IDX仅允许 INT32两个支持列表完成校验并额外约束unpermutedTokensGrad与permutedTokens、permutedTokensGradOut之间数据类型必须一致OP_CHECK_DTYPE_NOT_SAME。aclnnMoeTokenUnpermuteGrad 第二段接口参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnMoeTokenUnpermuteGradGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream约束说明确定性计算aclnnMoeTokenUnpermuteGrad默认确定性实现多次运行结果可复现。tokens_num 表示输入的 token 数量hidden_size 表示词向量维度。通过 paddedMode 区分两种模式paddedMode 等于 true 时每个专家固定能够处理 capacity 个 token当前版本未开放paddedMode 等于 false 时每个 token 固定被 topK_num 个专家处理这也是当前唯一支持的路径。topK_num 上限Atlas A2 训练系列产品/Atlas A2 推理系列产品、Atlas A3 训练系列产品/Atlas A3 推理系列产品上topK_num 512。该限制在 tiling 实现中同样有硬校验tokensNum 0 || topK 0 || hiddenSize 0 || numOutTokens 0或topK 512时直接返回GRAPH_FAILED。Ascend 950 转调说明Ascend 950PR/Ascend 950DT 上调用本接口时框架内部会转调用 aclnnMoeFinalizeRoutingV2Grad 接口如果出现参数错误提示请参考以下参数映射关系permutedTokensOptional 输入等同于 aclnnMoeFinalizeRoutingV2Grad 接口的 expandedXOptional 输入unpermutedTokensGrad 输入等同于 aclnnMoeFinalizeRoutingV2Grad 接口的 gradY 输入sortedIndices 输入等同于 aclnnMoeFinalizeRoutingV2Grad 接口的 expandedRowIdx 输入probsOptional 输入等同于 aclnnMoeFinalizeRoutingV2Grad 接口的 scalesOptional 输入paddedMode 输入等同于 aclnnMoeFinalizeRoutingV2Grad 接口的 dropPadMode 输入permutedTokensGradOut 输出等同于 aclnnMoeFinalizeRoutingV2Grad 接口的 gradExpandedXOut 输出probsGradOut 输出等同于 aclnnMoeFinalizeRoutingV2Grad 接口的 gradScalesOut 输出。这一转调机制同样体现在 op_api 源码中aclnnMoeTokenUnpermuteGradGetWorkspaceSize通过Ops::Transformer::AclnnUtil::IsRegbase()判断当前芯片架构非 regbase 架构走自动生成的aclnnInnerMoeTokenUnpermuteGradGetWorkspaceSizeregbase 架构Ascend 950则直接以l0op::MoeFinalizeRoutingV2Grad构建计算图并在此路径下额外拒绝paddedMode true返回ACLNN_ERR_PARAM_INVALID见 aclnn_moe_token_unpermute_grad.cpp。产品支持情况产品是否支持Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品支持Atlas A2 训练系列产品/Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持源码级实现剖析算子定义op_def在 moe_token_unpermute_grad_def.cpp 中算子注册为MoeTokenUnpermuteGrad输入permuted_tokens、probs均为OPTIONAL类型unpermuted_tokens_grad、sorted_indices为REQUIRED输出permuted_tokens_grad、probs_grad均为REQUIRED属性padded_mode为可选布尔属性默认falserestore_shape为可选 ListInt 属性默认{1, 1}所有输入声明了.AutoContiguous()与 op_api 层的 Contiguous 处理呼应编译配置AICore().AddConfig(ascend910b)与AddConfig(ascend910_93)对应op_host/config/ascend910b与ascend910_93两个目录下的 binary 与 simplified_key 配置。tiling 切分策略tiling 实现 分为核间切分与核内切分两步核间切分MoeTokenUnpermuteGradInitSplitInfo按 AIV 核数totalCoreNum对tokensNum均分前formerCoreNum个核各多处理 1 个 tokentokenNumEachCore tokenNumTailCore 1后续核处理tokenNumTailCore个每个核对应的 indices 数量为tokenNum * topK。核内切分根据probs是否存在走两条不同策略。CoreSplitInfoProbIsNone在probs为 None 时优先让 hidden 维全载入 UB剩余空间尽可能分满indicesNumPerLoop并在 UB 放不下时对 hidden 维按 32B 对齐切片循环CoreSplitInfoProbIsNotNone在存在probs时indicesReserveNum最大预留 256 个元素且必须是topK的整数倍INDICES_RESERVE_MAX_NUM 256同时保证inputReserveNum是topK的约数hidden 过大时按 512B 对齐切片。tiling 数据通过tilingKey区分四种组合0padded_modeFalse 且无 probs、1padded_modeFalse 且有 probs、10padded_modeTrue 且无 probs、11padded_modeTrue 且有 probs并在核函数中按TILING_KEY_IS分派。workspace 统一申请16 * 1024 * 102416MB的系统 workspacesysWorkspaceSize。AscendC kernel 计算核函数入口 moe_token_unpermute_grad.cpp 根据 tilingKey 实例化两个模板类MoeTokenUnpermuteGradProbNonemoe_token_unpermute_grad_prob_none.h对应probs为 None 的纯 scatter 场景。外层按 indices 数量循环搬入rowIdMapsortedIndices内层按 hidden 切片循环搬入unpermuted_tokens_grad的一块(indicesLoopNum, hidden)数据然后逐 index 读取映射值offset将数据DataCopyPad写到permuted_tokens_grad的offset * hiddenSize hiddenLoopOffset位置即实现公式permutedTokensGrad[sortedIndices[i]] unpermutedTokensGrad[i]。MoeTokenUnpermuteGradProbNotNonemoe_token_unpermute_grad_prob_not_none.h对应存在probs的完整梯度计算。该实现使用多块 UB 空间并借助EVENT_ID0~7事件做 V/MTE2/MTE3 之间的同步采用 ping-pong 双缓冲pingPongFlag*系列标志隐藏搬运与计算延迟将sortedIndices与probs通过GetValue预取到片上数组indicesArray/probsArray各最大 512 个元素避免频繁GetValue对permuted_tokens按inputReserveNum分批搬入 UB并处理截断场景inputLoopOffset numOutTokens时Duplicate填 0计算permutedTokensGradMuls(tmpBufferPermutedTokensGradFp32, tmpBufferUnpermutedFp32, probsArray[indicesIndex], hiddensizeLoopNum)即unpermutedTokensGrad * probs计算probsGrad先Mul得到permutedTokens * unpermutedTokensGrad再调用ReduceSumFunc沿 hidden 维做规约ReduceSumFunc定义于 moe_token_unpermute_grad_base.h按 hidden 长度分四档≥4096、≥512、≥64 用“二分累加BinaryAddFuncBlockReduceSumWholeReduceSum”组合64 直接WholeReduceSum规约结果逐 hidden 切片累加到tmpBufferProbGradReduceSumFp32后统一搬出FP16/BF16 输入统一Cast到 FP32 计算再Cast回原类型输出FP32 输入则用Copy保证中间计算精度。完整调用示例以下代码摘自 examples/test_aclnn_moe_token_unpermute_grad.cpp完整覆盖了资源初始化、张量构造、两段式调用、结果回拷与资源释放的整个生命周期。编译和执行过程可参考编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_moe_token_unpermute_grad.h #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的接口自定义构造 // 本示例对应 tokens_num1, topK_num3, hidden_size2 的场景 std::vectorint64_t permutedTokensShape {3, 2}; // (tokens_num*topK_num, hidden_size) std::vectorint64_t unpermutedTokensGradShape {1, 2}; // (tokens_num, hidden_size) std::vectorint64_t probsShape {1, 3}; // (tokens_num, topK_num) std::vectorint64_t sortedIndicesShape {3}; // (tokens_num*topK_num) 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; // 当前仅支持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 uint64_t workspaceSize 0; aclOpExecutor *executor; // 调用aclnnMoeTokenUnpermuteGrad第一段接口 // 注意restoreShapeOptional当前仅支持nullptr ret aclnnMoeTokenUnpermuteGradGetWorkspaceSize(permutedTokens, unpermutedTokensGrad, sortedIndices, probs, paddedMode, nullptr, permutedTokensGrad, probsGrad, workspaceSize, executor); CHECK_RET( ret ACL_SUCCESS, LOG_PRINT(aclnnMoeTokenUnpermuteGradGetWorkspaceSize 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); } // 调用aclnnMoeTokenUnpermuteGrad第二段接口 ret aclnnMoeTokenUnpermuteGrad(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMoeTokenUnpermuteGrad 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侧 PrintOutResult(permutedTokensGradShape, permutedTokensGradDeviceAddr); PrintOutResult(probsGradShape, probsGradDeviceAddr); // 6. 释放aclTensor和aclScalar 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; }示例中构造了tokens_num1, topK_num3, hidden_size2的极简场景sortedIndices {0,1,2}为顺序映射probs全 1、unpermutedTokensGrad全 1、permutedTokens全 1。按公式推算permutedTokensGrad每个元素应为unpermutedTokensGrad * probs 1probsGrad每个元素应为permutedTokens * unpermutedTokensGrad沿 hidden 维求和等于 2读者可将程序输出与这两个期望值比对验证。若传入非连续输入接口内部会先做 Contiguous 处理输出则按 ViewCopy 语义写回原始内存视图。测试与验证仓库为该算子提供了完整的测试资产可作为功能验证与二次开发的参考单算子调用样例examples/test_aclnn_moe_token_unpermute_grad.cpp即上文示例来源host 侧单测tests/ut/op_host 下的test_moe_token_unpermute_grad_infershape.cpp与test_moe_token_unpermute_grad_tiling.cpp分别验证 shape 推导与 tiling 数据op_api 单测tests/ut/op_host/op_api/test_aclnn_moe_token_unpermute_grad.cppkernel 单测tests/ut/op_kernel 下的test_moe_token_unpermute_grad.cpp与数据生成脚本gen_data.py、gen_tiling.pygolden 数据生成可参考 tests/assets/golden.pyST 测试tests/st/aclnnMoeTokenUnpermuteGrad 下的executor_aclnnMoeTokenUnpermuteGrad.py与atk_aclnnMoeTokenUnpermuteGrad.jsonATK 用例定义以及 tests/st/arch35/test_aclnn_moe_token_unpermute_grad_arch35.csv架构相关用例参数。小结aclnnMoeTokenUnpermuteGrad是 MoE 训练反向链路中的关键算子它通过sortedIndices完成梯度 scatter通过probs完成路由权重的梯度回传覆盖了有无路由权重两种典型场景。本文从接口原型、参数约束、计算公式、芯片差异Ascend 950 转调 MoeFinalizeRoutingV2Grad、tiling 切分与 AscendC kernel 的 ping-pong 双缓冲实现等多维度进行了剖析并给出了可直接运行的完整调用示例。开发者在集成该算子时重点关注paddedMode与restoreShapeOptional当前仅支持默认值、topK_num 512、sortedIndices取值范围与不重复性这三条约束即可避免绝大多数参数错误。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考