GatingOperation

Applicable Products

Hardware Model	Supported
Atlas 350 accelerator cards	x
Atlas A3 inference products/Atlas A3 training products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas training products	x
Atlas inference products	√
Atlas 200I/500 A2 inference products	x

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 = 1;
    int32_t cumSumNum = 0;
    bool cumSumInt64 = false;
    std::vector<int32_t> deviceExpert;
    uint8_t rsv[16] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Mandatory	Description
topkExpertNum	int32_t	1	When cumSumNum is 0, the value is 1. When cumSumNum is not 0, the value is in the range of (0, cumSumNum].	Yes	Number of experts selected for each token.
cumSumNum	int32_t	0	[0, 1024]	Yes	Total number of experts.
cumSumInt64	bool	false	true/false	No	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 >	-	When cumSumNum is set to 0, the list must be empty. When cumSumNum is not set to 0, the value range of each element in the list is [0, cumSumNum).	No	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}	[0]	No	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.

Constraints

Non-Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products support only TP.

Parent topic: atb/infer_op_params.h