MoeInitRouting

支持的产品型号

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：MoE的routing计算，根据aclnnMoeGatingTopKSoftmax的计算结果做routing处理。

• 计算公式：

  $$expandedExpertIdx,sortedRowIdx=keyValueSort(expertIdx,rowIdx)$$

  $$expandedRowIdx[sortedRowIdx[i]]=i$$

  $$expandedX[i]=x[sortedRowIdx[i]\%numRows]$$

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	MOE的输入即token特征输入，对应公式中的`x`。	FLOAT16、BFLOAT16、FLOAT32	ND
rowIdx	输入	指示每个位置对应的原始行位置，对应公式中的`rowIdx`。	INT32	ND
expertIdx	输入	aclnnMoeGatingTopKSoftmax的输出每一行特征对应的K个处理专家, 对应公式中的`expertIdx`。	INT32	ND
activeNum	属性	表示总的最大处理row数且大于等于0，expandedXOut只有这么多行是有效的。	INT64	-
expandedXOut	输出	根据expertIdx进行扩展过的特征，对应公式中的`expandedX`。	FLOAT16、BFLOAT16、FLOAT32	ND
expandedRowIdxOut	输出	expandedX和x的映射关系，对应公式中的`expandedRowIdx`。	INT32	ND
expandedExpertIdxOut	输出	输出expertIdx排序后的结果，对应公式中的`expandedExpertIdx`。	INT32	ND

约束说明

• expertIdx内的元素的值需要大于-2**24，不超过2**24，否则可能会存在精度问题。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_moe_init_routing	通过aclnnMoeInitRouting接口方式调用MoeInitRouting算子。