GatingOperation（代码开放）

产品支持情况

硬件型号	是否支持
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。

输出

参数	维度	数据类型	格式	描述
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。

约束说明

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

父主题： atb/infer_op_params.h