add custom group_topk · Pull Request !54 · MindSpore/ms_custom_ops

1. add custom group_topk
---
<!--  Thanks for sending a pull request!  Here are some tips for you:

1) If this is your first time, please read our contributor guidelines: https://gitee.com/mindspore/mindspore/blob/master/CONTRIBUTING.md

2) If you want to contribute your code but don't know who will review and merge, please add label `mindspore-assistant` to the pull request, we will find and do it as soon as possible.
-->

**What type of PR is this?**

/kind feature

---

**What does this PR do / why do we need it**:

新增GroupTok算子适配，支持动/静态图，具体算子规格如下：
# group_topk 算子

## 描述

group_topk算子将输入token中最后一维度的数据分为group_num个组，每个组取最大值，然后选出每个组最大值中的前k个，最后将非前k个组的数据全部置零。

## 输入参数

| Name            | DType                       | Shape                        | Optional | Inplace | Format | Description                                         |
|-----------------|-----------------------------|------------------------------|----------|---------|--------|-----------------------------------------------------|
| token           | Tensor[float16/bfloat16]    | (num_tokens, expert_num)/(batch_size, seq_len, expert_num)     | No       | Yes     | ND     | 输入tensor0，若为二维Tensor，维度0为token数，维度1为专家总数。<br>若为三维Tensor，维度0乘维度1为token数，维度2为专家总数。 |
| idx_arr         | Tensor[int32]               | (idx_arr_dim)                | No       | No      | ND     | 输入tensor1，一维Tensor，用于辅助计算，固定长度的等差序列。推荐使用固定长度1024，[0,1,2,...,1023]的等差序列。 |
| group_num       | int                         | -                            | No       | -       | -      | 每个token分组数量，取值范围为[1, expert_num]。 |
| k               | int                         | -                            | No       | -       | -      | 选择top K专家数量，取值范围为[1, group_num]。 |
| k_inner         | int                         | -                            | Yes      | -       | -      | 计算每组得分时取最大值的数量，默认值为1。 |

## 输出参数

该接口无输出返回值，计算结果将原地更新至输入token中。

## 规格约束

- 1 ≤ expert_num ≤ 1024;
- expert_num ≥ group_num ≥ k ≥ 1;
- expert_num能够被group_num整除;
- 1 ≤ k_inner ≤ expert_num/group_num;
- idx_arr_dim ≥ expert_num;

## 使用示例

```python
import numpy as np
import mindspore as ms
import ms_custom_ops

# 创建输入张量
token = ms.Tensor(np.random.rand(512, 256), ms.float16)
idx_arr = ms.Tensor(np.arange(1024, dtype=np.int32))
group_num = 8
k = 8
k_inner = 2

# 调用算子
ms_custom_ops.group_topk(token=token, idx_arr=idx_arr, group_num=group_num,
                         k=k, k_inner=k_inner)
```
---

**Which issue(s) this PR fixes**:

Fixes #ID2NM1

---

**Test Plan and Test result：What scenarios were tested, and what were the verification results（Function, performance, reliability, etc.）**：

包含功能和精度验证，以及异常场景报错校验
![输入图片说明](https://foruda.gitee.com/images/1761137308967054662/4730797f_7725520.png "屏幕截图")
![输入图片说明](https://foruda.gitee.com/images/1761137034062160412/298f277f_7725520.png "屏幕截图")
---

**Self-checklist**:（**请自检，在[ ]内打上x，我们将检视你的完成情况，否则会导致pr无法合入**）

+ - [x] **设计**：PR对应的方案是否已经经过Maintainer评审，方案检视意见是否均已答复并完成方案修改
+ - [x] **测试**：PR中的代码是否已有UT/ST测试用例进行充分的覆盖，新增测试用例是否随本PR一并上库或已经上库
+ - [x] **验证**：PR描述信息中是否已包含对该PR对应的Feature、Refactor、Bugfix的预期目标达成情况的详细验证结果描述
+ - [x] **接口**：是否涉及对外接口变更，相应变更已得到接口评审组织的通过，API对应的注释信息已经刷新正确
+ - [x] **文档**：是否涉及官网文档修改，如果涉及请及时提交资料到Doc仓

Type	Assign personnel	Status
Reviewer	@chengbin is designated as reviewer by CODEOWNERS line 9	In progress (0/1 )
Reviewer	fangwenyi	Completed (0/0 )
Tester	fangwenyi	Completed (0/0 )

MindSpore/ms_custom_ops .gitee-modal { width: 500px !important; }

Search

MindSpore/ms_custom_ops