算子名称：GetRoutingConfigV2

产品支持情况

产品	是否支持
Atlas A2 训练系列产品	是

功能说明

• 算子功能：moeinitrouting计算前的数据准备工作。

• 计算公式：

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
blockDim	输入	AI CORE的数量，比如：Ascend910B是40。	int64_t	-
stream	输入	Device端的stream	AclrtStream	-
indices	输入	专家索引，shape(token, topk)	int32_t	ND
scores	输入	专家得分，shape(token, topk)	bfloat16	ND
init_token_table	输出	当前卡上每个token对应的专家的token位置，shape为((expert, token))	int64	ND
final_token_table	输出	每条routing信息（token-expert）的编号， shape为(token, expert)	int32	ND
final_score_table	输出	每条routing信息的对应得分， shape(token, expert)	BFLOAT16	ND
token_idx_intra_expert	输出	中间结果，每条routing信息在其对应专家内部的编号，是initTokenTable的逆映射，shape(expert, token)	int32	ND
start_expert_id	输入	当前eprank上专家的起始id	int64	-
end_expert_id	输入	当前eprank上最后专家的id的后一位id	int64	-
local_expert_num	输入	每个eprank上的专家数	int64	-
token_num	输入	所有eprank总token数	int64	-
ub_max_token	输入	ub能存放的最大token, 1024	int64_t	-

约束说明

• 输入输出仅支持BFLOAT16类型。

价值/作用

为MoERouting计算前准备必要的数据。

设计方案

Tilling策略

• 分核策略： init_table：分割tokenNum init_list：单核计算 init_other：分割localExpertNum，得到每个expert的专属token编号 final：分割tokenNum

Kernel侧设计

进行Init和Process两个阶段，其中Process包括数据搬入（CopyIn）、计算（Compute）、数据搬出（CopyOut）三个阶段，实现RoutingConfig算子计算。

• 初始化（Init）
  • 计算loop_count
  • 建立GM Tensor映射，初始化队列缓冲区，初始化临时缓冲区
• 计算流程（Compute）
  • 调用CAST将输入转为float类型
  • 调用Ascendc命令完成计算
  • 调用CAST将fp32表示的结果转为bfloat16/half类型
• 数据搬入（CopyIn）搬出（CopyOut）
  • CopyIn: 将输入数据从inGM中搬入到inQueue
  • CopyOut: 将输出结果数据从outQueue搬出到outGM
• 流程调度（Process）
  • 遍历loop_count，按CopyIn->Compute->CopyOut的顺序调度