MoeTokenPermuteWithEp

产品支持情况

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

功能说明

算子功能：MoE的permute计算，根据索引indices将tokens和可选probs广播后排序并按照rangeOptional中范围切片。

计算公式：

• paddedModefalse时

  $$sortedIndicesFirst=argSort(indices)$$

  $$sortedIndicesOut=argSort(sortedIndicesFirst)$$

  当rangeOptional[0] <= sortedIndicesOut[i] < rangeOptional[1]时

  $$permuteTokensOut[sortedIndicesOut[i]-range[0]]=tokens[i//topK]$$

  $$permuteProbsOut[sortedIndicesOut[i]-rangeOptional[0]]=probsOptional[i]$$

  • paddedMode为true时

    $$permuteTokensOut[i]=tokens[indices[i]]$$

    $$sortedIndicesOut=indices$$

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
tokens	输入	permute中的输入tokens，公式中的`tokens`。	BFLOAT16、FLOAT16、FLOAT32	ND
indices	输入	输入tokens对应的专家索引，公式中的`indices`。	INT32、INT64	ND
probsOptional	输入	可选输入，输入tokens对应的专家概率，公式中的`probsOptional`。	BFLOAT16、FLOAT16、FLOAT32	ND
rangeOptional	属性	ep切分的有效范围。	aclIntArray	-
numOutTokens	属性	有效输出token数，在rangeOptional为空时生效。	INT64	-
paddedMode	属性	为true时表示indices已被填充为代表每个专家选中的token索引。	BOOL	-
permuteTokensOut	输出	indices进行扩展并排序过的tokens，公式中的`permuteTokensOut`。	BFLOAT16、FLOAT16、FLOAT32	ND
sortedIndicesOut	输出	排序后的输出结果。	INT32	ND
permuteProbsOut	输出	permute之后的输出。	BFLOAT16、FLOAT16、FLOAT32	ND

约束说明

• indices 要求元素个数小于16777215，值大于等于0且小于16777215(单点支持int32或int64的最大或最小值，其余值不在范围内排序结果不正确)。
• topK小于等于512。
• 不支持paddedMode为true。
• 当rangeOptional为空时，忽略probsOptional和permuteTokensOut，执行逻辑回退到aclnnMoeTokenPermute。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_moe_token_permute_with_ep.cpp	通过aclnnMoeTokenPermuteWithEp接口方式调用MoeTokenPermuteWithEp算子。