MoeTokenPermute

产品支持情况

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

功能说明

• 算子功能：MoE的permute计算，根据索引indices将tokens广播并排序。
• 计算公式：

  • paddedMode为false时

    $$sortedIndicesFirst=argSort(indices)$$

    $$sortedIndicesOut=argSort(sortedIndicesFirst)$$

    $$permuteTokens[sortedIndicesFirst[i]]=tokens[i//topK]$$

  • paddedMode为true时

    $$permuteTokensOut[i]=tokens[Indices[i]]$$

    $$sortedIndicesOut=indices$$

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
tokens	输入	输入token，对应公式中的`tokens`。	FLOAT16、BFLOAT16、FLOAT32、INT8	ND
indices	输入	输入indices，对应公式中的`Indices`。tokens数据类型为INT8时，取值范围为`[0, 10240)`。	INT32、INT64	ND
numOutTokens	属性	有效输出token数，设置为0时，表示不会删除任何token。不为0时，会按照numOutTokens进行切片丢弃按照indices排序好的token中超过numOutTokens的部分，为负数时按照切片索引为负数时处理。	INT64	-
paddedMode	属性	paddedMode为true时表示indices已被填充为代表每个专家选中的token索引，此时不对indices进行排序。目前仅支持paddedMode为false。	BOOL	-
permuteTokensOut	输出	根据indices进行扩展并排序过的tokens，对应公式中的`permuteTokensOut`。	FLOAT16、BFLOAT16、FLOAT32、INT8	ND
sortedIndicesOut	输出	permuteTokensOut和tokens的映射关系，对应公式中的`sortedIndicesOut`。	INT32	ND

约束说明

• indices要求元素个数小于16777215，值大于等于0小于16777215(单点支持int32或int64的最大或最小值，其余值不在范围内排序结果不正确)。在Ascend 950上调用aclnnMoeTokenPermute且tokens数据类型为INT8时，indices表示expert ID，取值范围为[0, 10240)。
• 调用aclnnMoeTokenPermute时，tokens必须为2D，indices必须为1D或2D。
• 调用aclnnMoeTokenPermuteV2时，tokens必须为2D，indices必须为1D或2D。
• 不支持paddedMode为true。
• Atlas A2 训练系列产品/Atlas A2 推理系列产品、Atlas A3 训练系列产品/Atlas A3 推理系列产品：topK小于等于512。
• Ascend 950PR/Ascend 950DT： 在调用本接口时，框架内部会转调用aclnnMoeInitRoutingV2接口，如果出现参数错误提示，请参考以下参数映射关系：
  • token输入等同于aclnnMoeInitRoutingV2接口的x输入。
  • indices输入等同于aclnnMoeInitRoutingV2接口的expertIdx输入。
  • numOutTokens输入等同于aclnnMoeInitRoutingV2接口的activeNum输入。
  • paddedMode输入等同于aclnnMoeInitRoutingV2接口的dropPadMode输入。
  • permuteTokensOut输出等同于aclnnMoeInitRoutingV2接口的expandedXOut输出。
  • sortedIndicesOut输出等同于aclnnMoeInitRoutingV2接口的expandedRowIdxOut输出。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_moe_token_permute	通过aclnnMoeTokenPermute接口方式调用MoeTokenPermute算子。