GroupTopkOperation（代码开放）

产品支持情况

功能说明

GroupTopk算子超参数。将输入tensor0中维度1（输入tensor0有2个维度：维度0和维度1）数据分groupNum个组，每组取最大值，然后选出每组最大值中前k个，最后将非前k个组的数据全部置零。

定义

 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};
};

参数列表

输入

输出

约束说明

1≤expert_num≤1024，expert_num≥groupNum≥k≥1，expert_num能被groupNum整除。

基础功能

• 功能概述

  将输入的各专家分数分组，每组内选取最大值，根据每组最大值大小选取topk组，其余组置零。

• 计算公式

  对于每个token，首先将token的数据均匀分groupNum组，

  1. 每组取最大值；
  2. 对每组最大值降序排序；
  3. 获取前k个组的index，然后将其余的index对应的组的数据置零。
• 计算图

  groupMultiFlag取0时，计算过程下图所示：

  图1 功能示例
  
• 参数列表

  见参数列表，需要满足groupMultiFlag置为UNDEFINED。

• 输入

  参数	维度	数据类型	格式	描述
  token	[num_tokens, expert_num]	float16/bf16	ND	二维tensor。 维度0为token数。 维度1为专家总数。
  idxArr	[1024]	int32	ND	一维tensor，用于辅助计算，固定长度1024，[0,1,2,...,1023]的等差序列。

• 输出

  参数	维度	数据类型	格式	描述
  output	[num_tokens, expert_num]	float16/bf16	ND	是对输入token原地写的输出。数据类型与输入token保持一致。

• 使用示例

  1 2 3 4

  // 参数构造 atb::infer::GroupTopkParam param; param.groupNum = 8; param.k = 1;

组内再排序和选取前n个值功能

• 功能概述

  将输入的各专家分数分组，每组内降序排列选取前n个值求和，根据和的大小选取topk组，其余组置零。

• 计算公式

  对于每个token，首先将token的数据均匀分groupNum组，

  1. 每组降序排列，取前n个值；
  2. 对这n个值求和；
  3. 各个组根据求得的和大小降序排序；
  4. 获取前k个组的index，然后将其余的index对应的组的数据置零。
  

  n取1时对应基础功能场景。

• 参数列表

  见参数列表，需要满足groupMultiFlag置为SUM_MULTI_MAX。

• 输入

  参数	维度	数据类型	格式	描述
  token	[num_tokens, expert_num]	float16/bf16	ND	二维tensor。 维度0为token数。 维度1为专家总数。
  idxArr	[1024]	int32	ND	一维tensor，用于辅助计算，固定长度1024，[0,1,2,...,1023]的等差序列。

• 输出

  参数	维度	数据类型	格式	描述
  output	[num_tokens, expert_num]	float16/bf16	ND	是对输入tensor0原地写的输出。数据类型与输入tensor0保持一致。

• 使用示例

  1 2 3 4 5 6

  // 参数构造 atb::infer::GroupTopkParam param; param.groupNum = 8; param.k = 1; param.groupMultiFlag = 1; param.n = 2;