算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载本文围绕 CANN ops-transformer 仓库中 moe_re_routing_v2/README.md 及其配套的 PyTorch 接口文档系统讲解 MoeReRoutingV2 算子的功能定位、数学原理、输入输出参数、约束条件并结合仓库源码算子 IR 原型、InferShape、Tiling、Kernel 模板与 Torch 扩展与图模式/单算子调用示例帮助读者在 Ascend 950 系列硬件上正确使用该算子完成 MoE 网络中 AlltoAll 之后的 Token 重排并让 TopK 权重与 Token 一一对应地同步搬运。产品支持情况MoeReRoutingV2 算子仅在 Ascend 950 系列上可用其余产品线暂不支持产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从 算子定义文件 可以看出注册算子时通过this-AICore().AddConfig(ascend950, regbaseCfg)仅为 ascend950 平台绑定 AICore 配置这从源码层面印证了上述产品支持矩阵。功能说明在 MoE 网络中的角色在 MoEMixture of Experts混合专家分布式推理网络中各张卡上的 token 需要先经过 AlltoAll 通信从其他卡上拿到需要计算的 token。此时每张卡拿到的是“按来源卡rank顺序”排列的 token 集合而后续专家计算需要按“专家expert顺序”组织数据。MoeReRoutingV2 正是完成这一关键重排的算子将 token 从 rank 序源卡顺序重新排列为 expert 序专家顺序同时输出每个 token 在原排布方式中的索引permute_token_idx与每个专家处理的 token 数expert_token_num供后续算子定位与切片使用。与 MoeReRouting 的差异相较于原 MoeReRouting 算子MoeReRoutingV2 新增了可选输入expert_topk_weight与可选输出permute_topk_weight支持对 TopK 权重按专家顺序进行重排使得permute_topk_weight与permute_tokens一一对应。expert_topk_weight与permute_topk_weight必须同时传入或同时不传入这一约束在 算子 IR 原型 的注释以及 Tiling 实现 中都有显式校验。计算公式通过双重求和计算当前 token 在源位置的偏移量$$ SrcOffset \sum_{i0}^{cur_rank} \left( \sum_{j0}^{cur_expert} {expert_token_num_per_rank}(i,j) \right) $$通过双重求和计算当前 token 在目标位置的偏移量$$ DstOffset \sum_{j0}^{cur_expert} \left( \sum_{i0}^{cur_rank} {expert_token_num_per_rank}(i,j) \right) $$符号约定SrcOffset指当前需要移动的 token 源偏移根据输入expert_token_num_per_rank的值进行计算DstOffset指当前需要移动的 token 目的偏移cur_rank是expert_token_num_per_rank的纵轴索引表示该 token 原本所在的卡cur_expert是expert_token_num_per_rank的横轴索引表示该 token 由卡上专家cur_expert计算。token、per_token_scales 和 expert_topk_weight 的搬运偏移量完全一致$$ permute_tokens[DstOffset k] tokens[SrcOffset k] $$$$ permute_per_token_scales[DstOffset k] per_token_scales[SrcOffset k] $$$$ permute_topk_weight[DstOffset k] expert_topk_weight[SrcOffset k] $$其中 k 表示当前 expert 下第 k 个 token 的偏移0 ≤ k currTokenNum。由于 topkWeight 与 token 一一对应搬运偏移量与 token 完全一致直接复用 token 的 SrcOffset 和 DstOffset。数据流定位在典型 MoE 分布式推理链路中MoeReRoutingV2 承接 AlltoAll 的输出按 rank 序的 token输出按 expert 序的 token 及其配套的 scale、索引、专家 token 数供后续专家 FFN 计算使用。结合仓库中相关算子如 moe_finalize_routing_v2、moe_token_permute 等可以推断该算子位于“跨卡通信完成之后、专家计算之前”的数据整理阶段是分布式 MoE 推理流水中的关键一步。参数说明以下为算子的完整输入、输出与属性参数表参数名输入/输出/属性描述数据类型数据格式tokens输入表示待重新排布的 token要求 2 维shape 为 [A, H]FLOAT16、BF16、INT8、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8、FLOAT4_E2M1、FLOAT4_E1M2NDexpert_token_num_per_rank输入二维矩阵元素 [i, j] 表示当前卡上从卡 i 获取到的专家 j 处理的 token 数要求 2 维shape 为 [N, E]取值必须大于 0INT32、INT64NDper_token_scales可选输入表示每个 token 对应的 scale需要随 token 同样进行重新排布。支持 1 维 shape [A]float32、2 维 shape [A, S]float32、3 维 shape [A, K/64, 2]float8_e8m0用于 FP8 量化 token。默认不输入表示不使用 scale此时输出permute_per_token_scales的值无意义FLOAT、FLOAT8_E8M0NDexpert_topk_weight可选输入表示每个 token 对应的 topk 权重值需要随 token 同样进行重新排布。要求 2 维shape 为 [A, 1]与permute_topk_weight联动必须同时传入或同时不传入FLOATNDpermute_tokens输出表示重新排布后的 tokenshape 为 [A, H]数据类型同 tokens与 tokens 一致NDpermute_per_token_scales输出表示重新排布后的 per_token_scales。per_token_scales输入时 shape 和数据类型与其一致未输入时 shape 为 [A]数据类型为 float32该输出无意义FLOAT、FLOAT8_E8M0NDpermute_token_idx输出表示每个 token 在原排布方式的索引shape 为 [A]INT32NDexpert_token_num输出表示每个专家处理的 token 数shape 为 [E]数据类型同expert_token_num_per_rankINT32、INT64NDpermute_topk_weight可选输出表示重新排布后的 topk 权重与 permute_tokens 一一对应。expert_topk_weight输入时必须同时输出shape 为 [A, 1]未输入时输出为空 tensorshape 为 (0,)。与expert_topk_weight联动必须同时传入或同时不传入FLOATNDexpert_token_num_type可选属性表示输出 expert_token_num 的模式。0 为 cumsum 模式1 为 count 模式默认值为 1INT64-idx_type可选属性表示输出 permute_token_idx 的索引类型。0 为 gather 索引1 为 scatter 索引默认值为 0INT64-参数取值细节expert_token_num_per_rank中行第 0 维对应 rank卡列第 1 维对应 expert专家语义为“当前卡从某卡获取的由某专家处理的 token 数量”与 SrcOffset/DstOffset 公式中的双重求和一一对应。permute_token_idx的输出语义由idx_type控制idx_type0输出 gather 索引取数索引idx_type1输出 scatter 索引写回索引。expert_topk_weight的数据类型在 Tiling 实现 中被硬性校验为DT_FLOAT且要求 2 维、第 0 维等于 token 总数 A、第 1 维必须为 1否则直接返回GRAPH_FAILED。约束说明Tensor 中 shape 使用的变量说明A表示 token 个数取值要求 Sum(expert_token_num_per_rank)AH表示 token 长度取值要求 0 H 16384N表示卡数取值无限制E表示卡上的专家数取值无限制。输入值域限制expert_token_num_type当前只支持为 1count 模式expert_topk_weight输入要求为 2 维shape 为 [A, 1]数据类型仅支持 FLOAT。输出类型限制expert_token_num类型应与输入的expert_token_num_per_rank类型保持一致。使用场景限制该接口支持推理场景下使用该接口支持单算子模式和 TorchAir 图模式调用默认支持确定性计算。源码级实现解析算子 IR 原型op_graph算子 IR 原型 通过REG_OP(MoeReRoutingV2)声明了完整的输入输出与属性契约REG_OP(MoeReRoutingV2) .INPUT(tokens, TensorType({DT_FLOAT16, DT_BF16, DT_INT8, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN, DT_HIFLOAT8, DT_FLOAT4_E2M1, DT_FLOAT4_E1M2})) .INPUT(expert_token_num_per_rank, TensorType({DT_INT32, DT_INT64})) .OPTIONAL_INPUT(per_token_scales, TensorType({DT_FLOAT, DT_FLOAT8_E8M0})) .OPTIONAL_INPUT(expert_topk_weight, TensorType({DT_FLOAT})) .OUTPUT(permute_tokens, TensorType({DT_FLOAT16, DT_BF16, DT_INT8, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN, DT_HIFLOAT8, DT_FLOAT4_E2M1, DT_FLOAT4_E1M2})) .OUTPUT(permute_per_token_scales, TensorType({DT_FLOAT, DT_FLOAT8_E8M0})) .OUTPUT(permute_token_idx, TensorType({DT_INT32})) .OUTPUT(expert_token_num, TensorType({DT_INT32, DT_INT64})) .OUTPUT(permute_topk_weight, TensorType({DT_FLOAT})) .ATTR(expert_token_num_type, Int, 1) .ATTR(idx_type, Int, 0) .OP_END_FACTORY_REG(MoeReRoutingV2)从该原型可见tokens与permute_tokens支持全系列的浮点/量化数据类型包括 FP8 的 E5M2/E4M3FN、HIFLOAT8 以及 FP4 的 E2M1/E1M2expert_token_num_per_rank与expert_token_num支持 INT32/INT64permute_token_idx固定为 INT32TopK 权重相关输入输出固定为 FLOAT。InferShape 与 InferDtypeop_hostInferShape 实现 定义了各输出的 shape 推导逻辑permute_tokens的 shape 与tokens完全一致permute_per_token_scales有per_token_scales输入时取其 shape否则退化为 1 维 [A]permute_token_idx为 1 维 [A]expert_token_num为 1 维 [E]取expert_token_num_per_rank的第 1 维permute_topk_weight有expert_topk_weight输入时取其 shape否则为 [0, 1] 的空 shape。dtype 推导方面permute_tokens跟随tokenspermute_per_token_scales跟随per_token_scales未输入时为 FLOATpermute_token_idx固定 INT32expert_token_num跟随expert_token_num_per_rankpermute_topk_weight跟随expert_topk_weight未输入时为 FLOAT。Tiling 与 Kernelop_host / op_kernelTiling 实现 通过继承原 MoeReRouting 的MoeReRoutingReTiling/MoeReRoutingRTiling完成核心切分逻辑并在GetShapeAttrsInfo中额外校验expert_topk_weight的 dtype 与 shape必须为 2 维、第 0 维等于 tokenSum、第 1 维等于 1同时校验“未输入expert_topk_weight时不得提供有效permute_topk_weight输出”。Tiling 模板按任务类型注册了多种组合如MoeReRoutingV2ReTiling10000、MoeReRoutingV2RTiling11000并在 tiling.h 中为不同 TilingKey 绑定对应的 TilingData 结构。Tiling 基类 在TilingPrepare4MoeReRoutingV2阶段读取平台信息AIV 核数、UB 内存大小、SoC 版本并写入编译信息为后续切分提供硬件参数。Kernel 入口 展示了算子内核的多模板分派策略根据是否有 scale、scale 的类型float 或 float8_e8m0以及 token 是否为 FP8/FP4 量化类型选择不同的 Regbase 模板实例如MoeReRoutingReRegbaseDTYPE_TOKENS, DTYPE_EXPERT_TOKEN_NUM_PER_RANK, float, false等。值得注意的是FP8/HIF8/FP4 这类 1 字节打包数据统一复用int8_t模板从而减少模板实例数量。算子配置与精度档位算子二进制配置 按 token 数据类型float16、bfloat16、int8、float8_e5m2、float8_e4m3fn 等× 计数数据类型int32/int64组合生成多档 bin如MoeReRoutingV2_fp16_s32、MoeReRoutingV2_bf16_s64、MoeReRoutingV2_fp8_e5m2_s32_fp8_e8m0等每个档位声明了输入输出的 dtype、ND 格式与可选参数类型运行时根据实际 dtype 组合匹配对应内核。调用方式图模式调用GEIR 构图图模式通过算子 IR 构图方式调用 MoeReRoutingV2完整样例见 test_geir_moe_re_routing_v2.cpp。其核心流程为使用 算子 IR 原型 生成算子对象op::MoeReRoutingV2(test_geir_moe_re_routing_v2)定义输入/输出 shape示例中 A11、H79、N11、E1、S1依次添加 4 个输入tokens、expert_token_num_per_rank、per_token_scales、expert_topk_weight与 5 个输出permute_tokens、permute_per_token_scales、permute_token_idx、expert_token_num、permute_topk_weight顺序严格匹配 proto.h设置属性expert_token_num_type1、idx_type0通过 GE Session 完成AddGraph → RunGraph并将输入/输出数据落盘为.bin文件tc_geir_test_moe_re_routing_v2_npu_input_*.bin/_output_*.bin便于后续比对验证。PyTorch API 调用单算子模式PyTorch 接口完整说明见 torchapi_moe_re_routing.md函数原型如下cann_ops_transformer.moe_re_routing( tokens: torch.Tensor, expert_token_num_per_rank: torch.Tensor, *, per_token_scales: Optional[torch.Tensor] None, expert_topk_weight: Optional[torch.Tensor] None, expert_token_num_type: int 1, idx_type: int 0, ) - Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]单算子模式调用示例import torch import torch_npu import cann_ops_transformer tokens_num 16384 tokens_length 7168 rank_num 16 expert_num 16 tokens torch.randint(low-10, high20, size(tokens_num, tokens_length), dtypetorch.int8).npu() expert_token_num_per_rank torch.ones(rank_num, expert_num, dtypetorch.int32).npu() per_token_scales torch.randn(tokens_num, dtypetorch.float32).npu() # 不传入 expert_topk_weightpermute_topk_weight 为空 tensor permute_tokens, permute_per_token_scales, permute_token_idx, expert_token_num, permute_topk_weight \ cann_ops_transformer.moe_re_routing(tokens, expert_token_num_per_rank, per_token_scalesper_token_scales, expert_token_num_type1, idx_type0) # 传入 expert_topk_weightpermute_topk_weight 为有效数据 expert_topk_weight torch.randn(tokens_num, 1, dtypetorch.float32).npu() permute_tokens, permute_per_token_scales, permute_token_idx, expert_token_num, permute_topk_weight \ cann_ops_transformer.moe_re_routing(tokens, expert_token_num_per_rank, per_token_scalesper_token_scales, expert_topk_weightexpert_topk_weight, expert_token_num_type1, idx_type0)注意示例中expert_token_num_per_rank torch.ones(rank_num, expert_num)使得每个 (rank, expert) 组合的 token 数均为 1此时 Sum(expert_token_num_per_rank) 16×16 256实际使用时该矩阵各元素为大于 0 的计数值且总和必须等于 Atoken 总数。PyTorch API 调用TorchAir 图模式TorchAir 图模式通过torch.compile将模型编译到 NPU 后端执行import torch import torch.nn as nn import torch_npu import torchair as tng from torchair.configs.compiler_config import CompilerConfig import cann_ops_transformer config CompilerConfig() config.experimental_config.keep_inference_input_mutations True npu_backend tng.get_npu_backend(compiler_configconfig) class MoeReRoutingModel(nn.Module): def __init__(self): super().__init__() def forward(self, tokens, expert_token_num_per_rank, *, per_token_scalesNone, expert_topk_weightNone, expert_token_num_type1, idx_type0): return cann_ops_transformer.moe_re_routing(tokens, expert_token_num_per_rank, per_token_scalesper_token_scales, expert_topk_weightexpert_topk_weight, expert_token_num_typeexpert_token_num_type, idx_typeidx_type) def main(): tokens_num 16384 tokens_length 7168 rank_num 16 expert_num 16 tokens torch.randint(low-10, high20, size(tokens_num, tokens_length), dtypetorch.int8).npu() expert_token_num_per_rank torch.ones(rank_num, expert_num, dtypetorch.int32).npu() per_token_scales torch.randn(tokens_num, dtypetorch.float32).npu() expert_topk_weight torch.randn(tokens_num, 1, dtypetorch.float32).npu() model MoeReRoutingModel().npu() model torch.compile(model, backendnpu_backend, dynamicFalse) permute_tokens, permute_per_token_scales, permute_token_idx, expert_token_num, permute_topk_weight \ model(tokens, expert_token_num_per_rank, per_token_scalesper_token_scales, expert_topk_weightexpert_topk_weight, expert_token_num_type1, idx_type0) if __name__ __main__: main()Torch 扩展的图转换细节图转换模块 展示了图模式下的底层转换逻辑将torch.ops.cann_ops_transformer.moe_re_routing.default通过register_fx_node_ge_converter注册为 GE 的MoeReRoutingV2算子节点对于 FP4 量化类型float4_e2m1/float4_e1m2若 tokens 以 uint8 存储会先执行Bitcast Reshape解包为 FP4 语义算子执行后再打包回 uint8保证外部接口对用户透明。此外Torch 前端封装 提供了完整的 Meta 注册与校验tokens/expert_token_num_per_rank要求至少 2 维expert_token_num_type与idx_type必须为 0 或 1并支持通过tokens_dtype参数指定 uint8 存储下的实际数据类型hifloat8、float4_e2m1fn_x2、float4_e1m2fn_x2。Meta 函数中同时给出了各输出的 shape/dtype 推导与 op_host 侧 InferShape 逻辑保持一致。测试与验证仓库在 tests/st/arch35/test_moe_re_routing_v2.csv 中提供了 Ascend 950arch35上的 ST 模糊测试用例覆盖了以下验证维度dtype 组合float16/bfloat16/int8 的 tokens 搭配 int32/int64 的计数矩阵scale 支持 float32 与 float8_e8m0可选输入开关per_token_scales分别以None和有效 tensor 形式出现expert_topk_weight始终传入并验证permute_topk_weight输出属性组合expert_token_num_type1与idx_type的 0/1 两种取值均有覆盖shape 多样性A 从个位数到上万如 (52, 21)、(35, 15267)、(28, 6057)N、E 也覆盖多种取值如 (4,2)、(23,2)、(14,1) 等精度判定以precision_tolerances(0.001, 0.001)与absolute_precision1.00E-08校验输出。同时 golden.py 提供了参照实现可用于离线核对 SrcOffset/DstOffset 公式与各输出之间的对应关系。op_host 侧还有 InferShape 单测 与 Tiling 单测分别验证 shape 推导与切分信息的正确性。小结MoeReRoutingV2 是 CANN ops-transformer 中面向 Ascend 950 系列硬件的 MoE 分布式推理数据整理算子在 AlltoAll 通信与专家计算之间承担 token/scale/topk 权重的按专家顺序重排任务。其核心要点可归纳为通过expert_token_num_per_rank双重求和得到每个 token 的源/目的偏移完成 rank 序到 expert 序的转换相比原 MoeReRouting 新增expert_topk_weight↔permute_topk_weight联动搬运且必须成对出现支持 FP16/BF16/INT8/FP8/FP4 等多种 token 量化类型expert_token_num_type当前仅支持 count 模式1提供图模式GEIR 构图与 PyTorch 单算子/TorchAir 图模式两种调用路径均可在仓库中找到对应的示例与实现源码。如需深入了解其计算细节与实现可继续阅读 算子 IR 原型、InferShape 实现、Tiling 实现 与 Kernel 实现。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer FusedExpert 算子全解析MoE 分布式推理中的专家 ID 与 TopK 权重融合拼接CANN ops transformer FusedExpert 算子全解析MoE 分布式推理中的专家 ID 与 TopK 权重融合拼接 导读 FusedEx算子库人工智能深度学习AscendCANN ops-transformer FfnWorkerBatching 算子Attention/FFN 分离部署下 MoE 场景的 token 按专家重排CANN ops transformer FfnWorkerBatching 算子Attention/FFN 分离部署下 MoE 场景的 token 按专家重算子库人工智能深度学习AscendCANN ops-transformer MoE 系列算子 MoeReRoutingAlltoAll 之后的 Token 按专家重排技术解析与使用指南CANN ops transformer MoE 系列算子 MoeReRoutingAlltoAll 之后的 Token 按专家重排技术解析与使用指南 Moe算子库人工智能深度学习Ascend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN ops-transformer MoeReRoutingV2 算子实战：MoE 分布式推理中的 Token 按专家顺序重排与 TopK 权重联动搬运
算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载本文围绕 CANN ops-transformer 仓库中 moe_re_routing_v2/README.md 及其配套的 PyTorch 接口文档系统讲解 MoeReRoutingV2 算子的功能定位、数学原理、输入输出参数、约束条件并结合仓库源码算子 IR 原型、InferShape、Tiling、Kernel 模板与 Torch 扩展与图模式/单算子调用示例帮助读者在 Ascend 950 系列硬件上正确使用该算子完成 MoE 网络中 AlltoAll 之后的 Token 重排并让 TopK 权重与 Token 一一对应地同步搬运。产品支持情况MoeReRoutingV2 算子仅在 Ascend 950 系列上可用其余产品线暂不支持产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从 算子定义文件 可以看出注册算子时通过this-AICore().AddConfig(ascend950, regbaseCfg)仅为 ascend950 平台绑定 AICore 配置这从源码层面印证了上述产品支持矩阵。功能说明在 MoE 网络中的角色在 MoEMixture of Experts混合专家分布式推理网络中各张卡上的 token 需要先经过 AlltoAll 通信从其他卡上拿到需要计算的 token。此时每张卡拿到的是“按来源卡rank顺序”排列的 token 集合而后续专家计算需要按“专家expert顺序”组织数据。MoeReRoutingV2 正是完成这一关键重排的算子将 token 从 rank 序源卡顺序重新排列为 expert 序专家顺序同时输出每个 token 在原排布方式中的索引permute_token_idx与每个专家处理的 token 数expert_token_num供后续算子定位与切片使用。与 MoeReRouting 的差异相较于原 MoeReRouting 算子MoeReRoutingV2 新增了可选输入expert_topk_weight与可选输出permute_topk_weight支持对 TopK 权重按专家顺序进行重排使得permute_topk_weight与permute_tokens一一对应。expert_topk_weight与permute_topk_weight必须同时传入或同时不传入这一约束在 算子 IR 原型 的注释以及 Tiling 实现 中都有显式校验。计算公式通过双重求和计算当前 token 在源位置的偏移量$$ SrcOffset \sum_{i0}^{cur_rank} \left( \sum_{j0}^{cur_expert} {expert_token_num_per_rank}(i,j) \right) $$通过双重求和计算当前 token 在目标位置的偏移量$$ DstOffset \sum_{j0}^{cur_expert} \left( \sum_{i0}^{cur_rank} {expert_token_num_per_rank}(i,j) \right) $$符号约定SrcOffset指当前需要移动的 token 源偏移根据输入expert_token_num_per_rank的值进行计算DstOffset指当前需要移动的 token 目的偏移cur_rank是expert_token_num_per_rank的纵轴索引表示该 token 原本所在的卡cur_expert是expert_token_num_per_rank的横轴索引表示该 token 由卡上专家cur_expert计算。token、per_token_scales 和 expert_topk_weight 的搬运偏移量完全一致$$ permute_tokens[DstOffset k] tokens[SrcOffset k] $$$$ permute_per_token_scales[DstOffset k] per_token_scales[SrcOffset k] $$$$ permute_topk_weight[DstOffset k] expert_topk_weight[SrcOffset k] $$其中 k 表示当前 expert 下第 k 个 token 的偏移0 ≤ k currTokenNum。由于 topkWeight 与 token 一一对应搬运偏移量与 token 完全一致直接复用 token 的 SrcOffset 和 DstOffset。数据流定位在典型 MoE 分布式推理链路中MoeReRoutingV2 承接 AlltoAll 的输出按 rank 序的 token输出按 expert 序的 token 及其配套的 scale、索引、专家 token 数供后续专家 FFN 计算使用。结合仓库中相关算子如 moe_finalize_routing_v2、moe_token_permute 等可以推断该算子位于“跨卡通信完成之后、专家计算之前”的数据整理阶段是分布式 MoE 推理流水中的关键一步。参数说明以下为算子的完整输入、输出与属性参数表参数名输入/输出/属性描述数据类型数据格式tokens输入表示待重新排布的 token要求 2 维shape 为 [A, H]FLOAT16、BF16、INT8、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8、FLOAT4_E2M1、FLOAT4_E1M2NDexpert_token_num_per_rank输入二维矩阵元素 [i, j] 表示当前卡上从卡 i 获取到的专家 j 处理的 token 数要求 2 维shape 为 [N, E]取值必须大于 0INT32、INT64NDper_token_scales可选输入表示每个 token 对应的 scale需要随 token 同样进行重新排布。支持 1 维 shape [A]float32、2 维 shape [A, S]float32、3 维 shape [A, K/64, 2]float8_e8m0用于 FP8 量化 token。默认不输入表示不使用 scale此时输出permute_per_token_scales的值无意义FLOAT、FLOAT8_E8M0NDexpert_topk_weight可选输入表示每个 token 对应的 topk 权重值需要随 token 同样进行重新排布。要求 2 维shape 为 [A, 1]与permute_topk_weight联动必须同时传入或同时不传入FLOATNDpermute_tokens输出表示重新排布后的 tokenshape 为 [A, H]数据类型同 tokens与 tokens 一致NDpermute_per_token_scales输出表示重新排布后的 per_token_scales。per_token_scales输入时 shape 和数据类型与其一致未输入时 shape 为 [A]数据类型为 float32该输出无意义FLOAT、FLOAT8_E8M0NDpermute_token_idx输出表示每个 token 在原排布方式的索引shape 为 [A]INT32NDexpert_token_num输出表示每个专家处理的 token 数shape 为 [E]数据类型同expert_token_num_per_rankINT32、INT64NDpermute_topk_weight可选输出表示重新排布后的 topk 权重与 permute_tokens 一一对应。expert_topk_weight输入时必须同时输出shape 为 [A, 1]未输入时输出为空 tensorshape 为 (0,)。与expert_topk_weight联动必须同时传入或同时不传入FLOATNDexpert_token_num_type可选属性表示输出 expert_token_num 的模式。0 为 cumsum 模式1 为 count 模式默认值为 1INT64-idx_type可选属性表示输出 permute_token_idx 的索引类型。0 为 gather 索引1 为 scatter 索引默认值为 0INT64-参数取值细节expert_token_num_per_rank中行第 0 维对应 rank卡列第 1 维对应 expert专家语义为“当前卡从某卡获取的由某专家处理的 token 数量”与 SrcOffset/DstOffset 公式中的双重求和一一对应。permute_token_idx的输出语义由idx_type控制idx_type0输出 gather 索引取数索引idx_type1输出 scatter 索引写回索引。expert_topk_weight的数据类型在 Tiling 实现 中被硬性校验为DT_FLOAT且要求 2 维、第 0 维等于 token 总数 A、第 1 维必须为 1否则直接返回GRAPH_FAILED。约束说明Tensor 中 shape 使用的变量说明A表示 token 个数取值要求 Sum(expert_token_num_per_rank)AH表示 token 长度取值要求 0 H 16384N表示卡数取值无限制E表示卡上的专家数取值无限制。输入值域限制expert_token_num_type当前只支持为 1count 模式expert_topk_weight输入要求为 2 维shape 为 [A, 1]数据类型仅支持 FLOAT。输出类型限制expert_token_num类型应与输入的expert_token_num_per_rank类型保持一致。使用场景限制该接口支持推理场景下使用该接口支持单算子模式和 TorchAir 图模式调用默认支持确定性计算。源码级实现解析算子 IR 原型op_graph算子 IR 原型 通过REG_OP(MoeReRoutingV2)声明了完整的输入输出与属性契约REG_OP(MoeReRoutingV2) .INPUT(tokens, TensorType({DT_FLOAT16, DT_BF16, DT_INT8, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN, DT_HIFLOAT8, DT_FLOAT4_E2M1, DT_FLOAT4_E1M2})) .INPUT(expert_token_num_per_rank, TensorType({DT_INT32, DT_INT64})) .OPTIONAL_INPUT(per_token_scales, TensorType({DT_FLOAT, DT_FLOAT8_E8M0})) .OPTIONAL_INPUT(expert_topk_weight, TensorType({DT_FLOAT})) .OUTPUT(permute_tokens, TensorType({DT_FLOAT16, DT_BF16, DT_INT8, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN, DT_HIFLOAT8, DT_FLOAT4_E2M1, DT_FLOAT4_E1M2})) .OUTPUT(permute_per_token_scales, TensorType({DT_FLOAT, DT_FLOAT8_E8M0})) .OUTPUT(permute_token_idx, TensorType({DT_INT32})) .OUTPUT(expert_token_num, TensorType({DT_INT32, DT_INT64})) .OUTPUT(permute_topk_weight, TensorType({DT_FLOAT})) .ATTR(expert_token_num_type, Int, 1) .ATTR(idx_type, Int, 0) .OP_END_FACTORY_REG(MoeReRoutingV2)从该原型可见tokens与permute_tokens支持全系列的浮点/量化数据类型包括 FP8 的 E5M2/E4M3FN、HIFLOAT8 以及 FP4 的 E2M1/E1M2expert_token_num_per_rank与expert_token_num支持 INT32/INT64permute_token_idx固定为 INT32TopK 权重相关输入输出固定为 FLOAT。InferShape 与 InferDtypeop_hostInferShape 实现 定义了各输出的 shape 推导逻辑permute_tokens的 shape 与tokens完全一致permute_per_token_scales有per_token_scales输入时取其 shape否则退化为 1 维 [A]permute_token_idx为 1 维 [A]expert_token_num为 1 维 [E]取expert_token_num_per_rank的第 1 维permute_topk_weight有expert_topk_weight输入时取其 shape否则为 [0, 1] 的空 shape。dtype 推导方面permute_tokens跟随tokenspermute_per_token_scales跟随per_token_scales未输入时为 FLOATpermute_token_idx固定 INT32expert_token_num跟随expert_token_num_per_rankpermute_topk_weight跟随expert_topk_weight未输入时为 FLOAT。Tiling 与 Kernelop_host / op_kernelTiling 实现 通过继承原 MoeReRouting 的MoeReRoutingReTiling/MoeReRoutingRTiling完成核心切分逻辑并在GetShapeAttrsInfo中额外校验expert_topk_weight的 dtype 与 shape必须为 2 维、第 0 维等于 tokenSum、第 1 维等于 1同时校验“未输入expert_topk_weight时不得提供有效permute_topk_weight输出”。Tiling 模板按任务类型注册了多种组合如MoeReRoutingV2ReTiling10000、MoeReRoutingV2RTiling11000并在 tiling.h 中为不同 TilingKey 绑定对应的 TilingData 结构。Tiling 基类 在TilingPrepare4MoeReRoutingV2阶段读取平台信息AIV 核数、UB 内存大小、SoC 版本并写入编译信息为后续切分提供硬件参数。Kernel 入口 展示了算子内核的多模板分派策略根据是否有 scale、scale 的类型float 或 float8_e8m0以及 token 是否为 FP8/FP4 量化类型选择不同的 Regbase 模板实例如MoeReRoutingReRegbaseDTYPE_TOKENS, DTYPE_EXPERT_TOKEN_NUM_PER_RANK, float, false等。值得注意的是FP8/HIF8/FP4 这类 1 字节打包数据统一复用int8_t模板从而减少模板实例数量。算子配置与精度档位算子二进制配置 按 token 数据类型float16、bfloat16、int8、float8_e5m2、float8_e4m3fn 等× 计数数据类型int32/int64组合生成多档 bin如MoeReRoutingV2_fp16_s32、MoeReRoutingV2_bf16_s64、MoeReRoutingV2_fp8_e5m2_s32_fp8_e8m0等每个档位声明了输入输出的 dtype、ND 格式与可选参数类型运行时根据实际 dtype 组合匹配对应内核。调用方式图模式调用GEIR 构图图模式通过算子 IR 构图方式调用 MoeReRoutingV2完整样例见 test_geir_moe_re_routing_v2.cpp。其核心流程为使用 算子 IR 原型 生成算子对象op::MoeReRoutingV2(test_geir_moe_re_routing_v2)定义输入/输出 shape示例中 A11、H79、N11、E1、S1依次添加 4 个输入tokens、expert_token_num_per_rank、per_token_scales、expert_topk_weight与 5 个输出permute_tokens、permute_per_token_scales、permute_token_idx、expert_token_num、permute_topk_weight顺序严格匹配 proto.h设置属性expert_token_num_type1、idx_type0通过 GE Session 完成AddGraph → RunGraph并将输入/输出数据落盘为.bin文件tc_geir_test_moe_re_routing_v2_npu_input_*.bin/_output_*.bin便于后续比对验证。PyTorch API 调用单算子模式PyTorch 接口完整说明见 torchapi_moe_re_routing.md函数原型如下cann_ops_transformer.moe_re_routing( tokens: torch.Tensor, expert_token_num_per_rank: torch.Tensor, *, per_token_scales: Optional[torch.Tensor] None, expert_topk_weight: Optional[torch.Tensor] None, expert_token_num_type: int 1, idx_type: int 0, ) - Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]单算子模式调用示例import torch import torch_npu import cann_ops_transformer tokens_num 16384 tokens_length 7168 rank_num 16 expert_num 16 tokens torch.randint(low-10, high20, size(tokens_num, tokens_length), dtypetorch.int8).npu() expert_token_num_per_rank torch.ones(rank_num, expert_num, dtypetorch.int32).npu() per_token_scales torch.randn(tokens_num, dtypetorch.float32).npu() # 不传入 expert_topk_weightpermute_topk_weight 为空 tensor permute_tokens, permute_per_token_scales, permute_token_idx, expert_token_num, permute_topk_weight \ cann_ops_transformer.moe_re_routing(tokens, expert_token_num_per_rank, per_token_scalesper_token_scales, expert_token_num_type1, idx_type0) # 传入 expert_topk_weightpermute_topk_weight 为有效数据 expert_topk_weight torch.randn(tokens_num, 1, dtypetorch.float32).npu() permute_tokens, permute_per_token_scales, permute_token_idx, expert_token_num, permute_topk_weight \ cann_ops_transformer.moe_re_routing(tokens, expert_token_num_per_rank, per_token_scalesper_token_scales, expert_topk_weightexpert_topk_weight, expert_token_num_type1, idx_type0)注意示例中expert_token_num_per_rank torch.ones(rank_num, expert_num)使得每个 (rank, expert) 组合的 token 数均为 1此时 Sum(expert_token_num_per_rank) 16×16 256实际使用时该矩阵各元素为大于 0 的计数值且总和必须等于 Atoken 总数。PyTorch API 调用TorchAir 图模式TorchAir 图模式通过torch.compile将模型编译到 NPU 后端执行import torch import torch.nn as nn import torch_npu import torchair as tng from torchair.configs.compiler_config import CompilerConfig import cann_ops_transformer config CompilerConfig() config.experimental_config.keep_inference_input_mutations True npu_backend tng.get_npu_backend(compiler_configconfig) class MoeReRoutingModel(nn.Module): def __init__(self): super().__init__() def forward(self, tokens, expert_token_num_per_rank, *, per_token_scalesNone, expert_topk_weightNone, expert_token_num_type1, idx_type0): return cann_ops_transformer.moe_re_routing(tokens, expert_token_num_per_rank, per_token_scalesper_token_scales, expert_topk_weightexpert_topk_weight, expert_token_num_typeexpert_token_num_type, idx_typeidx_type) def main(): tokens_num 16384 tokens_length 7168 rank_num 16 expert_num 16 tokens torch.randint(low-10, high20, size(tokens_num, tokens_length), dtypetorch.int8).npu() expert_token_num_per_rank torch.ones(rank_num, expert_num, dtypetorch.int32).npu() per_token_scales torch.randn(tokens_num, dtypetorch.float32).npu() expert_topk_weight torch.randn(tokens_num, 1, dtypetorch.float32).npu() model MoeReRoutingModel().npu() model torch.compile(model, backendnpu_backend, dynamicFalse) permute_tokens, permute_per_token_scales, permute_token_idx, expert_token_num, permute_topk_weight \ model(tokens, expert_token_num_per_rank, per_token_scalesper_token_scales, expert_topk_weightexpert_topk_weight, expert_token_num_type1, idx_type0) if __name__ __main__: main()Torch 扩展的图转换细节图转换模块 展示了图模式下的底层转换逻辑将torch.ops.cann_ops_transformer.moe_re_routing.default通过register_fx_node_ge_converter注册为 GE 的MoeReRoutingV2算子节点对于 FP4 量化类型float4_e2m1/float4_e1m2若 tokens 以 uint8 存储会先执行Bitcast Reshape解包为 FP4 语义算子执行后再打包回 uint8保证外部接口对用户透明。此外Torch 前端封装 提供了完整的 Meta 注册与校验tokens/expert_token_num_per_rank要求至少 2 维expert_token_num_type与idx_type必须为 0 或 1并支持通过tokens_dtype参数指定 uint8 存储下的实际数据类型hifloat8、float4_e2m1fn_x2、float4_e1m2fn_x2。Meta 函数中同时给出了各输出的 shape/dtype 推导与 op_host 侧 InferShape 逻辑保持一致。测试与验证仓库在 tests/st/arch35/test_moe_re_routing_v2.csv 中提供了 Ascend 950arch35上的 ST 模糊测试用例覆盖了以下验证维度dtype 组合float16/bfloat16/int8 的 tokens 搭配 int32/int64 的计数矩阵scale 支持 float32 与 float8_e8m0可选输入开关per_token_scales分别以None和有效 tensor 形式出现expert_topk_weight始终传入并验证permute_topk_weight输出属性组合expert_token_num_type1与idx_type的 0/1 两种取值均有覆盖shape 多样性A 从个位数到上万如 (52, 21)、(35, 15267)、(28, 6057)N、E 也覆盖多种取值如 (4,2)、(23,2)、(14,1) 等精度判定以precision_tolerances(0.001, 0.001)与absolute_precision1.00E-08校验输出。同时 golden.py 提供了参照实现可用于离线核对 SrcOffset/DstOffset 公式与各输出之间的对应关系。op_host 侧还有 InferShape 单测 与 Tiling 单测分别验证 shape 推导与切分信息的正确性。小结MoeReRoutingV2 是 CANN ops-transformer 中面向 Ascend 950 系列硬件的 MoE 分布式推理数据整理算子在 AlltoAll 通信与专家计算之间承担 token/scale/topk 权重的按专家顺序重排任务。其核心要点可归纳为通过expert_token_num_per_rank双重求和得到每个 token 的源/目的偏移完成 rank 序到 expert 序的转换相比原 MoeReRouting 新增expert_topk_weight↔permute_topk_weight联动搬运且必须成对出现支持 FP16/BF16/INT8/FP8/FP4 等多种 token 量化类型expert_token_num_type当前仅支持 count 模式1提供图模式GEIR 构图与 PyTorch 单算子/TorchAir 图模式两种调用路径均可在仓库中找到对应的示例与实现源码。如需深入了解其计算细节与实现可继续阅读 算子 IR 原型、InferShape 实现、Tiling 实现 与 Kernel 实现。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer FusedExpert 算子全解析MoE 分布式推理中的专家 ID 与 TopK 权重融合拼接CANN ops transformer FusedExpert 算子全解析MoE 分布式推理中的专家 ID 与 TopK 权重融合拼接 导读 FusedEx算子库人工智能深度学习AscendCANN ops-transformer FfnWorkerBatching 算子Attention/FFN 分离部署下 MoE 场景的 token 按专家重排CANN ops transformer FfnWorkerBatching 算子Attention/FFN 分离部署下 MoE 场景的 token 按专家重算子库人工智能深度学习AscendCANN ops-transformer MoE 系列算子 MoeReRoutingAlltoAll 之后的 Token 按专家重排技术解析与使用指南CANN ops transformer MoE 系列算子 MoeReRoutingAlltoAll 之后的 Token 按专家重排技术解析与使用指南 Moe算子库人工智能深度学习Ascend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
