算子名称：FinalRouting

产品支持情况

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

功能说明

• 算子功能：用于MoE模型的最终路由阶段，将各个Expert计算的结果按照评分加权合并，得到每个Token的最终输出，完成MoE的Combine阶段。

• 计算公式： 对于每个Token t：

$$out[t]=\sum _e(in[token\_table[t,e]]\cdot score\_table[t,e])$$

其中，仅当 token_table[t, e] >= 0 时参与计算。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
blockDim	输入	AI CORE的数量，比如：Ascend910B是40。	int64_t	-
in	输入	expert输出张量， shape为(expert_num*token_num, hidden_size)	BFLOAT16	ND
token_table	输入	token到expert的映射表， shape为(token_num, expert_num)	int32_t	ND
score_table	输入	每个token在每个expert的评分，shape为(token_num, expert_num)	BFLOAT16	ND
out	输出	加权合并后的输出张量，shape为(token_num, hidden_size)	BFLOAT16	ND

约束说明

• token_table中小于0的值表示该expert对此token无效

调用说明

torch.ops.npu_ops_transformer_ext.final_routing(block_dim, input, token_table,score_table, output)