CANN ops-transformer 的 moe_token_permute 算子：基于 torch_npu 的 MoE Token 重排与 MX 量化实现

CANN ops-transformer 的 moe_token_permute 算子：基于 torch_npu 的 MoE Token 重排与 MX 量化实现

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer

导读

moe_token_permute是 CANN ops-transformer 算子库中面向 MoE（Mixture of Experts）训练/推理场景的核心前置算子，它以torch_npu扩展接口的形式提供：根据专家索引indices将输入tokens按专家排序，并返回排序后的 token 与原始 token 的行映射关系。本文以 moe_token_permute 官方 Torch API 文档 为主体，结合 Python 封装层、C++ 接入层、算子定义与 Kernel 实现，完整讲解函数原型、参数语义、返回值 shape 推导、MXFP8/MXFP4 量化模式以及源码级实现原理，帮助读者在 NPU 上正确、高效地使用该接口完成 MoE 路由后的 token 重排。

功能说明：MoE 中 Token 重排的语义

在 MoE 模型中，Router（门控网络）会为每个 token 产生一个或多个专家索引，随后需要把属于同一专家的 token 聚合到一起，才能送入对应的专家网络计算。moe_token_permute正是完成这一步"重排"的高效算子：

根据indices将输入tokens扩展（每个 token 可按K个索引重复），并按照专家索引排序，返回排序后的 token（permuted_tokens）及其与原始 token 的映射关系（sorted_indices）。

算子仓库 README（moe_token_permute/README.md）中给出了精确的计算公式（paddedMode=false时）：

