MoeTokenUnpermuteWithEp

产品支持情况

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

功能说明

• 算子功能：根据sortedIndices存储的下标位置，去获取permutedTokens中的输入数据与probs相乘，并进行合并累加。
• 计算公式：

$$sortedIndices=sortedIndices[rangeOptional[0]<=i<rangeOptional[1]]$$

（1）probs非None计算公式如下，其中$i\in 0,1,2,...,num\_tokens-1$，$j\in 0,1,2,...,numTopk-1$，$k\in 0,1,2,...,num\_tokens\ast numTopk$：

$$permutedTokens=permutedTokens.indexSelect(0,sortedIndices)$$

$$permutedToken{s}_k=permutedToken{s}_k\ast prob{s}_{i,j}$$

$$ou{t}_i=\sum _{k=i\ast numTopk}^{(i+1)\ast numTopk-1}permutedToken{s}_k$$

（2）probs为None计算公式如下，其中$i\in 0,1,2,...,num\_tokens-1$，$j\in 0,1,2,...,numTopk-1$：

$$permutedTokens=permutedTokens.indexSelect(0,sortedIndices)$$

$$ou{t}_i=\sum _{k=i\ast numTopk}^{(i+1)\ast numTopk-1}permutedToken{s}_k$$

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
permutedTokens	输入	表示经过扩展并排序过的tokens，对应公式中的`permutedTokens`。	FLOAT16、BFLOAT16、FLOAT32	ND
sortedIndices	输入	表示需要计算的数据在permutedTokens中的位置，对应公式中的`sortedIndices`。	INT32	ND
probsOptional	可选输入	表示输入tokens对应的专家概率，对应公式中的`probs`。	FLOAT16、BFLOAT16、FLOAT32	ND
numTopk	属性	被选中的专家个数。	INT64	-
rangeOptional	属性	ep切分的有效范围，size为2。	aclIntArray*	-
paddedMode	属性	目前仅支持false。	BOOL	-
restoreShapeOptional	属性	目前仅支持nullptr。	aclIntArray*	-
out	输出	表示permutedTokens反重排的输出结果，对应公式中的`out`。	FLOAT16、BFLOAT16、FLOAT32	ND

• Kirin X90/Kirin 9030 处理器系列产品: 不支持BFLOAT16。

约束说明

• numTopk <= 512。
• 不支持paddedMode为True。
• 当rangeOptional为空时，忽略numTopk，执行逻辑回退到aclnnMoeTokenUnpermute。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_moe_token_unpermute_with_ep	通过aclnnMoeTokenUnpermuteWithEp接口方式调用MoeTokenUnpermuteWithEp算子。