GroupTopkOperation
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
struct GroupTopkParam {
int32_t groupNum = 1;
int32_t k = 0;
uint8_t rsv[16] = {0};
};
Parameters
Member |
Type |
Default Value |
Value Range |
Description |
Required (Yes/No) |
|---|---|---|---|---|---|
groupNum |
int32_t |
1 |
[1. Total number of experts] |
Number of token groups. Note: The total number of experts is the value of inTensor0.dims[1]. |
No |
k |
int32_t |
0 |
[1, groupNum] |
Number of top K experts. |
Yes |
rsv[16] |
uint8_t |
{0} |
- |
Reserved |
- |
Input
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
token |
[d0,d1] |
float16/bf16 |
ND |
Input tensor 0, 2D tensor. Dimension 0 indicates the number of tokens, and dimension 1 indicates the total number of experts. |
idxArr |
[1024] |
int32 |
ND |
Input tensor 1, 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 |
|---|---|---|---|---|
output0 |
[d0,d1] |
float16/bf16 |
ND |
Output tensor 0. There is only one output tensor, which is the output of the input tensor 0 written in place. Has the same data type as input tensor 1. |
Restrictions
- Ensure that inTensor0Desc.shape.dims[1] can be exactly divided by groupNum .
- 1 ≤ d1 ≤ 1024, d1 ≥ groupNum ≥ k ≥ 1, where d1 can be exactly divided by groupNum.
- Only the
Atlas A2 inference products is supported.