sortedIndicesFirst = argSort(indices) sortedIndicesOut = argSort(sortedIndicesFirst) permuteTokens[sortedIndicesFirst[i]] = tokens[i // topK]

即：先对展平后的indices做一次稳定的argSort得到按专家分组的行顺序，再取该顺序的逆映射作为sorted_indices输出，最终按sortedIndicesFirst[i] // topK从原始 token 中取值填入输出。仓库测试资产的 golden 参考实现（golden.py）用 PyTorch 原语还原了这一过程：torch.argsort(expert_idx.reshape(-1), stable=True)得到排序后的行号，expanded_row_idx[sorted_row_idx] = torch.arange(...)构造逆映射，再通过input_x[selected_indices]完成 token 收集，可作为理解算子行为的最小可读参考。

值得一提的是，indices可以是一维[N]（每个 token 对应 1 个专家，即 topK=1），也可以是二维[N, K]（每个 token 对应 K 个专家）。当K > 1时，token 会按索引重复，最终输出行数等于indices.numel()。

产品支持情况

根据官方文档，该接口在不同产品上的支持情况如下：

从算子注册配置（moe_token_permute_def.cpp）可以看到，AICore 侧仅注册了ascend910b与ascend910_93两个 AI Core 配置，与文档所述 Atlas A2/A3 及 Ascend 950 系列支持范围一致。

函数原型与 Python 侧注册

接口的完整函数原型如下：

cann_ops_transformer.ops.moe_token_permute( tokens, indices, num_out_tokens=None, padded_mode=False, quant_mode=-1, ) -> (Tensor, Tensor, Tensor)

在 Python 侧（moe_token_permute.py），该接口通过torch.ops.cann_ops_transformer.moe_token_permute分发，其自定义算子 schema 为：

moe_token_permute(Tensor tokens, Tensor indices, int? num_out_tokens=None, bool padded_mode=False, int quant_mode=-1) -> (Tensor, Tensor, Tensor)

底层 C++ 封装（csrc/moe_token_permute.cpp）会完成输入合法性检查、num_out_tokens归一化、量化模式解析（非 Ascend 950 一律强制quant_mode=-1），最终调用aclnnMoeTokenPermuteV2执行。该接口属于cann_ops_transformer扩展包，使用前需要安装该包并导入torch_npu。

参数说明

参数细节与产品差异

Atlas A3 训练系列产品 / Atlas A3 推理系列产品、Atlas A2 训练系列产品 / Atlas A2 推理系列产品：quant_mode仅支持-1，即按非量化路径执行。

Ascend 950PR / Ascend 950DT：

• quant_mode支持-1 / 2 / 3 / 9；
• 取值为9（MXFP4）时，隐藏维H必须为偶数；
• 取值为2 / 3 / 9时，不支持 autograd（即输入tokens不能带有requires_grad=True）。

这一差异在 C++ 侧有明确对应实现（csrc/moe_token_permute.cpp）：ResolveQuantMode首先通过aclrtGetSocName判断是否 Ascend 950，非 950 平台强制回退为-1，并检查tokens.requires_grad()与量化模式互斥。

num_out_tokens的语义与 Python 封装中的_actual_num_out_tokens（moe_token_permute.py）以及 C++ 侧GetActualNumOutTokens（csrc/moe_token_permute.cpp）完全一致：

num_out_tokens 为 None 或 0：有效输出行数 = F num_out_tokens > 0： 有效输出行数 = min(num_out_tokens, F) num_out_tokens < 0： 有效输出行数 = max(F + num_out_tokens, 0)

返回值说明

接口固定返回三个 Tensor：

permuted_tokens, sorted_indices, expanded_scale

记F = indices.numel()、H = tokens.shape[1]，则各量化模式的输出规格如下：

其中：

CeilDiv(a, b) = (a + b - 1) // b AlignUp(a, b) = CeilDiv(a, b) * b

各输出含义：

• permuted_tokens：根据indices扩展并按专家索引排序后的 token；仅该输出与expanded_scale的第一维受num_out_tokens影响。
• sorted_indices：permuted_tokens与原始tokens的行映射关系，长度始终为indices.numel()（即不受num_out_tokens截断影响），dtype 固定为 int32。
• expanded_scale：量化输出对应的 per-token 分块 scale；非量化模式返回空 Tensor。

[!NOTE] MXFP4 类型当前通过 uint8 Tensor 承载物理存储。调用 aclnn 接口时，permuted_tokens会被解释为ACL_FLOAT4_E2M1，expanded_scale会被解释为ACL_FLOAT8_E8M0。因此表中的 uint8 表示 PyTorch 侧的物理存储类型，不表示量化数据的逻辑类型。

上述 shape 计算逻辑在 Python 侧的 Meta 函数moe_token_permute_meta（moe_token_permute.py）中有逐字对应的实现，例如 MXFP8 的 scale 列数计算为_align_up(_ceil_div(H, 32), 2)，MXFP4 的 scale 物理 shape 为(M, _ceil_div(H, 64), 2)，与文档表格完全一致，可作为 shape 推导的权威参考。

约束说明

• indices元素个数必须小于16777215，元素值必须大于等于0且小于16777215。这一上限在 tiling 侧同样存在（moe_token_permute_tiling.cpp 中的SORT_LIMIT_LENGTH = 16777215），超出范围时排序结果不正确。
• MXFP8 的每 32 个量化值共享一个 E8M0 scale，scale 数量向 2 对齐（即AlignUp(CeilDiv(H, 32), 2)）。
• MXFP4 的两个 4 bit 值打包在一个 uint8 中，因此permuted_tokens的 PyTorch 物理隐藏维为H / 2。
• padded_mode当前仅支持False。
• 在 Atlas A2 / A3 系列上，topK 需小于等于 512（来自 README.md 约束说明，与 tiling 中MAX_INDICES_NUM = 512、SPLIT_K_THRESHOLD = 512相印证）。
• 输入 shape 约束：tokens必须为 2D，indices必须为 1D 或 2D（csrc/moe_token_permute.cpp 与 Meta 函数中均有检查）。

确定性计算

该接口默认支持确定性计算，即在相同输入下多次调用会得到一致的输出顺序（排序使用稳定排序语义，golden 实现中亦使用stable=True的argsort）。这在分布式训练中有利于保证梯度与调试结果的可复现性。

调用示例

非量化调用

import torch import torch_npu from cann_ops_transformer.ops import moe_token_permute tokens = torch.randn(4, 128, dtype=torch.float16, device="npu") indices = torch.tensor( [[1, 0], [0, 1], [1, 0], [0, 1]], dtype=torch.int32, device="npu", ) permuted_tokens, sorted_indices, expanded_scale = moe_token_permute( tokens, indices, num_out_tokens=6, ) print(permuted_tokens.shape) # torch.Size([6, 128]) print(sorted_indices.shape) # torch.Size([8]) print(expanded_scale.shape) # torch.Size([0])

该示例中indices为[4, 2]，即F = 8，topK=2，每个 token 会按两个专家索引重复参与排序；num_out_tokens=6对permuted_tokens第一维截断至 6，而sorted_indices仍保留全部 8 个元素。

Ascend 950 MXFP8 调用（Ascend 950PR / Ascend 950DT）

permuted_tokens, sorted_indices, expanded_scale = moe_token_permute( tokens, indices, quant_mode=2, ) print(permuted_tokens.shape) # torch.Size([8, 128]) print(expanded_scale.shape) # torch.Size([8, 4])

这里H = 128，MXFP8 E5M2 模式下的 scale 列数为AlignUp(CeilDiv(128, 32), 2) = AlignUp(4, 2) = 4，故expanded_scale形状为[8, 4]，且未指定num_out_tokens时M = F = 8。如需 MXFP4（quant_mode=9）调用，请确保H为偶数。

源码级实现解析

算子定义与 shape 推导

算子注册（moe_token_permute_def.cpp）明确了输入输出 dtype 组合：tokens支持 bf16 / fp16 / fp32 / int8，indices支持 int64 / int32，sorted_indices固定 int32，属性num_out_tokens（默认 0）与padded_mode（默认 False）均为可选。infershape 逻辑（moe_token_permute_infershape.cpp）展示了如何从indices的[N]/[N, K]形状推导输出长度：sortedIndicesLen = topK * N，再结合numOutTokens按与文档一致的规则裁剪第一维。

NPU Kernel：排序 + 索引拷贝的两段式流水

Kernel 入口（moe_token_permute.cpp）体现了算子内部的计算拆解：通过TILING_KEY在单核/多核排序与是否按隐藏维拆分（SpiltD）之间选择不同执行路径，核心流程为：

1. 第一趟排序：对indices使用MoeSortOneCore/MoeSortMultiCore（实现见 moe_sort_one_core_token_permute.h 与 moe_sort_multi_core_token_permute.h）按专家索引稳定排序；
2. 第二趟排序：对第一趟结果再做一次排序得到逆映射sortedIndices；
3. 索引拷贝：通过MoeindexCopyOp/MoeindexCopySpiltDOp（moe_index_copy.h）按映射关系把tokens收集到permuted_tokens输出。

tiling 侧（moe_token_permute_tiling.cpp）负责根据 token 数量、隐藏维、topK 等参数决策排序模式与切分策略，并在对应单测 test_moe_token_permute_tiling.cpp 中验证。

与 aclnn 接口的关系

该 Torch API 在底层转发至aclnnMoeTokenPermuteV2（aclnn_moe_token_permute_v2.h）。而根据 README.md 的说明，在 Ascend 950 上框架内部会进一步转调aclnnMoeInitRoutingV2/aclnnMoeInitRoutingV3接口（见 aclnn_moe_token_permute.cpp），参数映射关系如下，遇到参数错误提示时可据此排查：

• tokens输入 对应MoeInitRoutingV2的x输入；
• indices输入 对应expertIdx输入；
• numOutTokens输入 对应activeNum输入；
• paddedMode输入 对应dropPadMode输入；
• permuteTokensOut输出 对应expandedXOut输出；
• sortedIndicesOut输出 对应expandedRowIdxOut输出。

总结与使用建议

moe_token_permute是 CANN ops-transformer 中实现 MoE 路由后 token 重排的标准接口，具有以下使用要点：

1. 先确认平台：Atlas A2/A3 系列仅支持非量化路径（quant_mode=-1），MXFP8/MXFP4 量化仅适用于 Ascend 950PR / 950DT，且量化模式不支持 autograd；
2. 善用num_out_tokens：可通过正数截断或负数尾部裁剪来控制有效输出行数，且它只影响permuted_tokens与expanded_scale的第一维，不影响sorted_indices长度；
3. 理解量化布局：MXFP8 每 32 个值共享一个 E8M0 scale 且列数向 2 对齐，MXFP4 以 uint8 物理承载、隐藏维减半，配合expanded_scale即可在 permute 阶段直接产出可用于 MX 量化的输出，减少后续二次量化开销。

相关源码与测试可继续在仓库中查阅：Torch API 文档、aclnn 接口文档、aclnnMoeTokenPermuteV2 文档、C 语言调用示例、kernel 单测 与 ST 执行器。

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer