主要功能为将token和专家的映射关系反转为专家与token的映射关系。算子输入为MoE模型每个token选中专家的索引，算子输出为MoE模型每个专家对应的token的索引。
该算子支持TP和EP场景，当参数“deviceExpert”为空时，为TP场景，否则为EP场景。
1 2 3 4 5 6 7 | 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 |
|
是 |
每个token选中的专家数。 |
cumSumNum |
int32_t |
0 |
[0, 1024] |
是 |
专家总数。 |
cumSumInt64 |
bool |
false |
true/false |
否 |
输出的cumSum的类型是否为int64。 当为false时，输出的cumSum类型为int32。 |
deviceExpert |
std::vector< int32_t > |
- |
|
否 |
当前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。 |
||||
参数 |
维度 |
数据类型 |
格式 |
描述 |
|---|---|---|---|---|
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。 |
||||
非