---
title: GatingOperation
description: "| 硬件型号 | 是否支持 |"
url: https://www.hiascend.com/document/detail/zh/canncommercial/latest/API/ascendtb/ascendtb_01_0105.html
sourcePath: /source/zh/canncommercial/900/API/ascendtb/ascendtb_01_0105.html
indexId: 3677ad5fe9cec93e821169b7615736b86c725b613e8fa42bae88edfeb1531c7f64
---
# GatingOperation

#### 产品支持情况

| 硬件型号 | 是否支持 |
| --- | --- |
| Atlas 350 加速卡 | x |
| Atlas A3 推理系列产品 / Atlas A3 训练系列产品 | √ |
| Atlas A2 训练系列产品 / Atlas A2 推理系列产品 | √ |
| Atlas 训练系列产品 | x |
| Atlas 推理系列产品 | √ |
| Atlas 200I/500 A2 推理产品 | x |


#### 功能说明

主要功能为将token和专家的映射关系反转为专家与token的映射关系。算子输入为MoE模型每个token选中专家的索引，算子输出为MoE模型每个专家对应的token的索引。

该算子支持TP和EP场景，当参数“deviceExpert”为空时，为TP场景，否则为EP场景。


#### 定义

```
struct GatingParam {
    int32_t topkExpertNum = 1;
    int32_t cumSumNum = 0;
    bool cumSumInt64 = false;
    std::vector<int32_t> deviceExpert;
    uint8_t rsv[16] = {0};
};
```


#### 参数列表

| 成员名称 | 类型 | 默认值 | 取值范围 | 是否必选 | 描述 |
| --- | --- | --- | --- | --- | --- |
| topkExpertNum | int32\_t | 1 | 当cumSumNum为0时，取值为1 当cumSumNum不为0时，为(0, cumSumNum] | 是 | 每个token选中的专家数。 |
| cumSumNum | int32\_t | 0 | [0, 1024] | 是 | 专家总数。 |
| cumSumInt64 | bool | false | true/false | 否 | 输出的cumSum的类型是否为int64。 当为false时，输出的cumSum类型为int32。 |
| deviceExpert | std::vector< int32\_t > | \- | 当cumSumNum为0时，列表必须为空 当cumSumNum不为0时，列表中各个元素取值为[0, cumSumNum) | 否 | 当前device上的专家索引列表。 列表中各个元素取值范围为[0, cumSumNum)，且其中元素值不可重复。 当cumSumNum为0时，必须为空。 |
| rsv[16] | uint8\_t | {0} | [0] | 否 | 预留参数。 |


#### 输入

| 参数 | 维度 | 数据类型 | 格式 | 描述 |
| --- | --- | --- | --- | --- |
| topk | [tokenNum \* topkExpertNum] | int32 | ND | 输入tensor。每个token选中的专家的index。值域为[0, cumSumNum \- 1]；当cumSumNum为0时，值域为[0, 11300)。 |
| idxArr | [tokenNum \* topkExpertNum] | int32 | ND | 输入tensor。每个token原始的index，具体的值为[0,1,2,3,...]。 |
| tokenNum表示token个数，tokenNum = batch \* seqlen。 | tokenNum表示token个数，tokenNum = batch \* seqlen。 | tokenNum表示token个数，tokenNum = batch \* seqlen。 | tokenNum表示token个数，tokenNum = batch \* seqlen。 | tokenNum表示token个数，tokenNum = batch \* seqlen。 |


#### 输出

| 参数 | 维度 | 数据类型 | 格式 | 描述 |
| --- | --- | --- | --- | --- |
| tokenIndex | [tokenNum \* topkExpertNum] | int32 | ND | 输出tensor。token重排以后原始的索引值。 |
| cumSum | [expertNum] | int32/int64 | ND | 输出tensor。每个专家被选中的次数。当cumSumNum为0时，expertNum值为1；当deviceExpert不为空时，expertNum值为deviceExpert的元素个数，否则，expertNum值为cumSumNum，当cumSumInt64为True时，输出为int64类型。 |
| originalIndex | [tokenNum \* topkExpertNum] | int32 | ND | 输出tensor。token重排以后token的索引值。 |
| validIndex | [1] | int32 | ND | 输出tensor。当deviceExpert不为空时输出。 |
| tokenNum表示token个数，tokenNum = batch \* seqlen。 | tokenNum表示token个数，tokenNum = batch \* seqlen。 | tokenNum表示token个数，tokenNum = batch \* seqlen。 | tokenNum表示token个数，tokenNum = batch \* seqlen。 | tokenNum表示token个数，tokenNum = batch \* seqlen。 |


#### 约束说明

非 Atlas A2 训练系列产品 / Atlas A2 推理系列产品 和 Atlas A3 推理系列产品 / Atlas A3 训练系列产品 仅支持TP场景。