主要功能为将token和专家的映射关系反转为专家与token的映射关系。算子输入为MoE模型每个token选中专家的索引，算子输出为MoE模型每个专家对应的token的索引。
该算子支持TP和EP场景，当参数“deviceExpert”为空时，为TP场景，否则为EP场景。
struct GatingParam {
int32_t topkExpertNum = 0;
int32_t cumSumNum = 0;
bool cumSumInt64 = false;
std::vector<int32_t> deviceExpert;
};
成员名称 |
类型 |
默认值 |
描述 |
|---|---|---|---|
topkExpertNum |
int32_t |
0 |
每个token选中的专家数。 默认值为0。 当cumSumNum为0时，取值为1；否则，取值范围为(0, cumSumNum]。 |
cumSumNum |
int32_t |
0 |
专家总数。 默认值为0。 取值范围为[0, 200]。 |
cumSumInt64 |
bool |
false |
输出的cumSum的类型是否为int64。 默认值为false。 当为false时，输出的cumSum类型为int32。 |
deviceExpert |
std::vector< int32_t > |
- |
当前device上的专家索引列表。 默认为空。 列表中各个元素取值范围为[0, cumSumNum)，且其中元素值不可重复。 当cumSumNum为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。 |
||||
参数 |
维度 |
数据类型 |
格式 |
描述 |
|---|---|---|---|---|
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。 |
||||
非