CANN ops-transformer MoE 路由反向算子 moe_init_routing_grad 解析公式、参数与 NPU 实现【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读moe_init_routing_grad是 CANN ops-transformer 仓库中 MoEMixture of Experts混合专家路由算子moe_init_routing的反向算子负责把专家扩展后张量expanded_x的梯度按行映射回原始输入x的梯度。本文以 torchapi_moe_init_routing_grad.md 为主体结合算子定义、infershape、kernel 与 torch 扩展源码完整讲解其数学语义、API 参数、三种场景Dropless / Active / DropPad的计算路径、约束条件与单算子/自动反向两种调用方式帮助读者在训练 MoE 模型时正确使用并理解其底层行为。一、背景MoE 路由与梯度回传在 MoE 类大模型中每个 token行会由 gating 网络选出 top-K 个专家随后数据需要被路由到对应专家参与计算。正向算子moe_init_routing在仓库中对应 moe_init_routing_v4 模块底层为aclnnMoeInitRoutingV2完成这一过程输入为原始 token 特征xshape(NUM_ROWS, H)和每个 token 选出的专家索引expert_idxshape(NUM_ROWS, K)对expert_idx排序后按 gather 或 scatter 索引把x搬运为扩展张量expanded_x每行重复 K 次共NUM_ROWS * K行同时输出expanded_row_idx记录expanded_x与x的行映射关系DropPad 场景下还会按专家容量expert_capacity对齐每个专家的 token 数超出部分被丢弃Drop不足部分用 0 填充Pad。训练场景下损失需要回传到原始输入x因此必须实现反向算子将expanded_x的梯度grad_expanded_x沿expanded_row_idx指引的行映射归约回grad_x。这正是moe_init_routing_grad的职责它封装底层aclnnMoeInitRoutingV2Grad接口。二、功能说明与计算公式接口功能moe_init_routing_grad是moe_init_routing的反向算子封装aclnnMoeInitRoutingV2Grad接口。该接口根据正向输出的expanded_row_idx将grad_expanded_x即expanded_x的梯度按行映射回原始输入x的梯度grad_x。计算公式$$ grad_x[i] \sum_{t0}^{K-1} grad_expanded_x[expanded_row_idx[i \cdot K t]] $$其中i范围为[0, NUM_ROWS)K为top_kexpanded_row_idx为正向moe_init_routing的输出。DropPad 场景下expanded_row_idx值为 -1 的项不参与累加。算例推演Kernel 源码 moe_init_routing_v2_grad_with_dropless.h 中的注释给出一个直观例子expanded_x: [6, 4]、expanded_row_idx: [6]值为4 1 5 3 0 2、grad_x: [3, 4]、K2、NUM_ROWS3。此时grad_x[0] grad_expanded_x[4] grad_expanded_x[1]对应expanded_row_idx[0]4、expanded_row_idx[1]1grad_x[1] grad_expanded_x[5] grad_expanded_x[3]对应expanded_row_idx[2]5、expanded_row_idx[3]3grad_x[2] grad_expanded_x[0] grad_expanded_x[2]对应expanded_row_idx[4]0、expanded_row_idx[5]2。即每个原始 token 行的梯度等于其 K 个专家副本对应行的梯度之和这与 MoE 反向传播中同一 token 被多个专家复用、梯度需累加的语义完全一致。三、产品支持情况该接口底层aclnnMoeInitRoutingV2Grad的产品支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持从算子注册代码 moe_init_routing_v2_grad_def.cpp 也可以印证OpAICoreConfig分别向ascend910b、ascend910_93注册了moe_init_routing_v2_grad实现向ascend950注册了moe_init_routing_v2_grad_apt实现op_host/config 目录下也对应存在ascend910b、ascend910_93、ascend950三份算子二进制与 simplified_key 配置。四、函数原型与参数说明函数原型cann_ops_transformer.moe_init_routing_grad( grad_expanded_x: torch.Tensor, expanded_row_idx: torch.Tensor, top_k: int, drop_pad_mode: int 0, active_num: int 0, ) - torch.Tensor该原型与 torch 扩展中的 schema 定义一致见 moe_init_routing_grad.pymoe_init_routing_grad(Tensor grad_expanded_x, Tensor expanded_row_idx, int top_k, int drop_pad_mode0, int active_num0) - Tensor参数说明参数名参数类型可选/必选描述数据类型维度(shape)grad_expanded_xTensor必选正向输出expanded_x的梯度。Dropless 场景下为 2 维张量shape 为(active_num, H)或(NUM_ROWS * K, H)DropPad 场景下为 3 维张量shape 为(expert_num, expert_capacity, H)。float16、bfloat16、float32NDexpanded_row_idxTensor必选正向moe_init_routing的输出expanded_row_idx表示expanded_x与x的行映射关系。为 1 维张量shape 为(NUM_ROWS * K,)。int32NDtop_kint必选正向输入expert_idx的第二维大小 K即每个 token 选出的 top-K 专家个数。取值需大于 0。--drop_pad_modeint可选默认值为 0表示丢弃/填充模式需与正向保持一致。0 表示 Dropless 场景1 表示 DropPad 场景。--active_numint可选默认值为 0表示正向moe_init_routing的active_num参数值。取值需大于等于 0当drop_pad_mode0时生效0 表示非 Active 场景大于 0 表示 Active 场景此时grad_expanded_x的第 0 维大小必须等于active_num。--参数校验在 C wrapper moe_init_routing_grad.cpp 中有完整实现每项均有TORCH_CHECK兜底例如grad_expanded_x必须是 2D 或 3D、dtype 必须是 float16/float32/bfloat16expanded_row_idx必须是 1D int32drop_pad_mode只能是 0 或 1top_k 0active_num 0且drop_pad_mode1时强制grad_expanded_x为 3D。五、返回值说明参数名参数类型可选/必选描述数据类型维度(shape)grad_xTensor必选输入x的梯度。shape 为(NUM_ROWS, H)其中NUM_ROWS expanded_row_idx.numel() // top_kH为grad_expanded_x的最后一维DropPad 场景下为第三维。数据类型与grad_expanded_x一致。与 grad_expanded_x 一致NDgrad_x的 shape 推导逻辑可以从两处源码确认host 侧 infershapemoe_init_routing_v2_grad_infershape.cpp 中输出第 0 维gradXDim0 expandedRowIdxDim0 / topK第 1 维取grad_expanded_x的最后一维DropPad 时为第 2 维否则为第 1 维最终输出固定为 2 维torch 侧 meta 注册moe_init_routing_grad.py 中grad_x_dim0 expanded_row_idx.numel() // top_k、grad_x_dim1按drop_pad_mode 1取grad_expanded_x.size(2)否则取size(1)并返回new_empty张量。六、三种计算场景与 Kernel 实现原理moe_init_routing_grad依据drop_pad_mode与active_num的组合在 NPU 上会落入三种不同的计算路径。算子定义 moe_init_routing_v2_grad_def.cpp 声明了两个必选输入grad_expanded_xfloat16/bfloat16/float32ND 格式与expanded_row_idxint32一个输出grad_x以及三个属性top_k必选、drop_pad_mode默认 0、active_num默认 0。Kernel 入口 moe_init_routing_v2_grad.cpp 为纯 AIV 核KERNEL_TYPE_AIV_ONLY通过 tiling key 在 9 种组合3 种场景 × float/float16/bfloat16中分派到对应的 Compute 类场景触发条件Compute 类关键行为Droplessdrop_pad_mode0且active_num0MoeInitRoutingV2GradDroplessComputemoe_init_routing_v2_grad_with_dropless.hexpanded_row_idx全部有效直接按 K 累加Activedrop_pad_mode0且active_num0MoeInitRoutingV2GradActivateComputemoe_init_routing_v2_grad_with_activate.hxRow active_num的项跳过不累加DropPaddrop_pad_mode1MoeInitRoutingV2GradPositionPad0Computemoe_init_routing_v2_grad_with_pos_drop_and_pad_zero.hxRow -1的项跳过对应正向填充 0 的位置核心累加策略二分累加Binary Accumulation以 Dropless 场景为例GradProcessmoe_init_routing_v2_grad_with_dropless.h将一个 token 的 K 行梯度累加分为 4 步S1先把若干被加数行通过CopyInWithCastFloat拷入 on-chip 的binBuff累加 buffer并做MTE2_V事件同步S2当 K 超过累加 buffer 数量时多出的行直接在对应 buffer 上BinaryAddWithMovIn累加S3其余加数行依次累加进对应 bufferV_S事件同步保证写回顺序S4对 buffer 做指数间隔的递归二分累加stride 逐轮翻倍interval stride * baseStride最终将结果CopyOut到grad_x对应行。Processmoe_init_routing_v2_grad_with_dropless.h按 core 数切分 tokenperCoreElements再按 token 组循环每组用Duplicate将 buffer 清零后逐 token 处理实现多 core 并行。expanded_row_idx以 stride N 方式访问使得同一 token 的 K 个索引分散在不同 core 的独立累加路径上避免写冲突。Active 与 DropPad 场景的差别仅在跳过无效索引的判断上Active 场景xRow activeNum时跳过moe_init_routing_v2_grad_with_activate.hDropPad 场景xRow -1时跳过moe_init_routing_v2_grad_with_pos_drop_and_pad_zero.h且被跳过位置会用Duplicate将 buffer 置 0 后再继续其余累加流程完全一致。aclnn 接口层的空输出优化在 moe_init_routing_v2_grad.cpp 的l0op::MoeInitRoutingV2Grad中输出张量grad_x若为空grad_x-IsEmpty()会直接跳过算子逻辑返回不触发 AICore 任务下发这对应了文档示例中输出为空时不走算子的边界处理。七、约束说明使用该接口时需要满足以下约束与正向调用严格对齐该接口支持训练场景下使用。该接口支持单算子模式和图模式。grad_expanded_x在 Dropless 场景下必须是 2 维张量在 DropPad 场景下必须是 3 维张量。expanded_row_idx必须是 1 维张量其长度必须能被top_k整除。drop_pad_mode仅支持取值 0 或 1。top_k必须大于 0。active_num必须大于等于 0。Dropless 场景且active_num0时grad_expanded_x的第 0 维大小必须与expanded_row_idx的长度一致Dropless 场景且active_num0时grad_expanded_x的第 0 维大小必须等于active_num。grad_expanded_x的最后一维DropPad 场景下为第三维必须与输出grad_x的第二维一致。该算子为 moe_init_routing 的反向算子expanded_row_idx、top_k、drop_pad_mode、active_num需与正向调用保持一致。该反向算子仅支持非量化场景正向quant_mode-1不支持aclnnMoeInitRoutingV4特有特性scale、offset、topk_weight、x_dtype、row_idx_type非 0。当正向使用了这些特性时调用自动反向会抛出NotImplementedError。上述约束在 host 侧 moe_init_routing_v2_grad_infershape.cpp 的CheckParm中同样有图模式下的校验drop_pad_mode必须为 0 或 1、top_k必须大于 0、active_num必须大于等于 0、grad_expanded_x维度数必须匹配场景要求、expanded_row_idx必须为 1 维。八、确定性计算该接口默认支持确定性计算。对同一输入多次调用结果可复现适用于对数值可复现性有要求的训练调试场景。九、配套接口自动反向与手动调用该算子为 moe_init_routing 的反向算子。当正向moe_init_routing中x.requires_grad为 True且未使用aclnnMoeInitRoutingV4特有特性量化、scale、offset、topk_weight等时loss.backward()会自动触发本算子无需手动调用。仅在需要显式控制梯度的场景下保留手动调用路径。自动反向的约束细节在正向文档中有明确说明自动反向仅在正向退化为aclnnMoeInitRoutingV2场景时支持即scale不传入、offset不传入、topk_weight不传入、quant_mode-1、row_idx_type0、x_dtype为 None、drop_pad_mode为 0 或 1active_expert_range不影响反向。此外自动反向仅对x求梯度expert_idx为整数索引张量无梯度expanded_row_idx及其他整数/统计输出同样无梯度。十、调用示例单算子模式调用以下示例在 NPU 上完成一次正向路由 手动显式反向的完整流程来源torchapi_moe_init_routing_grad.mdimport torch import torch_npu from cann_ops_transformer.ops import moe_init_routing, moe_init_routing_grad torch_npu.npu.set_device(0) n 4 h 8 k 2 expert_num 8 drop_pad_mode 0 expert_tokens_num_type 1 expert_tokens_num_flag True quant_mode -1 active_expert_range [0, 4] row_idx_type 0 x torch.randn((n, h), dtypetorch.float32, devicenpu) expert_idx torch.randint(0, expert_num, (n, k), dtypetorch.int32, devicenpu) # 正向 expanded_x, expanded_row_idx, expert_tokens_count, expanded_scale, expanded_topk_weight \ moe_init_routing(x, expert_idx, active_num-1, expert_capacity-1, expert_numexpert_num, drop_pad_modedrop_pad_mode, expert_tokens_num_typeexpert_tokens_num_type, expert_tokens_num_flagexpert_tokens_num_flag, quant_modequant_mode, active_expert_rangeactive_expert_range, row_idx_typerow_idx_type) # 手动调用反向 grad_expanded_x torch.randn_like(expanded_x) grad_x moe_init_routing_grad(grad_expanded_x, expanded_row_idx, top_kk, drop_pad_modedrop_pad_mode, active_num0) print(grad_x.shape) # torch.Size([4, 8])注意示例中正向使用了row_idx_type0gather 索引且quant_mode-1非量化符合自动反向约束但此处选择了手动调用反向以演示显式控制梯度的路径。expanded_scale在非量化且未传scale时返回的数值无意义属正常现象。自动反向调用import torch import torch_npu from cann_ops_transformer.ops import moe_init_routing torch_npu.npu.set_device(0) n 4 h 8 k 2 expert_num 8 x torch.randn((n, h), dtypetorch.float32, devicenpu, requires_gradTrue) expert_idx torch.randint(0, expert_num, (n, k), dtypetorch.int32, devicenpu) # 正向x.requires_gradTrue 时自动启用 autograd expanded_x, expanded_row_idx, expert_tokens_count, expanded_scale, expanded_topk_weight \ moe_init_routing(x, expert_idx, active_num-1, expert_capacity-1, expert_numexpert_num, drop_pad_mode0, expert_tokens_num_type1, expert_tokens_num_flagTrue, quant_mode-1, active_expert_range[0, 4], row_idx_type0) # 反向自动触发 moe_init_routing_grad loss expanded_x.sum() loss.backward() print(x.grad.shape) # torch.Size([4, 8])自动反向通过 torch 扩展层完成moe_init_routing_grad.py注册了Meta实现用于 shape 推导PrivateUse1分发实现最终调用编译出的op_module.moe_init_routing_gradmoe_init_routing_grad.py该 C 函数完成参数校验与输出张量分配后经由ACLNN_CMD(aclnnMoeInitRoutingV2Grad, ...)下发 NPU 计算moe_init_routing_grad.cpp。图模式支持该接口支持 TorchAir 图模式调用。图模式下的调用方式与正向算子一致——将算子调用放入torch.compile(model, backendnpu_backend)的模型中执行正向 moe_init_routing 文档提供了完整的图模式示例可供参考反向算子随正向图一并编译执行。更多参考底层 aclnn 接口文档aclnnMoeInitRoutingV2Grad.mdC 调用示例test_aclnn_moe_init_routing_v2_grad.cpp模块级说明moe_init_routing_v2_grad/README.md单元测试host/op_apitest_aclnn_moe_init_routing_v2_grad.cpp、test_moe_init_routing_v2_grad_infershape.cpp单元测试kerneltest_moe_init_routing_v2_grad.cpp十一、总结moe_init_routing_grad是 CANN ops-transformer 中 MoE 路由反向链路的关键算子本质是一个基于expanded_row_idx行映射的分组归约reduction操作每个原始 token 的梯度是其 K 个专家副本梯度的累加和。在使用时需要抓住三个要点场景对齐drop_pad_mode、active_num必须与正向moe_init_routing调用保持一致并据此选择正确的grad_expanded_x维度2D 或 3D约束匹配仅非量化、不使用 V4 特有特性时支持自动反向否则需手动调用并处理NotImplementedError性能路径NPU kernel 通过多 core 切分、on-chip 累加 buffer 与二分累加策略高效完成 K 路梯度累加Active/DropPad 场景仅以索引有效性判断xRow active_num、xRow -1区分跳过逻辑整体计算路径统一且可确定性复现。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/20 16:53:48 · 阅读 1.2k CANN ops-transformer MoE 路由反向算子 moe_init_routing_grad 解析：公式、参数与 NPU 实现
CANN ops-transformer MoE 路由反向算子 moe_init_routing_grad 解析公式、参数与 NPU 实现【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读moe_init_routing_grad是 CANN ops-transformer 仓库中 MoEMixture of Experts混合专家路由算子moe_init_routing的反向算子负责把专家扩展后张量expanded_x的梯度按行映射回原始输入x的梯度。本文以 torchapi_moe_init_routing_grad.md 为主体结合算子定义、infershape、kernel 与 torch 扩展源码完整讲解其数学语义、API 参数、三种场景Dropless / Active / DropPad的计算路径、约束条件与单算子/自动反向两种调用方式帮助读者在训练 MoE 模型时正确使用并理解其底层行为。一、背景MoE 路由与梯度回传在 MoE 类大模型中每个 token行会由 gating 网络选出 top-K 个专家随后数据需要被路由到对应专家参与计算。正向算子moe_init_routing在仓库中对应 moe_init_routing_v4 模块底层为aclnnMoeInitRoutingV2完成这一过程输入为原始 token 特征xshape(NUM_ROWS, H)和每个 token 选出的专家索引expert_idxshape(NUM_ROWS, K)对expert_idx排序后按 gather 或 scatter 索引把x搬运为扩展张量expanded_x每行重复 K 次共NUM_ROWS * K行同时输出expanded_row_idx记录expanded_x与x的行映射关系DropPad 场景下还会按专家容量expert_capacity对齐每个专家的 token 数超出部分被丢弃Drop不足部分用 0 填充Pad。训练场景下损失需要回传到原始输入x因此必须实现反向算子将expanded_x的梯度grad_expanded_x沿expanded_row_idx指引的行映射归约回grad_x。这正是moe_init_routing_grad的职责它封装底层aclnnMoeInitRoutingV2Grad接口。二、功能说明与计算公式接口功能moe_init_routing_grad是moe_init_routing的反向算子封装aclnnMoeInitRoutingV2Grad接口。该接口根据正向输出的expanded_row_idx将grad_expanded_x即expanded_x的梯度按行映射回原始输入x的梯度grad_x。计算公式$$ grad_x[i] \sum_{t0}^{K-1} grad_expanded_x[expanded_row_idx[i \cdot K t]] $$其中i范围为[0, NUM_ROWS)K为top_kexpanded_row_idx为正向moe_init_routing的输出。DropPad 场景下expanded_row_idx值为 -1 的项不参与累加。算例推演Kernel 源码 moe_init_routing_v2_grad_with_dropless.h 中的注释给出一个直观例子expanded_x: [6, 4]、expanded_row_idx: [6]值为4 1 5 3 0 2、grad_x: [3, 4]、K2、NUM_ROWS3。此时grad_x[0] grad_expanded_x[4] grad_expanded_x[1]对应expanded_row_idx[0]4、expanded_row_idx[1]1grad_x[1] grad_expanded_x[5] grad_expanded_x[3]对应expanded_row_idx[2]5、expanded_row_idx[3]3grad_x[2] grad_expanded_x[0] grad_expanded_x[2]对应expanded_row_idx[4]0、expanded_row_idx[5]2。即每个原始 token 行的梯度等于其 K 个专家副本对应行的梯度之和这与 MoE 反向传播中同一 token 被多个专家复用、梯度需累加的语义完全一致。三、产品支持情况该接口底层aclnnMoeInitRoutingV2Grad的产品支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持从算子注册代码 moe_init_routing_v2_grad_def.cpp 也可以印证OpAICoreConfig分别向ascend910b、ascend910_93注册了moe_init_routing_v2_grad实现向ascend950注册了moe_init_routing_v2_grad_apt实现op_host/config 目录下也对应存在ascend910b、ascend910_93、ascend950三份算子二进制与 simplified_key 配置。四、函数原型与参数说明函数原型cann_ops_transformer.moe_init_routing_grad( grad_expanded_x: torch.Tensor, expanded_row_idx: torch.Tensor, top_k: int, drop_pad_mode: int 0, active_num: int 0, ) - torch.Tensor该原型与 torch 扩展中的 schema 定义一致见 moe_init_routing_grad.pymoe_init_routing_grad(Tensor grad_expanded_x, Tensor expanded_row_idx, int top_k, int drop_pad_mode0, int active_num0) - Tensor参数说明参数名参数类型可选/必选描述数据类型维度(shape)grad_expanded_xTensor必选正向输出expanded_x的梯度。Dropless 场景下为 2 维张量shape 为(active_num, H)或(NUM_ROWS * K, H)DropPad 场景下为 3 维张量shape 为(expert_num, expert_capacity, H)。float16、bfloat16、float32NDexpanded_row_idxTensor必选正向moe_init_routing的输出expanded_row_idx表示expanded_x与x的行映射关系。为 1 维张量shape 为(NUM_ROWS * K,)。int32NDtop_kint必选正向输入expert_idx的第二维大小 K即每个 token 选出的 top-K 专家个数。取值需大于 0。--drop_pad_modeint可选默认值为 0表示丢弃/填充模式需与正向保持一致。0 表示 Dropless 场景1 表示 DropPad 场景。--active_numint可选默认值为 0表示正向moe_init_routing的active_num参数值。取值需大于等于 0当drop_pad_mode0时生效0 表示非 Active 场景大于 0 表示 Active 场景此时grad_expanded_x的第 0 维大小必须等于active_num。--参数校验在 C wrapper moe_init_routing_grad.cpp 中有完整实现每项均有TORCH_CHECK兜底例如grad_expanded_x必须是 2D 或 3D、dtype 必须是 float16/float32/bfloat16expanded_row_idx必须是 1D int32drop_pad_mode只能是 0 或 1top_k 0active_num 0且drop_pad_mode1时强制grad_expanded_x为 3D。五、返回值说明参数名参数类型可选/必选描述数据类型维度(shape)grad_xTensor必选输入x的梯度。shape 为(NUM_ROWS, H)其中NUM_ROWS expanded_row_idx.numel() // top_kH为grad_expanded_x的最后一维DropPad 场景下为第三维。数据类型与grad_expanded_x一致。与 grad_expanded_x 一致NDgrad_x的 shape 推导逻辑可以从两处源码确认host 侧 infershapemoe_init_routing_v2_grad_infershape.cpp 中输出第 0 维gradXDim0 expandedRowIdxDim0 / topK第 1 维取grad_expanded_x的最后一维DropPad 时为第 2 维否则为第 1 维最终输出固定为 2 维torch 侧 meta 注册moe_init_routing_grad.py 中grad_x_dim0 expanded_row_idx.numel() // top_k、grad_x_dim1按drop_pad_mode 1取grad_expanded_x.size(2)否则取size(1)并返回new_empty张量。六、三种计算场景与 Kernel 实现原理moe_init_routing_grad依据drop_pad_mode与active_num的组合在 NPU 上会落入三种不同的计算路径。算子定义 moe_init_routing_v2_grad_def.cpp 声明了两个必选输入grad_expanded_xfloat16/bfloat16/float32ND 格式与expanded_row_idxint32一个输出grad_x以及三个属性top_k必选、drop_pad_mode默认 0、active_num默认 0。Kernel 入口 moe_init_routing_v2_grad.cpp 为纯 AIV 核KERNEL_TYPE_AIV_ONLY通过 tiling key 在 9 种组合3 种场景 × float/float16/bfloat16中分派到对应的 Compute 类场景触发条件Compute 类关键行为Droplessdrop_pad_mode0且active_num0MoeInitRoutingV2GradDroplessComputemoe_init_routing_v2_grad_with_dropless.hexpanded_row_idx全部有效直接按 K 累加Activedrop_pad_mode0且active_num0MoeInitRoutingV2GradActivateComputemoe_init_routing_v2_grad_with_activate.hxRow active_num的项跳过不累加DropPaddrop_pad_mode1MoeInitRoutingV2GradPositionPad0Computemoe_init_routing_v2_grad_with_pos_drop_and_pad_zero.hxRow -1的项跳过对应正向填充 0 的位置核心累加策略二分累加Binary Accumulation以 Dropless 场景为例GradProcessmoe_init_routing_v2_grad_with_dropless.h将一个 token 的 K 行梯度累加分为 4 步S1先把若干被加数行通过CopyInWithCastFloat拷入 on-chip 的binBuff累加 buffer并做MTE2_V事件同步S2当 K 超过累加 buffer 数量时多出的行直接在对应 buffer 上BinaryAddWithMovIn累加S3其余加数行依次累加进对应 bufferV_S事件同步保证写回顺序S4对 buffer 做指数间隔的递归二分累加stride 逐轮翻倍interval stride * baseStride最终将结果CopyOut到grad_x对应行。Processmoe_init_routing_v2_grad_with_dropless.h按 core 数切分 tokenperCoreElements再按 token 组循环每组用Duplicate将 buffer 清零后逐 token 处理实现多 core 并行。expanded_row_idx以 stride N 方式访问使得同一 token 的 K 个索引分散在不同 core 的独立累加路径上避免写冲突。Active 与 DropPad 场景的差别仅在跳过无效索引的判断上Active 场景xRow activeNum时跳过moe_init_routing_v2_grad_with_activate.hDropPad 场景xRow -1时跳过moe_init_routing_v2_grad_with_pos_drop_and_pad_zero.h且被跳过位置会用Duplicate将 buffer 置 0 后再继续其余累加流程完全一致。aclnn 接口层的空输出优化在 moe_init_routing_v2_grad.cpp 的l0op::MoeInitRoutingV2Grad中输出张量grad_x若为空grad_x-IsEmpty()会直接跳过算子逻辑返回不触发 AICore 任务下发这对应了文档示例中输出为空时不走算子的边界处理。七、约束说明使用该接口时需要满足以下约束与正向调用严格对齐该接口支持训练场景下使用。该接口支持单算子模式和图模式。grad_expanded_x在 Dropless 场景下必须是 2 维张量在 DropPad 场景下必须是 3 维张量。expanded_row_idx必须是 1 维张量其长度必须能被top_k整除。drop_pad_mode仅支持取值 0 或 1。top_k必须大于 0。active_num必须大于等于 0。Dropless 场景且active_num0时grad_expanded_x的第 0 维大小必须与expanded_row_idx的长度一致Dropless 场景且active_num0时grad_expanded_x的第 0 维大小必须等于active_num。grad_expanded_x的最后一维DropPad 场景下为第三维必须与输出grad_x的第二维一致。该算子为 moe_init_routing 的反向算子expanded_row_idx、top_k、drop_pad_mode、active_num需与正向调用保持一致。该反向算子仅支持非量化场景正向quant_mode-1不支持aclnnMoeInitRoutingV4特有特性scale、offset、topk_weight、x_dtype、row_idx_type非 0。当正向使用了这些特性时调用自动反向会抛出NotImplementedError。上述约束在 host 侧 moe_init_routing_v2_grad_infershape.cpp 的CheckParm中同样有图模式下的校验drop_pad_mode必须为 0 或 1、top_k必须大于 0、active_num必须大于等于 0、grad_expanded_x维度数必须匹配场景要求、expanded_row_idx必须为 1 维。八、确定性计算该接口默认支持确定性计算。对同一输入多次调用结果可复现适用于对数值可复现性有要求的训练调试场景。九、配套接口自动反向与手动调用该算子为 moe_init_routing 的反向算子。当正向moe_init_routing中x.requires_grad为 True且未使用aclnnMoeInitRoutingV4特有特性量化、scale、offset、topk_weight等时loss.backward()会自动触发本算子无需手动调用。仅在需要显式控制梯度的场景下保留手动调用路径。自动反向的约束细节在正向文档中有明确说明自动反向仅在正向退化为aclnnMoeInitRoutingV2场景时支持即scale不传入、offset不传入、topk_weight不传入、quant_mode-1、row_idx_type0、x_dtype为 None、drop_pad_mode为 0 或 1active_expert_range不影响反向。此外自动反向仅对x求梯度expert_idx为整数索引张量无梯度expanded_row_idx及其他整数/统计输出同样无梯度。十、调用示例单算子模式调用以下示例在 NPU 上完成一次正向路由 手动显式反向的完整流程来源torchapi_moe_init_routing_grad.mdimport torch import torch_npu from cann_ops_transformer.ops import moe_init_routing, moe_init_routing_grad torch_npu.npu.set_device(0) n 4 h 8 k 2 expert_num 8 drop_pad_mode 0 expert_tokens_num_type 1 expert_tokens_num_flag True quant_mode -1 active_expert_range [0, 4] row_idx_type 0 x torch.randn((n, h), dtypetorch.float32, devicenpu) expert_idx torch.randint(0, expert_num, (n, k), dtypetorch.int32, devicenpu) # 正向 expanded_x, expanded_row_idx, expert_tokens_count, expanded_scale, expanded_topk_weight \ moe_init_routing(x, expert_idx, active_num-1, expert_capacity-1, expert_numexpert_num, drop_pad_modedrop_pad_mode, expert_tokens_num_typeexpert_tokens_num_type, expert_tokens_num_flagexpert_tokens_num_flag, quant_modequant_mode, active_expert_rangeactive_expert_range, row_idx_typerow_idx_type) # 手动调用反向 grad_expanded_x torch.randn_like(expanded_x) grad_x moe_init_routing_grad(grad_expanded_x, expanded_row_idx, top_kk, drop_pad_modedrop_pad_mode, active_num0) print(grad_x.shape) # torch.Size([4, 8])注意示例中正向使用了row_idx_type0gather 索引且quant_mode-1非量化符合自动反向约束但此处选择了手动调用反向以演示显式控制梯度的路径。expanded_scale在非量化且未传scale时返回的数值无意义属正常现象。自动反向调用import torch import torch_npu from cann_ops_transformer.ops import moe_init_routing torch_npu.npu.set_device(0) n 4 h 8 k 2 expert_num 8 x torch.randn((n, h), dtypetorch.float32, devicenpu, requires_gradTrue) expert_idx torch.randint(0, expert_num, (n, k), dtypetorch.int32, devicenpu) # 正向x.requires_gradTrue 时自动启用 autograd expanded_x, expanded_row_idx, expert_tokens_count, expanded_scale, expanded_topk_weight \ moe_init_routing(x, expert_idx, active_num-1, expert_capacity-1, expert_numexpert_num, drop_pad_mode0, expert_tokens_num_type1, expert_tokens_num_flagTrue, quant_mode-1, active_expert_range[0, 4], row_idx_type0) # 反向自动触发 moe_init_routing_grad loss expanded_x.sum() loss.backward() print(x.grad.shape) # torch.Size([4, 8])自动反向通过 torch 扩展层完成moe_init_routing_grad.py注册了Meta实现用于 shape 推导PrivateUse1分发实现最终调用编译出的op_module.moe_init_routing_gradmoe_init_routing_grad.py该 C 函数完成参数校验与输出张量分配后经由ACLNN_CMD(aclnnMoeInitRoutingV2Grad, ...)下发 NPU 计算moe_init_routing_grad.cpp。图模式支持该接口支持 TorchAir 图模式调用。图模式下的调用方式与正向算子一致——将算子调用放入torch.compile(model, backendnpu_backend)的模型中执行正向 moe_init_routing 文档提供了完整的图模式示例可供参考反向算子随正向图一并编译执行。更多参考底层 aclnn 接口文档aclnnMoeInitRoutingV2Grad.mdC 调用示例test_aclnn_moe_init_routing_v2_grad.cpp模块级说明moe_init_routing_v2_grad/README.md单元测试host/op_apitest_aclnn_moe_init_routing_v2_grad.cpp、test_moe_init_routing_v2_grad_infershape.cpp单元测试kerneltest_moe_init_routing_v2_grad.cpp十一、总结moe_init_routing_grad是 CANN ops-transformer 中 MoE 路由反向链路的关键算子本质是一个基于expanded_row_idx行映射的分组归约reduction操作每个原始 token 的梯度是其 K 个专家副本梯度的累加和。在使用时需要抓住三个要点场景对齐drop_pad_mode、active_num必须与正向moe_init_routing调用保持一致并据此选择正确的grad_expanded_x维度2D 或 3D约束匹配仅非量化、不使用 V4 特有特性时支持自动反向否则需手动调用并处理NotImplementedError性能路径NPU kernel 通过多 core 切分、on-chip 累加 buffer 与二分累加策略高效完成 K 路梯度累加Active/DropPad 场景仅以索引有效性判断xRow active_num、xRow -1区分跳过逻辑整体计算路径统一且可确定性复现。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考