尧图网站建设 
CANN ops-transformer 算子详解MoeTokenPermuteWithRoutingMapGrad 反向传播算子原理与 aclnn 接口实践【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读MoeTokenPermuteWithRoutingMapGrad 是 CANN ops-transformer 算子库中 MoeTokenPermuteWithRoutingMap 正向算子的反向传播实现用于在 MoEMixture of Experts稀疏门控网络中将专家侧按 routing map 排列后的 token 梯度还原unpermute回原始 token 顺序并同时计算门控概率 probs 的梯度。本文从该算子的功能定义、计算公式、两段式 aclnn 接口、参数约束与源码实现五个维度展开帮助你理解如何在 Atlas 训练/推理系列产品上正确调用aclnnMoeTokenPermuteWithRoutingMapGrad并掌握其底层 tiling 与 kernel 计算原理。一、算子定位与产品支持情况1.1 算子功能定位MoeTokenPermuteWithRoutingMapGrad 实现了aclnnMoeTokenPermuteWithRoutingMap的反向传播。在 MoE 前向中token 会按照 routing maptoken 到 expert 的映射关系被打散permute到各专家反向时则需要把专家输出的梯度按原索引还原回 token 的原始排布并生成 probs门控权重的梯度供后续对 gate 网络做梯度更新。从仓库源码结构看该算子位于 moe/moe_token_permute_with_routing_map_grad 目录由 Host 侧算子定义、Tiling、aclnn 封装、Kernel 侧计算以及 ST/UT 测试等完整工程组成。1.2 产品支持情况以下支持矩阵与 README.md 及 aclnnMoeTokenPermuteWithRoutingMapGrad.md 保持一致产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×对应的算子定义文件 moe_token_permute_with_routing_map_grad_def.cpp 中同样通过AICore().AddConfig()声明了ascend910b、ascend910_93与ascend950三个平台的配置其中 ascend950 配置启用了动态编译、动态 shape 与动态 rank 支持。二、计算公式与梯度语义2.1 核心公式算子首先对 sortedIndices 做一次排序得到用于还原的索引对$$ permuteTokenId, outIndex sortedIndices.sort(dim-1) $$在 dropAndPaddropPaddedMode模式下定义每个专家的容量$$ capacity permutedTokenOutputGrad.size(0) / experts_num $$2.2 计算 probs 的梯度probs 不为 None 时先申请一个形状为(tokens_num, experts_num)的全零张量作为 probs 梯度容器$$ probsGradOutOptional zeros(tokens_num, experts_num) $$paddedMode 为 truedropAndPad 开启时直接按位置回填$$ probsGradOutOptional[sortedIndices[i], i/capacity] permutedProbsOutputGradOptional[i] $$paddedMode 为 false 时借助 routingMap 作为掩码做 masked scatter$$ probsGradOutOptional maskedscatter(probsGradOutOptional, routingMap, permutedProbsOutputGradOptional) $$2.3 计算 tokens 的梯度probs 为 None 时申请与原始 token 张量同形状的全零梯度张量然后按索引累加还原$$ tokensGradOut zeros(restoreShape, dtypepermutedTokens.dtype, devicepermutedTokens.device) $$$$ tokensGradOut[permuteTokenId[i]] permutedTokens[outIndex[i]] $$注意这里是累加语义即一个原始 token 位置可能被多个 permuted 槽位贡献梯度kernel 中通过 index add 类操作实现。2.4 从源码看两种模式的实现差异在 aclnn_moe_token_permute_with_routing_map_grad.cpp 中Host 侧根据dropAndPad分支选择不同流程dropAndPad false调用ProcessNonDropAndPadGradientstokens 梯度走MoeTokenPermuteWithRoutingMapGrad内核probs 梯度则通过Transpose Cast MaskedScatter Transpose组合完成见 ProcessNonDropAndPadGradients。dropAndPad true分别调用ProcessDropAndPadProbsGrad与ProcessDropAndPadTokensGrad。后者在 Atlas A2/A3 及昇腾 950 架构上若 index 为 int32 且满足一定条件会先对 sortedIndices 执行稳定排序再走InplaceIndexAddWithSorted完成带确定性保证的梯度累加见 ProcessDropAndPadTokensGrad。Kernel 侧入口 moe_token_permute_with_routing_map_grad.cpp 通过 tilingKey 分发TILING_KEY_IS(0)走KernelMoeTokenUnpermuteprobs 恒为 false 的 unpermute 还原TILING_KEY_IS(1000)即TILINGKEY_DROPPAD走KernelMoeTokenPermuteWithRoutingMap处理 dropPad 模式下的 probs 梯度。三、两段式 aclnn 接口与函数原型该算子遵循 CANN 两段式接口规范详见 两段式接口说明先调用GetWorkspaceSize接口完成入参校验、构建算子计算流程并获取 workspace 大小再调用执行接口真正下发计算。3.1 第一段接口获取 workspace 与执行器aclnnStatus aclnnMoeTokenPermuteWithRoutingMapGradGetWorkspaceSize( const aclTensor *permutedTokensOutputGrad, const aclTensor *permutedProbsOutputGradOptional, const aclTensor *sortedIndices, const aclTensor *routingMapOptional, int64_t experts_num, int64_t tokens_num, bool dropAndPad, aclTensor *tokenGradOut, aclTensor *probsGradOutOptional, uint64_t *workspaceSize, aclOpExecutor **executor)3.2 第二段接口执行计算aclnnStatus aclnnMoeTokenPermuteWithRoutingMapGrad( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)四、参数说明4.1 第一段接口参数明细以下参数表同时覆盖了 README 的算子级参数定义与 aclnn 接口文档 的接口级细节参数名输入/输出描述使用说明数据类型数据格式维度shapepermutedTokensOutputGrad输入正向输出 permutedTokens 的梯度仅支持 2D不支持空 tensortopK_num 表示每个 token 最多被选中的专家数上限capacity 表示每个专家选中的 token 数BFLOAT16、FLOAT16、FLOAT32ND非 dropAndPad(tokens_num * topK_num, hidden_size)dropAndPad(experts_num * capacity, hidden_size)permutedProbsOutputGradOptional可选输入正向输出 permutedProbs 的梯度不传则表示不需要计算 probsGradOutOptional1D 张量数据类型须与 permutedTokensOutputGrad 一致或当其为 BFLOAT16 时支持 FLOATBFLOAT16、FLOAT16、FLOAT32ND非 dropAndPad(tokens_num * topK_num)dropAndPad(experts_num * capacity)sortedIndices输入排序后的索引值非 dropAndPad 模式取值[0, tokens_num * topK_num - 1]或 -1-1 表示无效槽位不参与计算dropAndPad 模式取值[0, experts_num * capacity - 1]INT32ND非 dropAndPad(tokens_num * topK_num,)dropAndPad(experts_num * capacity,)routingMapOptional可选输入token 到 expert 的映射关系2D 张量非 dropAndPad 模式要求每行不超过 topK 个 true/1INT8 取值 0/1bool 取值 true/falseINT8、boolND(tokens_num, experts_num)experts_num输入参与运算的专家个数需大于 0INT64--tokens_num输入参与运算的 token 个数需大于 0INT64--dropAndPad输入true 开启 dropPaddedModefalse 关闭决定走 dropPad 还是非 dropPad 计算分支bool--tokenGradOut输出输入 permutedTokens 的梯度2D 张量BFLOAT16、FLOAT16、FLOAT32ND(tokens_num, hidden_size)probsGradOutOptional可选输出正向算子输入 probs 的梯度2D 张量数据类型同 permutedProbsOutputGradOptional同 permutedProbsOutputGradOptionalND(tokens_num, experts_num)workspaceSize输出需要在 Device 侧申请的 workspace 大小由第一段接口计算---executor输出包含算子计算流程的执行器供第二段接口使用---4.2 参数校验规则源码佐证aclnn_moe_token_permute_with_routing_map_grad.cpp 中的CheckParams依次完成四类校验Attr 合法性experts_num 0、tokens_num 0否则返回ACLNN_ERR_PARAM_INVALID见 checkAttrValid。空指针检查permutedTokenOutputGrad、sortedIndices、tokensGradOut为必选不可为空同时若routingMapOptional nullptr而permutedProbsOutputGradOptional ! nullptr会直接报错因为计算 probs 梯度依赖 routing map见 CheckNotNull。数据类型检查permutedTokens 侧支持 FLOAT16/FLOAT/BF16probs 梯度支持与 tokens 同类型或 tokens 为 BF16 时 probs 为 FLOAT混合场景routingMap 仅支持 INT8/boolsortedIndices 仅支持 INT32见 CheckDtypeValid。Shape 检查permutedTokenOutputGrad必须 2D、routingMap必须 2D、permutedProbsOutputGradOptional必须 1D见 CheckShapeValid。4.3 第一段接口返回值返回aclnnStatus状态码完整取值见 aclnn 返回码说明。第一段接口完成入参校验典型报错场景如下返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001必选输入、必选输出或必选属性传入空指针ACLNN_ERR_PARAM_INVALID161002输入/输出的数据类型、数据格式或 Shape 不在支持范围内4.4 第二段接口参数说明参数名输入/输出描述workspace输入Device 侧申请的 workspace 内存地址workspaceSize输入workspace 大小由第一段接口获取executor输入包含算子计算流程的 op 执行器stream输入指定执行任务的 Stream五、约束说明使用该算子前需要确认以下约束来自 README.md 与 aclnn 接口文档topK 上限非 dropPaddedMode 场景下topK_num 512。该约束在 tiling 实现 中以MAX_TOPK_NUM 512常量强校验超出会返回GRAPH_FAILED。不支持混合精度输入permutedTokenOutputGrad、permutedProbsOutputGradOptional、tokensGradOut、probsGradOutOptional需保持相同数据类型。唯一的混合场景例外是permutedTokenOutputGrad为 BFLOAT16 时probs 相关张量允许使用 FLOAT。确定性计算该算子默认是非确定性实现可通过aclrtCtxSetSysParamOpt开启确定性模式开启后 dropAndPad 分支会自动切换为带稳定排序的InplaceIndexAddWithSorted路径见源码 ProcessDropAndPadTokensGrad。空 tensor若permutedTokenOutputGrad或tokensGradOut为空 tensor第一段接口直接返回 workspaceSize 0见 aclnn 实现。六、Tiling 与 Kernel 实现原理6.1 Host 侧 Tiling 策略Tiling 入口位于 moe_token_permute_with_routing_map_grad_tiling.cpp根据padded_mode属性分流非 dropPad 模式复用 unpermute 的 tiling 逻辑源码注释明确permute_with_routing_map_grad 复用 unpermute 代码。TilingCompute依次完成Init读取平台核心数、UB 内存、输入 shape 与 dtype 大小、SetCoreNum按 tokensNum 与最大 AIV 核数取小、TilingHiddenSizehidden 维按 UB 可容纳大小切片512B 对齐、SetBufferNum双缓冲至多 4 buffer 自适应与TilingTokentoken 维按核数均匀切分并处理余数最终通过SetBlockDim(usedCoreNum)设置核数并写入 tilingData见 TilingCompute。dropPad 模式走TilingComputeDropPad从permuted_token_output_grad第 0 维除以experts_num推出capacity按(experts_num * capacity)与最大核数比较决定单核或全核并行并把capacity/expertNum/tokenNum/singleCoreLen/lastCoreLen写入 tilingDatatilingKey 固定为TILINGKEY_DROPPAD 1000见 TilingComputeDropPad。另外两种模式均会向系统申请 16MB 的 workspace见 Init。6.2 Kernel 侧计算流程Kernel 入口 moe_token_permute_with_routing_map_grad.cpp 通过GET_TILING_DATA_MEMBER读取两种 tiling 结构并按 tilingKey 分发到KernelMoeTokenUnpermute或KernelMoeTokenPermuteWithRoutingMap。以非 dropPad 的 tokens 梯度还原为例moe_token_unpermute_routing_map_grad.h 中的Process按 token 分片迭代CalPartOutToken对每个输出 token 遍历其 topK 个槽位读取sortedIndices中该槽位的索引-1 作为哨兵值直接跳过无效槽位不参与计算见 CalPartOutToken命中有效索引的槽位将其 hidden 维数据搬入DataCopy/DataCopyPad非 FLOAT 类型先Cast到 FLOAT 累加最终写回输出CopyOut时非 FLOAT 再Cast回原类型见 CopyOut。6.3 Shape 推导infershape 实现 根据permuted_token_output_grad第 1 维得到 hidden_size用第 0 维除以 tokens_num 得到 topK_num据此推导输出 shapetokensGradOut (tokensNum, hidden_size)probsGradOutOptional (tokensNum, topk_num)输出数据类型与输入 0 保持一致。七、调用示例与编译运行7.1 完整调用流程仓库提供可直接参考的完整样例 test_aclnn_moe_token_permute_with_routing_map_grad.cpp其调用骨架如下#include aclnnop/aclnn_moe_token_permute_with_routing_map_grad.h #include acl/acl.h #include aclnn/acl_meta.h // 1. 固定写法acl 初始化、设置 device、创建 stream int32_t deviceId 0; aclrtStream stream; aclInit(nullptr); aclrtSetDevice(deviceId); aclrtCreateStream(stream); // 2. 构造输入输出示意dropAndPadtrue 场景 int64_t num_token 4096; int64_t hidden_size 7168; int64_t num_expert 256; int64_t num_capacity 16; // permuted_output_grad: shape (num_expert * num_capacity, hidden_size) // permutedProbsOutputGradOptional: shape (num_expert * num_capacity,) // sortedIndices: shape (num_expert * num_capacity,)INT32 // routingMapOptional: shape (num_token, num_expert)INT8/bool // out (tokensGradOut): shape (num_token, hidden_size) // out2 (probsGradOutOptional): shape (num_token, num_expert) // 上述 aclTensor 均通过 aclCreateTensor 创建并已用 aclrtMemcpy 拷入 device 侧 // 3. 调用第一段接口获取 workspaceSize 与 executor uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; aclnnMoeTokenPermuteWithRoutingMapGradGetWorkspaceSize( permuted_output_grad, ppermutedProbsOutputGrad, sortedIndices, proutingMapOptional, num_expert, num_token, true, out, out2, workspaceSize, executor); // 4. 按 workspaceSize 申请 device 内存 void *workspaceAddr nullptr; if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 5. 调用第二段接口执行计算 aclnnMoeTokenPermuteWithRoutingMapGrad(workspaceAddr, workspaceSize, executor, stream); // 6. 同步等待任务执行结束 aclrtSynchronizeStream(stream); // 7. 结果拷贝回 host 并打印aclrtMemcpy ACL_MEMCPY_DEVICE_TO_HOST // 8. 释放资源aclDestroyTensor、aclrtFree、aclrtDestroyStream、aclrtResetDevice、aclFinalize样例中给出了 dropAndPad 模式的一组典型尺寸num_token 4096、hidden_size 7168、num_expert 256、num_capacity 16此时permuted_output_grad形状为(256*16, 7168)routingMapOptional形状为(4096, 256)。若你使用非 dropAndPad 模式则 permuted 输入的第 0 维应为tokens_num * topK_num且满足topK_num 512。完整的编译与运行方式请参考 编译与运行样例。7.2 测试与验证仓库为该算子提供了三层测试Host 侧单测infershape 与 tiling 单测位于 tests/ut/op_hostop_api 层单测见 test_aclnn_moe_token_permute_with_routing_map_grad.cppKernel 侧单测tests/ut/op_kernel 下的test_moe_token_permute_with_routing_map_grad.cpp验证向量级计算结果ST 系统测试tests/st/aclnnMoeTokenPermuteWithRoutingMapGrad 提供 ATK 用例描述文件atk_aclnnMoeTokenPermuteWithRoutingMapGrad.json与执行脚本executor_aclnnMoeTokenPermuteWithRoutingMapGrad.py覆盖端到端调用。7.3 平台二进制配置算子针对不同平台预编译了多种 dtype 组合的 kernel 二进制配置见 ascend910b 二进制清单ascend910_93下同构。该清单覆盖 bf16/fp16/fp32 与 routing map 为 int8/bool 的 6 种组合外加 tokensbf16、probsfp32 的两种混合精度组合*_mix共 8 个二进制条目每个条目同时声明了num_experts默认 1、tokens_num默认 0与padded_mode默认空三个属性与算子定义一一对应。八、总结MoeTokenPermuteWithRoutingMapGrad 是 MoE 反向链路中负责梯度还原 门控梯度回传的关键算子核心设计可归纳为三点双模式统一接口通过dropAndPad属性在 dropPaddedMode 与非 dropPaddedMode 之间切换两种模式分别对应带 capacity 的稠密排列还原与基于 routingMap 掩码的稀疏 scatter 还原两段式 aclnn 编程模型GetWorkspaceSize完成校验与计算图构建aclnnMoeTokenPermuteWithRoutingMapGrad完成实际执行配合 workspace 机制在 Host 侧组合出 transpose/masked_scatter/index_add 等原子能力确定性与性能兼顾默认非确定性实现开启aclrtCtxSetSysParamOpt确定性后自动切换稳定排序 InplaceIndexAddWithSorted路径同时通过 hidden/token 双维 tiling 与多 buffer 流水充分利用 AIV 多核能力。实际开发中建议优先参考 接口文档 核对 shape 与 dtype 约束再对照 调用示例 搭建调用骨架最后用仓库内的 ST/UT 用例做结果对齐验证。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-transformer 算子详解：MoeTokenPermuteWithRoutingMapGrad 反向传播算子原理与 aclnn 接口实践
CANN ops-transformer 算子详解MoeTokenPermuteWithRoutingMapGrad 反向传播算子原理与 aclnn 接口实践【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读MoeTokenPermuteWithRoutingMapGrad 是 CANN ops-transformer 算子库中 MoeTokenPermuteWithRoutingMap 正向算子的反向传播实现用于在 MoEMixture of Experts稀疏门控网络中将专家侧按 routing map 排列后的 token 梯度还原unpermute回原始 token 顺序并同时计算门控概率 probs 的梯度。本文从该算子的功能定义、计算公式、两段式 aclnn 接口、参数约束与源码实现五个维度展开帮助你理解如何在 Atlas 训练/推理系列产品上正确调用aclnnMoeTokenPermuteWithRoutingMapGrad并掌握其底层 tiling 与 kernel 计算原理。一、算子定位与产品支持情况1.1 算子功能定位MoeTokenPermuteWithRoutingMapGrad 实现了aclnnMoeTokenPermuteWithRoutingMap的反向传播。在 MoE 前向中token 会按照 routing maptoken 到 expert 的映射关系被打散permute到各专家反向时则需要把专家输出的梯度按原索引还原回 token 的原始排布并生成 probs门控权重的梯度供后续对 gate 网络做梯度更新。从仓库源码结构看该算子位于 moe/moe_token_permute_with_routing_map_grad 目录由 Host 侧算子定义、Tiling、aclnn 封装、Kernel 侧计算以及 ST/UT 测试等完整工程组成。1.2 产品支持情况以下支持矩阵与 README.md 及 aclnnMoeTokenPermuteWithRoutingMapGrad.md 保持一致产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×对应的算子定义文件 moe_token_permute_with_routing_map_grad_def.cpp 中同样通过AICore().AddConfig()声明了ascend910b、ascend910_93与ascend950三个平台的配置其中 ascend950 配置启用了动态编译、动态 shape 与动态 rank 支持。二、计算公式与梯度语义2.1 核心公式算子首先对 sortedIndices 做一次排序得到用于还原的索引对$$ permuteTokenId, outIndex sortedIndices.sort(dim-1) $$在 dropAndPaddropPaddedMode模式下定义每个专家的容量$$ capacity permutedTokenOutputGrad.size(0) / experts_num $$2.2 计算 probs 的梯度probs 不为 None 时先申请一个形状为(tokens_num, experts_num)的全零张量作为 probs 梯度容器$$ probsGradOutOptional zeros(tokens_num, experts_num) $$paddedMode 为 truedropAndPad 开启时直接按位置回填$$ probsGradOutOptional[sortedIndices[i], i/capacity] permutedProbsOutputGradOptional[i] $$paddedMode 为 false 时借助 routingMap 作为掩码做 masked scatter$$ probsGradOutOptional maskedscatter(probsGradOutOptional, routingMap, permutedProbsOutputGradOptional) $$2.3 计算 tokens 的梯度probs 为 None 时申请与原始 token 张量同形状的全零梯度张量然后按索引累加还原$$ tokensGradOut zeros(restoreShape, dtypepermutedTokens.dtype, devicepermutedTokens.device) $$$$ tokensGradOut[permuteTokenId[i]] permutedTokens[outIndex[i]] $$注意这里是累加语义即一个原始 token 位置可能被多个 permuted 槽位贡献梯度kernel 中通过 index add 类操作实现。2.4 从源码看两种模式的实现差异在 aclnn_moe_token_permute_with_routing_map_grad.cpp 中Host 侧根据dropAndPad分支选择不同流程dropAndPad false调用ProcessNonDropAndPadGradientstokens 梯度走MoeTokenPermuteWithRoutingMapGrad内核probs 梯度则通过Transpose Cast MaskedScatter Transpose组合完成见 ProcessNonDropAndPadGradients。dropAndPad true分别调用ProcessDropAndPadProbsGrad与ProcessDropAndPadTokensGrad。后者在 Atlas A2/A3 及昇腾 950 架构上若 index 为 int32 且满足一定条件会先对 sortedIndices 执行稳定排序再走InplaceIndexAddWithSorted完成带确定性保证的梯度累加见 ProcessDropAndPadTokensGrad。Kernel 侧入口 moe_token_permute_with_routing_map_grad.cpp 通过 tilingKey 分发TILING_KEY_IS(0)走KernelMoeTokenUnpermuteprobs 恒为 false 的 unpermute 还原TILING_KEY_IS(1000)即TILINGKEY_DROPPAD走KernelMoeTokenPermuteWithRoutingMap处理 dropPad 模式下的 probs 梯度。三、两段式 aclnn 接口与函数原型该算子遵循 CANN 两段式接口规范详见 两段式接口说明先调用GetWorkspaceSize接口完成入参校验、构建算子计算流程并获取 workspace 大小再调用执行接口真正下发计算。3.1 第一段接口获取 workspace 与执行器aclnnStatus aclnnMoeTokenPermuteWithRoutingMapGradGetWorkspaceSize( const aclTensor *permutedTokensOutputGrad, const aclTensor *permutedProbsOutputGradOptional, const aclTensor *sortedIndices, const aclTensor *routingMapOptional, int64_t experts_num, int64_t tokens_num, bool dropAndPad, aclTensor *tokenGradOut, aclTensor *probsGradOutOptional, uint64_t *workspaceSize, aclOpExecutor **executor)3.2 第二段接口执行计算aclnnStatus aclnnMoeTokenPermuteWithRoutingMapGrad( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)四、参数说明4.1 第一段接口参数明细以下参数表同时覆盖了 README 的算子级参数定义与 aclnn 接口文档 的接口级细节参数名输入/输出描述使用说明数据类型数据格式维度shapepermutedTokensOutputGrad输入正向输出 permutedTokens 的梯度仅支持 2D不支持空 tensortopK_num 表示每个 token 最多被选中的专家数上限capacity 表示每个专家选中的 token 数BFLOAT16、FLOAT16、FLOAT32ND非 dropAndPad(tokens_num * topK_num, hidden_size)dropAndPad(experts_num * capacity, hidden_size)permutedProbsOutputGradOptional可选输入正向输出 permutedProbs 的梯度不传则表示不需要计算 probsGradOutOptional1D 张量数据类型须与 permutedTokensOutputGrad 一致或当其为 BFLOAT16 时支持 FLOATBFLOAT16、FLOAT16、FLOAT32ND非 dropAndPad(tokens_num * topK_num)dropAndPad(experts_num * capacity)sortedIndices输入排序后的索引值非 dropAndPad 模式取值[0, tokens_num * topK_num - 1]或 -1-1 表示无效槽位不参与计算dropAndPad 模式取值[0, experts_num * capacity - 1]INT32ND非 dropAndPad(tokens_num * topK_num,)dropAndPad(experts_num * capacity,)routingMapOptional可选输入token 到 expert 的映射关系2D 张量非 dropAndPad 模式要求每行不超过 topK 个 true/1INT8 取值 0/1bool 取值 true/falseINT8、boolND(tokens_num, experts_num)experts_num输入参与运算的专家个数需大于 0INT64--tokens_num输入参与运算的 token 个数需大于 0INT64--dropAndPad输入true 开启 dropPaddedModefalse 关闭决定走 dropPad 还是非 dropPad 计算分支bool--tokenGradOut输出输入 permutedTokens 的梯度2D 张量BFLOAT16、FLOAT16、FLOAT32ND(tokens_num, hidden_size)probsGradOutOptional可选输出正向算子输入 probs 的梯度2D 张量数据类型同 permutedProbsOutputGradOptional同 permutedProbsOutputGradOptionalND(tokens_num, experts_num)workspaceSize输出需要在 Device 侧申请的 workspace 大小由第一段接口计算---executor输出包含算子计算流程的执行器供第二段接口使用---4.2 参数校验规则源码佐证aclnn_moe_token_permute_with_routing_map_grad.cpp 中的CheckParams依次完成四类校验Attr 合法性experts_num 0、tokens_num 0否则返回ACLNN_ERR_PARAM_INVALID见 checkAttrValid。空指针检查permutedTokenOutputGrad、sortedIndices、tokensGradOut为必选不可为空同时若routingMapOptional nullptr而permutedProbsOutputGradOptional ! nullptr会直接报错因为计算 probs 梯度依赖 routing map见 CheckNotNull。数据类型检查permutedTokens 侧支持 FLOAT16/FLOAT/BF16probs 梯度支持与 tokens 同类型或 tokens 为 BF16 时 probs 为 FLOAT混合场景routingMap 仅支持 INT8/boolsortedIndices 仅支持 INT32见 CheckDtypeValid。Shape 检查permutedTokenOutputGrad必须 2D、routingMap必须 2D、permutedProbsOutputGradOptional必须 1D见 CheckShapeValid。4.3 第一段接口返回值返回aclnnStatus状态码完整取值见 aclnn 返回码说明。第一段接口完成入参校验典型报错场景如下返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001必选输入、必选输出或必选属性传入空指针ACLNN_ERR_PARAM_INVALID161002输入/输出的数据类型、数据格式或 Shape 不在支持范围内4.4 第二段接口参数说明参数名输入/输出描述workspace输入Device 侧申请的 workspace 内存地址workspaceSize输入workspace 大小由第一段接口获取executor输入包含算子计算流程的 op 执行器stream输入指定执行任务的 Stream五、约束说明使用该算子前需要确认以下约束来自 README.md 与 aclnn 接口文档topK 上限非 dropPaddedMode 场景下topK_num 512。该约束在 tiling 实现 中以MAX_TOPK_NUM 512常量强校验超出会返回GRAPH_FAILED。不支持混合精度输入permutedTokenOutputGrad、permutedProbsOutputGradOptional、tokensGradOut、probsGradOutOptional需保持相同数据类型。唯一的混合场景例外是permutedTokenOutputGrad为 BFLOAT16 时probs 相关张量允许使用 FLOAT。确定性计算该算子默认是非确定性实现可通过aclrtCtxSetSysParamOpt开启确定性模式开启后 dropAndPad 分支会自动切换为带稳定排序的InplaceIndexAddWithSorted路径见源码 ProcessDropAndPadTokensGrad。空 tensor若permutedTokenOutputGrad或tokensGradOut为空 tensor第一段接口直接返回 workspaceSize 0见 aclnn 实现。六、Tiling 与 Kernel 实现原理6.1 Host 侧 Tiling 策略Tiling 入口位于 moe_token_permute_with_routing_map_grad_tiling.cpp根据padded_mode属性分流非 dropPad 模式复用 unpermute 的 tiling 逻辑源码注释明确permute_with_routing_map_grad 复用 unpermute 代码。TilingCompute依次完成Init读取平台核心数、UB 内存、输入 shape 与 dtype 大小、SetCoreNum按 tokensNum 与最大 AIV 核数取小、TilingHiddenSizehidden 维按 UB 可容纳大小切片512B 对齐、SetBufferNum双缓冲至多 4 buffer 自适应与TilingTokentoken 维按核数均匀切分并处理余数最终通过SetBlockDim(usedCoreNum)设置核数并写入 tilingData见 TilingCompute。dropPad 模式走TilingComputeDropPad从permuted_token_output_grad第 0 维除以experts_num推出capacity按(experts_num * capacity)与最大核数比较决定单核或全核并行并把capacity/expertNum/tokenNum/singleCoreLen/lastCoreLen写入 tilingDatatilingKey 固定为TILINGKEY_DROPPAD 1000见 TilingComputeDropPad。另外两种模式均会向系统申请 16MB 的 workspace见 Init。6.2 Kernel 侧计算流程Kernel 入口 moe_token_permute_with_routing_map_grad.cpp 通过GET_TILING_DATA_MEMBER读取两种 tiling 结构并按 tilingKey 分发到KernelMoeTokenUnpermute或KernelMoeTokenPermuteWithRoutingMap。以非 dropPad 的 tokens 梯度还原为例moe_token_unpermute_routing_map_grad.h 中的Process按 token 分片迭代CalPartOutToken对每个输出 token 遍历其 topK 个槽位读取sortedIndices中该槽位的索引-1 作为哨兵值直接跳过无效槽位不参与计算见 CalPartOutToken命中有效索引的槽位将其 hidden 维数据搬入DataCopy/DataCopyPad非 FLOAT 类型先Cast到 FLOAT 累加最终写回输出CopyOut时非 FLOAT 再Cast回原类型见 CopyOut。6.3 Shape 推导infershape 实现 根据permuted_token_output_grad第 1 维得到 hidden_size用第 0 维除以 tokens_num 得到 topK_num据此推导输出 shapetokensGradOut (tokensNum, hidden_size)probsGradOutOptional (tokensNum, topk_num)输出数据类型与输入 0 保持一致。七、调用示例与编译运行7.1 完整调用流程仓库提供可直接参考的完整样例 test_aclnn_moe_token_permute_with_routing_map_grad.cpp其调用骨架如下#include aclnnop/aclnn_moe_token_permute_with_routing_map_grad.h #include acl/acl.h #include aclnn/acl_meta.h // 1. 固定写法acl 初始化、设置 device、创建 stream int32_t deviceId 0; aclrtStream stream; aclInit(nullptr); aclrtSetDevice(deviceId); aclrtCreateStream(stream); // 2. 构造输入输出示意dropAndPadtrue 场景 int64_t num_token 4096; int64_t hidden_size 7168; int64_t num_expert 256; int64_t num_capacity 16; // permuted_output_grad: shape (num_expert * num_capacity, hidden_size) // permutedProbsOutputGradOptional: shape (num_expert * num_capacity,) // sortedIndices: shape (num_expert * num_capacity,)INT32 // routingMapOptional: shape (num_token, num_expert)INT8/bool // out (tokensGradOut): shape (num_token, hidden_size) // out2 (probsGradOutOptional): shape (num_token, num_expert) // 上述 aclTensor 均通过 aclCreateTensor 创建并已用 aclrtMemcpy 拷入 device 侧 // 3. 调用第一段接口获取 workspaceSize 与 executor uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; aclnnMoeTokenPermuteWithRoutingMapGradGetWorkspaceSize( permuted_output_grad, ppermutedProbsOutputGrad, sortedIndices, proutingMapOptional, num_expert, num_token, true, out, out2, workspaceSize, executor); // 4. 按 workspaceSize 申请 device 内存 void *workspaceAddr nullptr; if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 5. 调用第二段接口执行计算 aclnnMoeTokenPermuteWithRoutingMapGrad(workspaceAddr, workspaceSize, executor, stream); // 6. 同步等待任务执行结束 aclrtSynchronizeStream(stream); // 7. 结果拷贝回 host 并打印aclrtMemcpy ACL_MEMCPY_DEVICE_TO_HOST // 8. 释放资源aclDestroyTensor、aclrtFree、aclrtDestroyStream、aclrtResetDevice、aclFinalize样例中给出了 dropAndPad 模式的一组典型尺寸num_token 4096、hidden_size 7168、num_expert 256、num_capacity 16此时permuted_output_grad形状为(256*16, 7168)routingMapOptional形状为(4096, 256)。若你使用非 dropAndPad 模式则 permuted 输入的第 0 维应为tokens_num * topK_num且满足topK_num 512。完整的编译与运行方式请参考 编译与运行样例。7.2 测试与验证仓库为该算子提供了三层测试Host 侧单测infershape 与 tiling 单测位于 tests/ut/op_hostop_api 层单测见 test_aclnn_moe_token_permute_with_routing_map_grad.cppKernel 侧单测tests/ut/op_kernel 下的test_moe_token_permute_with_routing_map_grad.cpp验证向量级计算结果ST 系统测试tests/st/aclnnMoeTokenPermuteWithRoutingMapGrad 提供 ATK 用例描述文件atk_aclnnMoeTokenPermuteWithRoutingMapGrad.json与执行脚本executor_aclnnMoeTokenPermuteWithRoutingMapGrad.py覆盖端到端调用。7.3 平台二进制配置算子针对不同平台预编译了多种 dtype 组合的 kernel 二进制配置见 ascend910b 二进制清单ascend910_93下同构。该清单覆盖 bf16/fp16/fp32 与 routing map 为 int8/bool 的 6 种组合外加 tokensbf16、probsfp32 的两种混合精度组合*_mix共 8 个二进制条目每个条目同时声明了num_experts默认 1、tokens_num默认 0与padded_mode默认空三个属性与算子定义一一对应。八、总结MoeTokenPermuteWithRoutingMapGrad 是 MoE 反向链路中负责梯度还原 门控梯度回传的关键算子核心设计可归纳为三点双模式统一接口通过dropAndPad属性在 dropPaddedMode 与非 dropPaddedMode 之间切换两种模式分别对应带 capacity 的稠密排列还原与基于 routingMap 掩码的稀疏 scatter 还原两段式 aclnn 编程模型GetWorkspaceSize完成校验与计算图构建aclnnMoeTokenPermuteWithRoutingMapGrad完成实际执行配合 workspace 机制在 Host 侧组合出 transpose/masked_scatter/index_add 等原子能力确定性与性能兼顾默认非确定性实现开启aclrtCtxSetSysParamOpt确定性后自动切换稳定排序 InplaceIndexAddWithSorted路径同时通过 hidden/token 双维 tiling 与多 buffer 流水充分利用 AIV 多核能力。实际开发中建议优先参考 接口文档 核对 shape 与 dtype 约束再对照 调用示例 搭建调用骨架最后用仓库内的 ST/UT 用例做结果对齐验证。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考