torch_npu.npu_moe_compute_expert_tokens

功能描述

• 算子功能：MoE计算中，通过二分查找的方式查找每个专家处理的最后一行的位置。
• 计算公式：

  

接口原型

npu_moe_compute_expert_tokens(Tensor sorted_expert_for_source_row, int num_expert) -> Tensor

参数说明

• sorted_expert_for_source_row：必选参数，经过专家处理过的结果，要求是一个1D的Tensor，数据类型支持INT32，数据格式要求为ND。shape大小需要小于2147483647。

• num_expert：必选参数，总专家数。

输出说明

expertTokens：Device侧的aclTensor，公式中的输出，要求的是一个1D的Tensor，数据类型与sorted_expert_for_source_row保持一致。

约束说明

该融合算子仅在推理场景使用。

支持的型号

Atlas A2训练系列产品/Atlas 800I A2推理产品

调用示例

import torch
import torch_npu
sorted_experts = torch.tensor([3,3,4,5,6,7], dtype=torch.int32)
num_experts = 5
output = torch_npu.npu_moe_compute_expert_tokens(sorted_experts.npu(), num_experts)