GroupTopkOperation

Applicable Products

Description

(Hyperparameter of the GroupTopk operator) Divides the data of dimension 1 in input tensor 0 (which has two dimensions: dimension 0 and dimension 1) into groupNum groups, selects the maximum value of each group and the first k values in each group, and sets all data of other groups to 0.

Definition

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

struct GroupTopkParam {
    int32_t groupNum = 1;
    int32_t k = 0;
    enum GroupMultiFlag : uint16_t {
        UNDEFINED = 0, 
        SUM_MULTI_MAX  
    };
    uint16_t n = 1; 
    uint8_t rsv[12] = {0};
};

Parameters

Input

Output

Constraints

1 ≤ expert_num ≤ 1024 and expert_num ≥ groupNum ≥ k ≥ 1, where expert_num can be exactly divided by groupNum.

Basic Functions

• Function

  Groups the input expert scores, selects the maximum value in each group, selects the top k groups based on the maximum value in each group, and sets the scores of other groups to 0.

• Formula

  For each token, the token data is evenly divided into groupNum groups.

  1. Obtain the maximum values of each group.
  2. Sort the maximum values of each group in descending order.
  3. Obtain the indexes of the first k groups and set the data of the groups corresponding to the other indexes to 0.
• Computational Graph

  When groupMultiFlag is set to 0, the calculation process is as follows:

  Figure 1 Function example
  
• Parameters

  For details, see Parameters. The value of groupMultiFlag must be UNDEFINED.

• Input

  Parameter	Dimension	Data Type	Format	Description
  token	[num_tokens, expert_num]	float16/bf16	ND	2D tensor. Dimension 0: number of tokens. Dimension 1: total number of experts.
  idxArr	[1024]	int32	ND	1D tensor, used for auxiliary computation. It is an arithmetic sequence [0, 1, 2, ...,1023] with a fixed length of 1024.

• Output

  Parameter	Dimension	Data Type	Format	Description
  output	[num_tokens, expert_num]	float16/bf16	ND	Output of the input token written in place. The data type is the same as that of the input token.

• Example

  1 2 3 4

  // Parameter construction atb::infer::GroupTopkParam param; param.groupNum = 8; param.k = 1;

Re-sorting and Selecting the First N Values in a Group

• Function

  Groups the input expert scores, sorts the scores in each group in descending order, selects the first n values, sums the values, selects the top k groups based on the sum, and sets the scores of other groups to 0.

• Formula

  For each token, the token data is evenly divided into groupNum groups.

  1. Sort the values in each group in descending order and obtain the first n values.
  2. Sum the n values.
  3. Sort the groups in descending order based on the obtained sum.
  4. Obtain the indexes of the first k groups and set the data of the groups corresponding to the other indexes to 0.
  

  When n is set to 1, the Basic Functions scenario applies.

• Parameters

  For details, see Parameters. groupMultiFlag must be set to SUM_MULTI_MAX.

• Input

  Parameter	Dimension	Data Type	Format	Description
  token	[num_tokens, expert_num]	float16/bf16	ND	2D tensor. Dimension 0: number of tokens. Dimension 1: total number of experts.
  idxArr	[1024]	int32	ND	1D tensor, used for auxiliary computation. It is an arithmetic sequence [0, 1, 2, ...,1023] with a fixed length of 1024.

• Output

  Parameter	Dimension	Data Type	Format	Description
  output	[num_tokens, expert_num]	float16/bf16	ND	Output of input tensor 0 written in place. The data type is the same as that of input tensor 0.

• Example

  1 2 3 4 5 6

  // Parameter construction atb::infer::GroupTopkParam param; param.groupNum = 8; param.k = 1; param.groupMultiFlag = 1; param.n = 2;