CANN/cann-bench MoE门控算子

MoeGatingTopKSoftmax 算子 API 描述

【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力，涵盖算子生成、算子优化等领域，支撑模型选型、训练效果评估，统一量化评估标准，识别Agent能力短板，构建CANN领域评测平台，推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench

1. 算子简介

MoE 门控网络中 Softmax 和 TopK 的融合算子，先对输入执行 Softmax 归一化，再取 TopK 个最大值，同时输出专家索引和行位置索引。

主要应用场景：

• Mixture of Experts (MoE) 模型中的门控路由选择
• 稀疏激活 Transformer 中的专家选择策略
• MoE 层中 token 到专家的路由分配

算子特征：

• 难度等级：L3（LayoutTransform）
• 支持 2D 或 3D 输入，输出 TopK 结果、专家索引、行索引三个输出

2. 算子定义

数学公式

$$ softmaxOut = softmax(x, axis=-1) $$

$$ yOut, expertIdxOut = topK(softmaxOut, k) $$

$$ rowIdxRange = arange(expertIdxOut.shape[0] \times expertIdxOut.shape[1]) $$

$$ rowIdxOut = rowIdxRange.reshape([expertIdxOut.shape[1], expertIdxOut.shape[0]]).transpose(1, 0) $$

处理流程

1. 对输入 $x$ 沿最后一维执行 Softmax：$\text{softmaxOut}[i] = \frac{e^{x_i}}{\sum_j e^{x_j}}$
2. 从 Softmax 结果中选取前 $k$ 个最大值作为输出 $yOut$，对应索引为 $expertIdxOut$
3. 计算 $rowIdxOut$，表示展平后的全局位置索引（用于指示每个输出位置在展平数组中的位置）
4. 如果finished中对应行为 True，则 $expertIdxOut$ 中填入 $num_expert$（即 $x$ 的最后一维大小）

3. 接口规范

算子原型

cann_bench.moe_gating_top_k_softmax(Tensor x, Tensor? finished=None, int k=1) -> (Tensor y, Tensor expert_idx, Tensor row_idx)

输入参数说明

输出

数据类型

规则与约束

• 输入x支持 2D 或 3D 张量
• k必须为正整数，且满足 0 < k <= x最后一维大小，k <= 1024
• finished为可选参数，若提供则 shape 必须为 x_shape[:-1]，dtype 必须为 bool
• Softmax 和 TopK 均沿最后一维（dim=-1）计算
• 当finished中某行为 True 时，该行对应的expert_idx值会被设为num_expert

4. 精度要求

采用生态算子精度标准进行验证。

误差指标：

1. 平均相对误差（MERE）：采样点中相对误差平均值

   $$ \text{MERE} = \text{avg}(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)+\text{1e-7}}) $$

2. 最大相对误差（MARE）：采样点中相对误差最大值

   $$ \text{MARE} = \max(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)+\text{1e-7}}) $$

通过标准：

当平均相对误差 MERE < Threshold，最大相对误差 MARE < 10 * Threshold 时判定为通过。

5. 标准 Golden 代码

import torch import numpy as np """ MoeGatingTopKSoftmax算子Torch Golden参考实现 公式: softmaxOut = softmax(x, axis=-1) yOut, expertIdxOut = topK(softmaxOut, k) rowIdxRange = arange(expertIdxOut.shape[0] * expertIdxOut.shape[1]) rowIdxOut = rowIdxRange.reshape([expertIdxOut.shape[1], expertIdxOut.shape[0]]).transpose(1, 0) 注意: row_idx是展平后的全局位置索引，而非行号 """ def moe_gating_top_k_softmax( x: torch.Tensor, finished: torch.Tensor = None, k: int = 1 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ MoE门控网络中Softmax和TopK的融合 Args: x: 输入张量，shape (..., E) finished: 可选，标记哪些行参与计算，bool类型，shape x_shape[:-1] k: topK数量 Returns: y: topK值，shape (..., k) expert_idx: topK索引（专家序号），shape (..., k)，int32 row_idx: 展平后的全局位置索引，shape (..., k)，int32 """ # Softmax沿最后一维 softmax_out = torch.nn.functional.softmax(x, dim=-1) # TopK沿最后一维 values, indices = torch.topk(softmax_out, k, dim=-1) # 计算row_idx，严格对标op-plugin测试代码 # 公式: rowIdxRange = arange(shape[0] * shape[1]) # rowIdxOut = rowIdxRange.reshape([shape[1], shape[0]]).transpose(1, 0) # 注意: row_idx是展平后的全局位置索引 output_shape = indices.shape if len(output_shape) == 2: # 2D: (N, k) # row_idx = arange(N*k).reshape(k, N).transpose(1, 0) -> (N, k) row_idx_range = torch.arange(output_shape[0] * output_shape[1], dtype=torch.int32) row_idx = row_idx_range.reshape(output_shape[1], output_shape[0]).transpose(0, 1) else: # 3D: (B, N, k) # 先把(B, N)看作整体，计算展平后的索引 # row_idx_range = arange(B*N*k) # reshape成(k, B*N)，transpose成(B*N, k) # 再reshape成(B, N, k) row_idx_range = torch.arange(output_shape[0] * output_shape[1] * output_shape[2], dtype=torch.int32) row_idx = row_idx_range.reshape(output_shape[2], output_shape[0] * output_shape[1]).transpose(0, 1) row_idx = row_idx.reshape(output_shape) # 处理finished参数 if finished is not None: num_expert = x.shape[-1] finished_expanded = finished.unsqueeze(-1).expand_as(indices) indices = torch.where(finished_expanded, num_expert, indices) return values, indices.to(torch.int32), row_idx.to(torch.int32)

6. 额外信息

算子调用示例

import torch import cann_bench # 2D 输入 x = torch.randn(1024, 64, dtype=torch.float16, device="npu") y, expert_idx, row_idx = cann_bench.moe_gating_top_k_softmax(x, k=8) # 带 finished 参数 finished = torch.randint(0, 2, (1024,), dtype=torch.bool, device="npu") y, expert_idx, row_idx = cann_bench.moe_gating_top_k_softmax(x, finished=finished, k=8) # 3D 输入 x_3d = torch.randn(4, 1024, 64, dtype=torch.float32, device="npu") y, expert_idx, row_idx = cann_bench.moe_gating_top_k_softmax(x_3d, k=4)

【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力，涵盖算子生成、算子优化等领域，支撑模型选型、训练效果评估，统一量化评估标准，识别Agent能力短板，构建CANN领域评测平台，推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench