CANN ops-transformer MoeTokenPermuteWithEp 算子深度解析MoE EP 场景下的 Token Permute 实现与 aclnn 接口调用指南【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformerMoeTokenPermuteWithEp 是 CANN ops-transformer 中面向 MoEMixture of Experts专家并行EP场景的 permute 算子它根据专家索引 indices 将 tokens 与可选的 probs 广播、排序并按 rangeOptional 指定的 EP 有效范围切片输出。本文基于仓库中的 README.md 与 aclnnMoeTokenPermuteWithEp.md结合 Host 侧算子定义、Tiling 逻辑与 Kernel 实现源码完整讲解算子功能、参数语义、约束、两段式 aclnn 调用方式及源码级实现原理帮助你直接上手在 Atlas A2/A3 系列产品上完成该算子的开发与验证。产品支持情况MoeTokenPermuteWithEp 算子对产品的支持情况如下与 aclnnMoeTokenPermuteWithEp.md 及算子定义中AICore().AddConfig声明一致见 moe_token_permute_with_ep_def.cpp产品是否支持Ascend 950PR/Ascend 950DT×Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×即该算子仅面向 Atlas A2 与 Atlas A3 训练/推理系列产品提供算力配置ascend910b与ascend910_93。功能说明EP 场景下的 Token Permute 计算MoE 模型中稀疏门控网络会将 token 路由到不同的专家expert。当采用专家并行EP时每个设备只持有部分专家需要把输入 token 按照其专家索引重新排列permute使属于本设备所辖专家的 token 聚拢到一起才能喂给本地专家计算计算完成后再通过 unpermute 还原顺序。MoeTokenPermuteWithEp 就是完成这一路由排列 EP 范围切片的算子。它的计算逻辑受属性paddedMode控制分两种模式paddedMode 为 false非填充模式当前唯一支持的模式此时 indices 的每个元素表示对应 token 被路由到的专家编号。令sortedIndicesFirst argSort(indices)对展平后的 indices 做稳定排序得到的位置序列再对其做一次 argSort 得到逆置换$$ sortedIndicesFirstargSort(indices) $$$$ sortedIndicesOutargSort(sortedIndicesFirst) $$当rangeOptional[0] sortedIndicesOut[i] rangeOptional[1]时README 公式中的 range 指 rangeOptional注意与 aclnn 文档中rangeOptional[1]的表述等价即落在 EP 有效区间内的元素$$ permuteTokensOut[sortedIndicesOut[i]-range[0]]tokens[i//topK] $$$$ permuteProbsOut[sortedIndicesOut[i]-rangeOptional[0]]probsOptional[i] $$其中 topK 表示每个 token 选择的专家数量若 indices 为 2 维topK 等于 indices 最后一维的大小若 indices 为 1 维则 topK 为 1见 aclnnMoeTokenPermuteWithEp.md。该公式在仓库的 ST 测试参考实现中也有完全对应的逻辑见 executor_aclnnMoeTokenPermuteWithEp.py先对展平 indices 做torch.argsort(..., stableTrue)再对其结果做一次 argsort 得到sorted_indices1随后用range[0]:range[1]切片并index_select出 tokens 与 probs。paddedMode 为 true填充模式暂不支持$$ permuteTokensOut[i]tokens[indices[i]] $$$$ sortedIndicesOutindices $$即 indices 已被填充为代表每个专家选中的 token 索引此时不做排序。目前算子明确不支持该模式详见下文约束说明。参数说明算子接口的参数完整说明如下表来自 README.md并结合 aclnnMoeTokenPermuteWithEp.md 补充了 shape 与使用说明参数名输入/输出/属性描述数据类型数据格式tokens输入permute 中的输入 tokens公式中的tokens要求 2D、shape 为 (num_tokens, hidden_size)支持空 tensorBFLOAT16、FLOAT16、FLOAT32NDindices输入输入 tokens 对应的专家索引公式中的indicespaddedMode 为 false 时表示每个输入 token 对应的 topK 个处理专家索引shape 为 (num_tokens, topK_num) 或 (num_tokens)支持空 tensorINT32、INT64NDprobsOptional输入可选输入输入 tokens 对应的专家概率公式中的probsOptionalshape 与 indices 相同传入空则不输出 permuteProbsOutBFLOAT16、FLOAT16、FLOAT32NDrangeOptional属性ep 切分的有效范围size 为 2为空时回退到 aclnnMoeTokenPermuteaclIntArray-numOutTokens属性有效输出 token 数在 rangeOptional 为空时生效0 表示不删除任何 token大于 0 时保留排序后前 numOutTokens 个 token小于 0 时按负切片索引处理INT64-paddedMode属性为 true 时表示 indices 已被填充为代表每个专家选中的 token 索引暂不支持false 表示对 indices 排序BOOL-permuteTokensOut输出indices 进行扩展并排序过的 tokens公式中的permuteTokensOutshape 为 (rangeOptional[1] - rangeOptional[0], hidden_size)数据类型同 tokensBFLOAT16、FLOAT16、FLOAT32NDsortedIndicesOut输出排序后的输出结果1D、shape 为 (num_tokens * topK_num)为 permuteTokensOut 与 tokens 的映射关系INT32NDpermuteProbsOut输出permute 之后的输出1D、shape 为 (rangeOptional[1] - rangeOptional[0])数据类型同 tokensBFLOAT16、FLOAT16、FLOAT32ND以上类型声明与 Host 侧算子注册完全一致在 moe_token_permute_with_ep_def.cpp 中tokens/probs/permute_tokens/permute_probs 均支持 BF16、FP16、FP32 且格式为 NDindices 支持 INT32/INT64sorted_indices 固定为 INT32三个属性range、num_out_tokens、padded_mode分别以ListInt({0, 0})、Int(0)、Bool(false)为默认值注册。补充说明数据类型间的组合关系是固定配对的tokens、probs、permuteTokensOut、permuteProbsOut 四者数据类型一致indices 与 sortedIndicesOut 按BF16/FP16/FP32 × INT64/INT32组合出 6 种二进制版本具体可见 moe_token_permute_with_ep_binary.jsonascend910_93与ascend910b目录下各有一份相同的组合清单。probsOptional、rangeOptional、permuteTokensOut、sortedIndicesOut、permuteProbsOut 均支持空不传但 probsOptional 为空时同时要求不输出 permuteProbsOutpermuteTokensOut、sortedIndicesOut、permuteProbsOut 输出均要求为连续 Tensor不支持非连续输入输入侧 tokens/indices/probs 则支持非连续 Tensor可参考仓库通用文档 non_contiguous_tensor.md。约束说明使用 MoeTokenPermuteWithEp 时需遵守以下约束来自 README.md 约束说明并在 Tiling 源码中得到印证indices 元素规模与取值indices 元素个数小于16777215值大于等于 0 且小于16777215。Tiling 侧的检查在 moe_token_permute_with_ep_tiling.cpp 中实现SORT_LIMIT_LENGTH 16777215不满足时返回GRAPH_FAILED。README 同时提醒indices 仅支持 int32 或 int64 的最大或最小值范围内的取值其余值不在范围内时排序结果不正确。topK 上限topK 小于等于512。Tiling 中以MAX_INDICES_NUM 512校验moe_token_permute_with_ep_tiling.cpp要求0 topK 512。paddedMode 限制不支持 paddedMode 为true。Tiling 中CheckAndGetAttrsInfo对paddedMode true直接报错Currently only support padded_mode is falsemoe_token_permute_with_ep_tiling.cpp。rangeOptional 为空时的回退行为当 rangeOptional 为空时忽略 probsOptional 和 permuteTokensOut执行逻辑回退到 aclnnMoeTokenPermute。此时由num_out_tokens属性决定有效输出 token 数。range 取值范围range 的 start/end 必须落在[0, totalLength]内totalLength num_tokens × topK见 moe_token_permute_with_ep_tiling.cpp。调用说明两段式 aclnn 接口MoeTokenPermuteWithEp 采用 CANN 算子的标准两段式接口调用方式先调用aclnnMoeTokenPermuteWithEpGetWorkspaceSize获取计算所需 workspace 大小与包含算子计算流程的执行器再调用aclnnMoeTokenPermuteWithEp执行计算。// 第一段获取 workspace 大小与执行器 aclnnStatus aclnnMoeTokenPermuteWithEpGetWorkspaceSize( const aclTensor *tokens, const aclTensor *indices, const aclTensor *probsOptional, const aclIntArray *rangeOptional, int64_t numOutTokens, bool paddedMode, const aclTensor *permuteTokensOut, const aclTensor *sortedIndicesOut, const aclTensor *permuteProbsOut, uint64_t *workspaceSize, aclOpExecutor **executor);// 第二段执行计算 aclnnStatus aclnnMoeTokenPermuteWithEp( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);第一段接口参数与返回值第一段接口的 9 个算子参数tokens、indices、probsOptional、rangeOptional、numOutTokens、paddedMode、permuteTokensOut、sortedIndicesOut、permuteProbsOut语义与上一节表格完全一致另有 2 个输出参数workspaceSize输出返回需要在 Device 侧申请的 workspace 大小字节数。executor输出返回 op 执行器包含算子计算流程。返回值类型为aclnnStatus状态码说明可参考仓库通用文档 aclnn_return_code.md。第一段接口完成入参校验以下场景会报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入和输出的 Tensor 是空指针ACLNN_ERR_PARAM_INVALID161002输入和输出的数据类型或数据格式不在支持的范围内ACLNN_ERR_INNER_TILING_ERROR561002tokens 的 shape 维度不为 2indices 的 shape 不为 1D 或 2D或 paddedMode 为 false 时 indices 的 shape 第一维与 tokens 的第一维不相等paddedMode 为 true暂不支持这些校验逻辑对应 Host 侧 Tiling 的输入检查CheckInputShapemoe_token_permute_with_ep_tiling.cpp要求 indices 为 1D 或 2D、tokens 第 0 维与 indices 第 0 维相等、probs若传入的维度与 indices 一致且第 0 维相等、2D 情况下 indices 与 probs 第 1 维相等。第二段接口参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream另外aclnnMoeTokenPermuteWithEp 默认是确定性计算实现详见 aclnnMoeTokenPermuteWithEp.md 与通用文档 determinism_compute.md。完整调用示例与运行验证仓库在 test_aclnn_moe_token_permute_with_ep.cpp 中提供了完整可参考的示例核心流程如下省略头文件与工具函数后int main() { // 1. device/stream 初始化aclInit / aclrtSetDevice / aclrtCreateStream int32_t deviceId 0; aclrtStream stream; Init(deviceId, stream); // 2. 构造输入与输出 shape std::vectorint64_t xShape {3, 4}; // tokens: (num_tokens3, hidden4) std::vectorint64_t idxShape {3, 2}; // indices: (3, topK2) std::vectorint64_t probsShape {3, 2}; // probs: (3, 2) std::vectorint64_t expandedXOutShape {4, 4}; // permuteTokensOut: (range[1]-range[0], hidden) std::vectorint64_t idxOutShape {6}; // sortedIndicesOut: (num_tokens*topK) std::vectorint64_t expandedProbsOutShape {4}; // permuteProbsOut: (range[1]-range[0]) std::vectorfloat xHostData {0.1, 0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.2, 0.3, 0.3, 0.3, 0.3}; std::vectorint indicesHostData {1, 2, 3, 1, 2, 3}; // token0-专家1,2; token1-专家3,1; token2-专家2,3 std::vectorfloat probsHostData {0.5, 0.3, 0.4, 0.2, 0.5, 0.4}; std::vectorint64_t rangeHostData {1, 5}; // EP 有效范围 [1,5) int64_t numTokenOut 6; bool padMode false; // 用 aclrtMalloc aclrtMemcpy 将数据搬入 Device并 aclCreateTensor 创建 aclTensor // tokens 用 ACL_BF16indices 用 ACL_INT32probs 用 ACL_BF16输出同理 // 属性 range 用 aclCreateIntArray 创建 aclIntArray // 3. 第一段接口计算 workspace 与 executor uint64_t workspaceSize 0; aclOpExecutor* executor; aclnnMoeTokenPermuteWithEpGetWorkspaceSize(x, indices, probs, range, numTokenOut, padMode, expandedXOut, sortedIndicesOut, expandedProbsOut, workspaceSize, executor); // 按 workspaceSize 申请 device 内存 void* workspaceAddr nullptr; if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 4. 第二段接口执行计算 aclnnMoeTokenPermuteWithEp(workspaceAddr, workspaceSize, executor, stream); // 5. aclrtSynchronizeStream 同步随后 aclrtMemcpy 将三个输出拷回 Host 并打印 // 6. 释放 aclTensoraclDestroyTensor、device 内存aclrtFree、stream 与 device 资源 return 0; }示例要点头文件包含acl/acl.h与aclnnop/aclnn_moe_token_permute_with_ep.htokens/indices/probs 均可作为非连续 Tensor 传入但三个输出 Tensor 必须是连续内存示例中使用aclCreateTensor传入连续 strides 构造示例中numTokenOut 6与idxOutShape {6}num_tokens×topK6一致而 range[1,5) 决定 permuteTokensOut 与 permuteProbsOut 的实际切片行数为 4shape 分别为 {4,4} 与 {4}编译与运行样例的整体流程环境准备、编译命令、执行方式请参考仓库通用文档 compile_and_run_sample.md。源码级实现解析Host 侧算子注册与 shape 推导算子注册MoeTokenPermuteWithEp通过OpDef注册输入tokens必选、indices必选、probs可选与输出permute_tokens、sorted_indices、permute_probs并声明ascend910b、ascend910_93两个 AICore 配置见 moe_token_permute_with_ep_def.cpp。形状推导InferShape4MoeTokenPermuteWithEp根据 indices 维度确定 topK1D 时 topK12D 时取第 1 维推导出sortedIndicesLen topK × N随后依据 range 或 numOutTokens 计算permuteTokensOut第 0 维outTokens (start end) ? sortedIndicesLen (end - start) : (end - start)并 clamp 到[0, sortedIndicesLen]见 moe_token_permute_with_ep_infershape.cpp。数据类型推导则固定输出 sorted_indices 为 INT32permute_tokens/permute_probs 继承输入类型。Host 侧Tiling 策略Tiling 逻辑在 moe_token_permute_with_ep_tiling.cpp 中实现可以从中看到算子的核心设计排序模式选择根据总元素数totalLength n × topK与单次排序容量sortLoopMaxElement由 UB 大小推算比较totalLength sortLoopMaxElement时选择单核排序tilingKey 加 1SORT_ONE_CORE_MODE否则选择多核排序tilingKey 加 2SORT_MULTI_CORE_MODE见Tiling4VBSComputeL696-L711。多核切分多核模式下核数按4^x次幂向上取整CeilLog4且不超过物理 AIV 核数每个核的元素数按 32 对齐SORT32_ALIGN_ELEMENT并计算各核 loop 数与最后一核的 workspace 索引见Tinlig4VBSMultiCoreComputeL658-L694。中间归并当参与归并的核队列数大于 4 时需要增加中间归并核Tiling4VMSMiddleComputeL713-L723体现了多核排序 分层归并的 VBS/VMS 设计。IndexCopy 切分Tiling4IndexCopyCompute负责把按排序结果搬移 token阶段按 token 数、topK、hiddencols切分到多核与多 loop并依据 UB 容量计算onceUbTokenNums、onceIndicesTokenNums等参数当单个 token 的 hidden 过大放不下双 buffer 时会进入按 D 维拆分SPILT_D_MODE的模式见CalculateUBParamsL741-L776。numOutTokens 使能位当numOutTokens ! n × topK时tilingKey 追加ENABLE_NUMOUTTOKENS值为 4见 L838-L840。workspace 计算workspace 排序临时空间totalLength × 4 × sizeof(float) 多核同步空间coreNum × 32 × 2 16MB 固定缓冲见GetWorkspaceSizeL621-L629同时通过SetScheduleMode(1)声明算子需要独占全核并涉及核间同步L631-L643。Kernel 侧执行流水Kernel 入口为 moe_token_permute_with_ep.cpp 中的moe_token_permute_with_ep函数AIC 核直接返回仅 AIV 核执行。核心流程通过GENERAL_OP_IMPL宏串联三个阶段对应 L32-L47第一轮排序对 indices 执行 argSort得到sortedIndicesFirst对应MoeSortOneCore或MoeSortMultiCore第二轮排序对第一轮结果再次执行 argSort得到sortedIndicesOut逆置换两轮排序间复用 workspaceIndexCopy依据排序结果把 tokens按i // topK广播与 probs 搬移到 permute 输出同时按 range 切片。GENERAL_OP_IMPL的模板参数组合对应 8 种 tilingKey单核/多核排序 × 是否按 D 拆分 × 是否启用 numOutTokens 切片分别实例化MoeindexCopyOp/MoeindexCopySpiltDOp第三个模板参数 true 表示启用 numOutTokens 切片。算子实现按排序与搬移两个 Kernel 函数族拆分在 moe_sort_one_core_with_ep.h、moe_sort_multi_core_with_ep.h、moe_mrgsort_with_ep.h、moe_mrgsort_out_with_ep.h、moe_index_copy_with_ep.h 与 moe_index_copy_spilt_d_with_ep.h 中公共常量如单次排序 32 元素、32 字节块对齐等定义在 moe_token_permute_with_ep_common.h。测试与验证仓库为算子提供了完整的测试体系ST系统测试测试定义见 atk_aclnnMoeTokenPermuteWithEp.json其参考实现 executor_aclnnMoeTokenPermuteWithEp.py 用torch.argsort(..., stableTrue)复现了算子的两轮排序 range 切片 index_select 逻辑可作为理解算子语义与自行验证输出正确性的参考基准。UT单元测试包括 Host 侧 shape/tiling 单测test_moe_token_permute_with_ep_infershape.cpp、test_moe_token_permute_with_ep_tiling.cpp与 Kernel 单测test_moe_token_permute_with_ep.cpp单测构建入口见 tests/ut/CMakeLists.txt。总结MoeTokenPermuteWithEp 是 CANN ops-transformer 面向 MoE 专家并行场景的关键算子它以两轮 argSort 求逆置换 按 range 切片的方式完成 token/prob 的按专家聚拢支持单核/多核排序、按 D 维拆分与 numOutTokens 切片等多种执行模式并支持在 rangeOptional 为空时回退到 aclnnMoeTokenPermute。在 Atlas A2/A3 系列产品上开发者可通过 aclnnMoeTokenPermuteWithEp 两段式接口直接调用并结合 示例代码 与 ST 参考实现快速完成功能验证与集成。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-transformer MoeTokenPermuteWithEp 算子深度解析：MoE EP 场景下的 Token Permute 实现与 aclnn 接口调用指南
CANN ops-transformer MoeTokenPermuteWithEp 算子深度解析MoE EP 场景下的 Token Permute 实现与 aclnn 接口调用指南【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformerMoeTokenPermuteWithEp 是 CANN ops-transformer 中面向 MoEMixture of Experts专家并行EP场景的 permute 算子它根据专家索引 indices 将 tokens 与可选的 probs 广播、排序并按 rangeOptional 指定的 EP 有效范围切片输出。本文基于仓库中的 README.md 与 aclnnMoeTokenPermuteWithEp.md结合 Host 侧算子定义、Tiling 逻辑与 Kernel 实现源码完整讲解算子功能、参数语义、约束、两段式 aclnn 调用方式及源码级实现原理帮助你直接上手在 Atlas A2/A3 系列产品上完成该算子的开发与验证。产品支持情况MoeTokenPermuteWithEp 算子对产品的支持情况如下与 aclnnMoeTokenPermuteWithEp.md 及算子定义中AICore().AddConfig声明一致见 moe_token_permute_with_ep_def.cpp产品是否支持Ascend 950PR/Ascend 950DT×Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×即该算子仅面向 Atlas A2 与 Atlas A3 训练/推理系列产品提供算力配置ascend910b与ascend910_93。功能说明EP 场景下的 Token Permute 计算MoE 模型中稀疏门控网络会将 token 路由到不同的专家expert。当采用专家并行EP时每个设备只持有部分专家需要把输入 token 按照其专家索引重新排列permute使属于本设备所辖专家的 token 聚拢到一起才能喂给本地专家计算计算完成后再通过 unpermute 还原顺序。MoeTokenPermuteWithEp 就是完成这一路由排列 EP 范围切片的算子。它的计算逻辑受属性paddedMode控制分两种模式paddedMode 为 false非填充模式当前唯一支持的模式此时 indices 的每个元素表示对应 token 被路由到的专家编号。令sortedIndicesFirst argSort(indices)对展平后的 indices 做稳定排序得到的位置序列再对其做一次 argSort 得到逆置换$$ sortedIndicesFirstargSort(indices) $$$$ sortedIndicesOutargSort(sortedIndicesFirst) $$当rangeOptional[0] sortedIndicesOut[i] rangeOptional[1]时README 公式中的 range 指 rangeOptional注意与 aclnn 文档中rangeOptional[1]的表述等价即落在 EP 有效区间内的元素$$ permuteTokensOut[sortedIndicesOut[i]-range[0]]tokens[i//topK] $$$$ permuteProbsOut[sortedIndicesOut[i]-rangeOptional[0]]probsOptional[i] $$其中 topK 表示每个 token 选择的专家数量若 indices 为 2 维topK 等于 indices 最后一维的大小若 indices 为 1 维则 topK 为 1见 aclnnMoeTokenPermuteWithEp.md。该公式在仓库的 ST 测试参考实现中也有完全对应的逻辑见 executor_aclnnMoeTokenPermuteWithEp.py先对展平 indices 做torch.argsort(..., stableTrue)再对其结果做一次 argsort 得到sorted_indices1随后用range[0]:range[1]切片并index_select出 tokens 与 probs。paddedMode 为 true填充模式暂不支持$$ permuteTokensOut[i]tokens[indices[i]] $$$$ sortedIndicesOutindices $$即 indices 已被填充为代表每个专家选中的 token 索引此时不做排序。目前算子明确不支持该模式详见下文约束说明。参数说明算子接口的参数完整说明如下表来自 README.md并结合 aclnnMoeTokenPermuteWithEp.md 补充了 shape 与使用说明参数名输入/输出/属性描述数据类型数据格式tokens输入permute 中的输入 tokens公式中的tokens要求 2D、shape 为 (num_tokens, hidden_size)支持空 tensorBFLOAT16、FLOAT16、FLOAT32NDindices输入输入 tokens 对应的专家索引公式中的indicespaddedMode 为 false 时表示每个输入 token 对应的 topK 个处理专家索引shape 为 (num_tokens, topK_num) 或 (num_tokens)支持空 tensorINT32、INT64NDprobsOptional输入可选输入输入 tokens 对应的专家概率公式中的probsOptionalshape 与 indices 相同传入空则不输出 permuteProbsOutBFLOAT16、FLOAT16、FLOAT32NDrangeOptional属性ep 切分的有效范围size 为 2为空时回退到 aclnnMoeTokenPermuteaclIntArray-numOutTokens属性有效输出 token 数在 rangeOptional 为空时生效0 表示不删除任何 token大于 0 时保留排序后前 numOutTokens 个 token小于 0 时按负切片索引处理INT64-paddedMode属性为 true 时表示 indices 已被填充为代表每个专家选中的 token 索引暂不支持false 表示对 indices 排序BOOL-permuteTokensOut输出indices 进行扩展并排序过的 tokens公式中的permuteTokensOutshape 为 (rangeOptional[1] - rangeOptional[0], hidden_size)数据类型同 tokensBFLOAT16、FLOAT16、FLOAT32NDsortedIndicesOut输出排序后的输出结果1D、shape 为 (num_tokens * topK_num)为 permuteTokensOut 与 tokens 的映射关系INT32NDpermuteProbsOut输出permute 之后的输出1D、shape 为 (rangeOptional[1] - rangeOptional[0])数据类型同 tokensBFLOAT16、FLOAT16、FLOAT32ND以上类型声明与 Host 侧算子注册完全一致在 moe_token_permute_with_ep_def.cpp 中tokens/probs/permute_tokens/permute_probs 均支持 BF16、FP16、FP32 且格式为 NDindices 支持 INT32/INT64sorted_indices 固定为 INT32三个属性range、num_out_tokens、padded_mode分别以ListInt({0, 0})、Int(0)、Bool(false)为默认值注册。补充说明数据类型间的组合关系是固定配对的tokens、probs、permuteTokensOut、permuteProbsOut 四者数据类型一致indices 与 sortedIndicesOut 按BF16/FP16/FP32 × INT64/INT32组合出 6 种二进制版本具体可见 moe_token_permute_with_ep_binary.jsonascend910_93与ascend910b目录下各有一份相同的组合清单。probsOptional、rangeOptional、permuteTokensOut、sortedIndicesOut、permuteProbsOut 均支持空不传但 probsOptional 为空时同时要求不输出 permuteProbsOutpermuteTokensOut、sortedIndicesOut、permuteProbsOut 输出均要求为连续 Tensor不支持非连续输入输入侧 tokens/indices/probs 则支持非连续 Tensor可参考仓库通用文档 non_contiguous_tensor.md。约束说明使用 MoeTokenPermuteWithEp 时需遵守以下约束来自 README.md 约束说明并在 Tiling 源码中得到印证indices 元素规模与取值indices 元素个数小于16777215值大于等于 0 且小于16777215。Tiling 侧的检查在 moe_token_permute_with_ep_tiling.cpp 中实现SORT_LIMIT_LENGTH 16777215不满足时返回GRAPH_FAILED。README 同时提醒indices 仅支持 int32 或 int64 的最大或最小值范围内的取值其余值不在范围内时排序结果不正确。topK 上限topK 小于等于512。Tiling 中以MAX_INDICES_NUM 512校验moe_token_permute_with_ep_tiling.cpp要求0 topK 512。paddedMode 限制不支持 paddedMode 为true。Tiling 中CheckAndGetAttrsInfo对paddedMode true直接报错Currently only support padded_mode is falsemoe_token_permute_with_ep_tiling.cpp。rangeOptional 为空时的回退行为当 rangeOptional 为空时忽略 probsOptional 和 permuteTokensOut执行逻辑回退到 aclnnMoeTokenPermute。此时由num_out_tokens属性决定有效输出 token 数。range 取值范围range 的 start/end 必须落在[0, totalLength]内totalLength num_tokens × topK见 moe_token_permute_with_ep_tiling.cpp。调用说明两段式 aclnn 接口MoeTokenPermuteWithEp 采用 CANN 算子的标准两段式接口调用方式先调用aclnnMoeTokenPermuteWithEpGetWorkspaceSize获取计算所需 workspace 大小与包含算子计算流程的执行器再调用aclnnMoeTokenPermuteWithEp执行计算。// 第一段获取 workspace 大小与执行器 aclnnStatus aclnnMoeTokenPermuteWithEpGetWorkspaceSize( const aclTensor *tokens, const aclTensor *indices, const aclTensor *probsOptional, const aclIntArray *rangeOptional, int64_t numOutTokens, bool paddedMode, const aclTensor *permuteTokensOut, const aclTensor *sortedIndicesOut, const aclTensor *permuteProbsOut, uint64_t *workspaceSize, aclOpExecutor **executor);// 第二段执行计算 aclnnStatus aclnnMoeTokenPermuteWithEp( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);第一段接口参数与返回值第一段接口的 9 个算子参数tokens、indices、probsOptional、rangeOptional、numOutTokens、paddedMode、permuteTokensOut、sortedIndicesOut、permuteProbsOut语义与上一节表格完全一致另有 2 个输出参数workspaceSize输出返回需要在 Device 侧申请的 workspace 大小字节数。executor输出返回 op 执行器包含算子计算流程。返回值类型为aclnnStatus状态码说明可参考仓库通用文档 aclnn_return_code.md。第一段接口完成入参校验以下场景会报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入和输出的 Tensor 是空指针ACLNN_ERR_PARAM_INVALID161002输入和输出的数据类型或数据格式不在支持的范围内ACLNN_ERR_INNER_TILING_ERROR561002tokens 的 shape 维度不为 2indices 的 shape 不为 1D 或 2D或 paddedMode 为 false 时 indices 的 shape 第一维与 tokens 的第一维不相等paddedMode 为 true暂不支持这些校验逻辑对应 Host 侧 Tiling 的输入检查CheckInputShapemoe_token_permute_with_ep_tiling.cpp要求 indices 为 1D 或 2D、tokens 第 0 维与 indices 第 0 维相等、probs若传入的维度与 indices 一致且第 0 维相等、2D 情况下 indices 与 probs 第 1 维相等。第二段接口参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream另外aclnnMoeTokenPermuteWithEp 默认是确定性计算实现详见 aclnnMoeTokenPermuteWithEp.md 与通用文档 determinism_compute.md。完整调用示例与运行验证仓库在 test_aclnn_moe_token_permute_with_ep.cpp 中提供了完整可参考的示例核心流程如下省略头文件与工具函数后int main() { // 1. device/stream 初始化aclInit / aclrtSetDevice / aclrtCreateStream int32_t deviceId 0; aclrtStream stream; Init(deviceId, stream); // 2. 构造输入与输出 shape std::vectorint64_t xShape {3, 4}; // tokens: (num_tokens3, hidden4) std::vectorint64_t idxShape {3, 2}; // indices: (3, topK2) std::vectorint64_t probsShape {3, 2}; // probs: (3, 2) std::vectorint64_t expandedXOutShape {4, 4}; // permuteTokensOut: (range[1]-range[0], hidden) std::vectorint64_t idxOutShape {6}; // sortedIndicesOut: (num_tokens*topK) std::vectorint64_t expandedProbsOutShape {4}; // permuteProbsOut: (range[1]-range[0]) std::vectorfloat xHostData {0.1, 0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.2, 0.3, 0.3, 0.3, 0.3}; std::vectorint indicesHostData {1, 2, 3, 1, 2, 3}; // token0-专家1,2; token1-专家3,1; token2-专家2,3 std::vectorfloat probsHostData {0.5, 0.3, 0.4, 0.2, 0.5, 0.4}; std::vectorint64_t rangeHostData {1, 5}; // EP 有效范围 [1,5) int64_t numTokenOut 6; bool padMode false; // 用 aclrtMalloc aclrtMemcpy 将数据搬入 Device并 aclCreateTensor 创建 aclTensor // tokens 用 ACL_BF16indices 用 ACL_INT32probs 用 ACL_BF16输出同理 // 属性 range 用 aclCreateIntArray 创建 aclIntArray // 3. 第一段接口计算 workspace 与 executor uint64_t workspaceSize 0; aclOpExecutor* executor; aclnnMoeTokenPermuteWithEpGetWorkspaceSize(x, indices, probs, range, numTokenOut, padMode, expandedXOut, sortedIndicesOut, expandedProbsOut, workspaceSize, executor); // 按 workspaceSize 申请 device 内存 void* workspaceAddr nullptr; if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 4. 第二段接口执行计算 aclnnMoeTokenPermuteWithEp(workspaceAddr, workspaceSize, executor, stream); // 5. aclrtSynchronizeStream 同步随后 aclrtMemcpy 将三个输出拷回 Host 并打印 // 6. 释放 aclTensoraclDestroyTensor、device 内存aclrtFree、stream 与 device 资源 return 0; }示例要点头文件包含acl/acl.h与aclnnop/aclnn_moe_token_permute_with_ep.htokens/indices/probs 均可作为非连续 Tensor 传入但三个输出 Tensor 必须是连续内存示例中使用aclCreateTensor传入连续 strides 构造示例中numTokenOut 6与idxOutShape {6}num_tokens×topK6一致而 range[1,5) 决定 permuteTokensOut 与 permuteProbsOut 的实际切片行数为 4shape 分别为 {4,4} 与 {4}编译与运行样例的整体流程环境准备、编译命令、执行方式请参考仓库通用文档 compile_and_run_sample.md。源码级实现解析Host 侧算子注册与 shape 推导算子注册MoeTokenPermuteWithEp通过OpDef注册输入tokens必选、indices必选、probs可选与输出permute_tokens、sorted_indices、permute_probs并声明ascend910b、ascend910_93两个 AICore 配置见 moe_token_permute_with_ep_def.cpp。形状推导InferShape4MoeTokenPermuteWithEp根据 indices 维度确定 topK1D 时 topK12D 时取第 1 维推导出sortedIndicesLen topK × N随后依据 range 或 numOutTokens 计算permuteTokensOut第 0 维outTokens (start end) ? sortedIndicesLen (end - start) : (end - start)并 clamp 到[0, sortedIndicesLen]见 moe_token_permute_with_ep_infershape.cpp。数据类型推导则固定输出 sorted_indices 为 INT32permute_tokens/permute_probs 继承输入类型。Host 侧Tiling 策略Tiling 逻辑在 moe_token_permute_with_ep_tiling.cpp 中实现可以从中看到算子的核心设计排序模式选择根据总元素数totalLength n × topK与单次排序容量sortLoopMaxElement由 UB 大小推算比较totalLength sortLoopMaxElement时选择单核排序tilingKey 加 1SORT_ONE_CORE_MODE否则选择多核排序tilingKey 加 2SORT_MULTI_CORE_MODE见Tiling4VBSComputeL696-L711。多核切分多核模式下核数按4^x次幂向上取整CeilLog4且不超过物理 AIV 核数每个核的元素数按 32 对齐SORT32_ALIGN_ELEMENT并计算各核 loop 数与最后一核的 workspace 索引见Tinlig4VBSMultiCoreComputeL658-L694。中间归并当参与归并的核队列数大于 4 时需要增加中间归并核Tiling4VMSMiddleComputeL713-L723体现了多核排序 分层归并的 VBS/VMS 设计。IndexCopy 切分Tiling4IndexCopyCompute负责把按排序结果搬移 token阶段按 token 数、topK、hiddencols切分到多核与多 loop并依据 UB 容量计算onceUbTokenNums、onceIndicesTokenNums等参数当单个 token 的 hidden 过大放不下双 buffer 时会进入按 D 维拆分SPILT_D_MODE的模式见CalculateUBParamsL741-L776。numOutTokens 使能位当numOutTokens ! n × topK时tilingKey 追加ENABLE_NUMOUTTOKENS值为 4见 L838-L840。workspace 计算workspace 排序临时空间totalLength × 4 × sizeof(float) 多核同步空间coreNum × 32 × 2 16MB 固定缓冲见GetWorkspaceSizeL621-L629同时通过SetScheduleMode(1)声明算子需要独占全核并涉及核间同步L631-L643。Kernel 侧执行流水Kernel 入口为 moe_token_permute_with_ep.cpp 中的moe_token_permute_with_ep函数AIC 核直接返回仅 AIV 核执行。核心流程通过GENERAL_OP_IMPL宏串联三个阶段对应 L32-L47第一轮排序对 indices 执行 argSort得到sortedIndicesFirst对应MoeSortOneCore或MoeSortMultiCore第二轮排序对第一轮结果再次执行 argSort得到sortedIndicesOut逆置换两轮排序间复用 workspaceIndexCopy依据排序结果把 tokens按i // topK广播与 probs 搬移到 permute 输出同时按 range 切片。GENERAL_OP_IMPL的模板参数组合对应 8 种 tilingKey单核/多核排序 × 是否按 D 拆分 × 是否启用 numOutTokens 切片分别实例化MoeindexCopyOp/MoeindexCopySpiltDOp第三个模板参数 true 表示启用 numOutTokens 切片。算子实现按排序与搬移两个 Kernel 函数族拆分在 moe_sort_one_core_with_ep.h、moe_sort_multi_core_with_ep.h、moe_mrgsort_with_ep.h、moe_mrgsort_out_with_ep.h、moe_index_copy_with_ep.h 与 moe_index_copy_spilt_d_with_ep.h 中公共常量如单次排序 32 元素、32 字节块对齐等定义在 moe_token_permute_with_ep_common.h。测试与验证仓库为算子提供了完整的测试体系ST系统测试测试定义见 atk_aclnnMoeTokenPermuteWithEp.json其参考实现 executor_aclnnMoeTokenPermuteWithEp.py 用torch.argsort(..., stableTrue)复现了算子的两轮排序 range 切片 index_select 逻辑可作为理解算子语义与自行验证输出正确性的参考基准。UT单元测试包括 Host 侧 shape/tiling 单测test_moe_token_permute_with_ep_infershape.cpp、test_moe_token_permute_with_ep_tiling.cpp与 Kernel 单测test_moe_token_permute_with_ep.cpp单测构建入口见 tests/ut/CMakeLists.txt。总结MoeTokenPermuteWithEp 是 CANN ops-transformer 面向 MoE 专家并行场景的关键算子它以两轮 argSort 求逆置换 按 range 切片的方式完成 token/prob 的按专家聚拢支持单核/多核排序、按 D 维拆分与 numOutTokens 切片等多种执行模式并支持在 rangeOptional 为空时回退到 aclnnMoeTokenPermute。在 Atlas A2/A3 系列产品上开发者可通过 aclnnMoeTokenPermuteWithEp 两段式接口直接调用并结合 示例代码 与 ST 参考实现快速完成功能验证与集成。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考