GroupedMatmulWithRoutingOperation

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	x
Atlas 200I/500 A2 inference products	x

Description

Multiplies the weights of top K experts by token activations in MOE. This operator does not support deterministic computing.

Application Scenarios

UP: The token value is allocated to the top K experts based on the weight of each expert.

DOWN: The computation results of top K experts are re-integrated based on the weight of each expert.

Computation Process

UP
Formula:

$\text{[math]}$

AcTensor is the activated token value, Expertindex is the expert-token index, and ExpertWeight is the expert weight. The process is as follows:
1. Obtain the activated token value of the current expert from AcTensor based on of the current expert.
2. Obtain the weight of the current expert and perform BatchMatmul with the activated token value of the current expert.
DOWN
Formula:

$\text{[math]}$

AcTensor is the activated token value, Expertindex is the expert-token index, and ExpertWeight is the expert weight. The process is as follows:
1. Obtain the activated token value of the current expert from AcTensor based on of the current expert.
2. Obtain the weight of the current expert and perform BatchMatmul with the activated token value of the current expert.
3. Add the results of the current token on all experts.

Definition

      
           struct GroupedMatmulWithRoutingParam {
    enum GroupedMatmulType : int {
        GROUPED_MATMUL_UP = 0, 
        GROUPED_MATMUL_DOWN  
    };
    bool transposeB = true;
    int32_t topK = 0;
    GroupedMatmulType groupedMatmulType = GROUPED_MATMUL_UP;
    aclDataType outDataType = ACL_DT_UNDEFINED;
    uint8_t rsv[16] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Required (Yes/No)	Description
groupedMatmulType	GroupedMatmulType	GROUP_MATMUL_UP	GROUP_MATMUL_UP GROUP_MATMUL_DOWN	Yes	GroupedMatmulType type: GROUP_MATMUL_UP: UP. GROUP_MATMUL_DOWN: DOWN.
transposeB	bool	true	false/true	No	Whether to transpose matrix B.
topK	int32_t	0	Greater than 0	Yes	Number of top K experts.
outDataType	aclDataType	ACL_DT_UNDEFINED	ACL_DT_UNDEFINED ACL_FLOAT16 ACL_BF16	Yes in the dequantization scenario	In non-quantization scenarios, the data type must be the same as the input data type. In the dequantization scenario, the output data type must be set. The output types can be ACL_FLOAT16 or ACL_BF16.
rsv[16]	uint8_t	{0}	[0]	-	Reserved

Input and output

Non-quantization scenario

Input

Parameter	Dimension	Data Type	Format	Description
AcTensor	[num_tokens, hidden_size_in]	float16/bf16	ND	Activation value
ExpertWeight	[num_experts, hidden_size_in, hidden_size_out]	float16/bf16	ND	Expert weight. If transposeB is set to true, the locations of hidden_size_in and hidden_size_out are exchanged.
ExpertCount	[num_experts]	int32	ND	Number of expert tokens.
Expertindex	[num_tokens*topK]	int32	ND	Expert token index.

Output

Parameter	Dimension	Data Type	Format	Description
Result	[num_tokens*Topk, hidden_size_out]	float16/bf16	ND	Output

Quantization scenario

Input

Parameter	Dimension	Data Type	Format	Description
AcTensor	[num_tokens, hidden_size_in]	int8	ND	Activation value
ExpertWeight	[num_experts, hidden_size_in, hidden_size_out]	int8	ND/NZ	Expert weight If transposeB is set to true, the locations of hidden_size_in and hidden_size_out are exchanged.
ExpertCount	[num_experts]	int32	ND	Number of expert tokens.
Expertindex	[num_tokens*topK]	int32	ND	Expert token index.
nscale	[num_experts, hidden_size_out]	float	ND	Dequantization coefficient in the ExpertWeight direction
mscale	[num_tokens]	float	ND	Dequantization coefficient in the AcTensor direction

Output

Parameter	Dimension	Data Type	Format	Description
Result	[num_tokens*Topk, hidden_size_out]	float16/bf16	ND	Output

DOWN

Non-quantization scenario

Input

Parameter	Dimension	Data Type	Format	Description
AcTensor	[num_tokens*topK, hidden_size_in]	float16/bf16	ND	Activation value
ExpertWeight	[num_experts, hidden_size_in,hidden_size_out]	float16/bf16	ND	Expert weight If transposeB is set to true, the locations of hidden_size_in and hidden_size_out are exchanged.
ExpertCount	[num_experts]	int32	ND	Number of expert tokens.
Expertindex	[num_tokens*topK]	int32	ND	Expert token index.

Output

Parameter	Dimension	Data Type	Format	Description
Result	[num_tokens, hidden_size_out]	float16/bf16	ND	Output

Quantization scenario

Input

Parameter	Dimension	Data Type	Format	Description
AcTensor	[num_tokens*topK, hidden_size_in]	int8	ND	Activation value
ExpertWeight	[num_experts, hidden_size_in,hidden_size_out]	int8	ND/NZ	Expert weight. If transposeB is set to true, the locations of hidden_size_in and hidden_size_out are exchanged.
ExpertCount	[num_experts]	int32	ND	Number of expert tokens.
Expertindex	[num_tokens*topK]	int32	ND	Expert token index.
nscale	[num_experts, hidden_size_out]	float	ND	Dequantization coefficient in the ExpertWeight direction
mscale	[num_tokens*topK]	float	ND	Dequantization coefficient in the AcTensor direction

Output

Parameter	Dimension	Data Type	Format	Description
Result	[num_tokens, hidden_size_out]	float16/bf16	ND	Output

Constraints

If the output is quantized, the input must be of the int8 type. If the output is not quantized, the input must be of the bf16 or float16 type.
The value of top K cannot exceed the number of experts.
Value range of nTokens: [128, 512].
Value range of nExperts: [128, 256].
Value range of topK: [2, 10].
In UP mode, the value of hiddenSizeIn is in the range of [32, 5120] and must be an integral multiple of 32. In DOWN mode, the value of hiddenSizeIn is in the range of [32, 256] and must be an integral multiple of 32.
In UP mode, the value of hiddenSizeOut is in the range of [32, 256] and must be an integral multiple of 32. In DOWN mode, the value of hiddenSizeOut is in the range of [32, 5120] and must be an integral multiple of 32.

Parent topic: atb/infer_op_params.h