GatingOperation
Description
Reverses the mapping between tokens and experts to the mapping between experts and tokens. The operator input is the index of each selected expert in the MoE model, and the operator output is the index of the token corresponding to each expert in the MoE model.
This operator supports the TP and EP scenarios. If deviceExpert is empty, the TP scenario is used. Otherwise, the EP scenario is used.
Definition
struct GatingParam {
int32_t topkExpertNum = 0;
int32_t cumSumNum = 0;
bool cumSumInt64 = false;
std::vector<int32_t> deviceExpert;
uint8_t rsv[16] = {0};
};
Parameters
Member |
Type |
Default Value |
Description |
|---|---|---|---|
topkExpertNum |
int32_t |
0 |
Number of experts selected for each token. When cumSumNum is 0, the value is 1. Otherwise, the value range is (0, cumSumNum]. |
cumSumNum |
int32_t |
0 |
Total number of experts. The value range is [0, 200]. |
cumSumInt64 |
bool |
false |
Whether the output cumSum is of the int64 type. When the value is false, the output cumSum is of the int32 type. |
deviceExpert |
std::vector< int32_t > |
- |
List of expert indexes on the current device. The value range of each element in the list is [0, cumSumNum), and the element values must be unique. When cumSumNum is set to 0, the value of this parameter must be null. |
rsv[16] |
uint8_t |
{0} |
Reserved |
Input
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
topk |
[tokenNum * topkExpertNum] |
int32 |
ND |
Input tensor Index of the expert selected for each token. The value range is [0, cumSumNum - 1]. When cumSumNum is 0, the value range is [0, 11300). |
idxArr |
[tokenNum * topkExpertNum] |
int32 |
ND |
Input tensor Original index of each token. The value range is [0,1,2,3, ...]. |
tokenNum indicates the number of tokens. tokenNum = batch * seqlen. |
||||
Output
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
tokenIndex |
[tokenNum * topkExpertNum] |
int32 |
ND |
Output tensor. Original index value after token rearrangement. |
cumSum |
[expertNum] |
int32/int64 |
ND |
Output tensor. Number of times that each expert is selected. When cumSumNum is 0, the value of expertNum is 1. When deviceExpert is not empty, the value of expertNum is the number of elements of deviceExpert. Otherwise, the value of expertNum is cumSumNum. When cumSumInt64 is True, the output is of the int64 type. |
originalIndex |
[tokenNum * topkExpertNum] |
int32 |
ND |
Output tensor. Token index after token rearrangement. |
validIndex |
[1] |
int32 |
ND |
Output tensor. Output when deviceExpert is not empty. |
tokenNum indicates the number of tokens. tokenNum = batch * seqlen. |
||||
Restrictions
Products except for the