算子名称：InitRouing
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Atlas A2 训练系列产品 | 是 |
功能说明
- 算子功能：初始化路由，将输入token按expert分组并连续排列。
参数说明
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|---|---|---|---|---|
| blockDim | 输入 | AI CORE的数量，比如：Ascend910B是40。 | int64_t | - |
| stream | 输入 | Device端的stream | AclrtStream | - |
| in | 输入 | 公式中的输入张量x，shape为(token_num, hidden_size) | BFLOAT16 | ND |
| token_table | 输入 | token到expert的映射表， shape为(expert_num, token_num) | int32_t | ND |
| token_list | 输入 | 每个expert累计token数量，shape为(expert_num) | int64_t | ND |
| out | 输出 | 同一expert的token连续排列，shape为(expert_num*topk, hidden_size) | BFLOAT16 | ND |
| expert_num | 输入 | expert数量 | int64_t | ND |
| copy_byte | 输入 | 每个token拷贝字节数，等于hidden_size * element_size | int64_t | - |
约束说明
- token_list最后一个元素的值为总token数
- copy_byte按64bytes对齐
调用说明
torch.ops.npu_ops_transformer_ext.init_routing(block_dim, input, token_table,token_list, output)