aclnnMoeTokenUnpermuteWithRoutingMap

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	√
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

• 算子功能：对经过aclnnMoeTokenpermuteWithRoutingMap处理的permutedTokens，累加回原unpermutedTokens。根据sortedIndices存储的下标，获取permutedTokens中存储的输入数据；如果存在probs数据，permutedTokens会与probs相乘，最后进行累加求和，并输出计算结果。

• 计算公式：

  $$topK\_num= permutedTokens.size(0) // routingMapOptional.size(0)$$ $$numExperts = probs.size(1)$$ $$numTokens = probs.size(0)$$ $$capacity = sortedIndices.size(0) // numExperts$$

  (1)probs不为None，padMode为true时：

  $$permutedProbs [i//capacity,sortedIndices[i]]=probs[i]$$ $$permutedTokens = permutedTokens * permutedProbs$$ $$unpermutedTokens= zeros(restoreShape, dtype=permutedTokens.dtype, device=permutedTokens.device)$$ $$permuteTokenId, outIndex= sortedIndices.sort(dim=-1)$$ $$unpermutedTokens[permuteTokenId[i]] += permutedTokens[outIndex[i]]$$

  (2)probs不为None，padMode为false时:

  $$permutedProbs = probs.T.maskedSelect(routingMap.T)$$ $$permutedTokens = permutedTokens * permutedProbs$$ $$unpermutedTokens= zeros(restoreShape, dtype=permutedTokens.dtype, device=permutedTokens.device)$$ $$unpermutedTokens[i//topK\_num] += permutedTokens[sortedIndices[i]]$$

  (3)probs为None,padMode为true时:

  $$permuteTokenId, outIndex= sortedIndices.sort(dim=-1)$$ $$unpermutedTokens[permuteTokenId[i]] += permutedTokens[outIndex[i]]$$

  (4)probs为None,padMode为false时:

  $$unpermutedTokens[i//topK\_num] += permutedTokens[sortedIndices[i]]$$

函数原型

每个算子分为undefined，必须先调用“aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnMoeTokenUnpermuteWithRoutingMap”接口执行计算。

• aclnnStatus aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize(const aclTensor *permutedTokens, const aclTensor *sortedIndices, const aclTensor* routingMapOptional, const aclTensor *probsOptional, bool paddedMode, const aclIntArray *restoreShapeOptional, aclTensor *unpermutedTokens, aclTensor *outIndex, aclTensor *permuteTokenId, aclTensor *permuteProbs, uint64_t *workspaceSize, aclOpExecutor **executor);
• aclnnStatus aclnnMoeTokenUnpermuteWithRoutingMap(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize

• 参数说明：

  • permutedTokens（aclTensor*，计算输入）：Device侧的aclTensor，输入token，要求为一个维度为2D的Tensor，当paddedMode为false时，shape为（tokens_num * topK_num， hidden_size），当paddedMode为true时，shape为（experts_num* capacity， hidden_size），capacity表示每个专家能够处理的token个数，数据类型支持BFLOAT16、FLOAT16、FLOAT，undefined要求为ND。支持undefined。
  • sortedIndices（aclTensor *，计算输入）：Device侧的aclTensor，非droppad模式要求shape为一个1D的（tokens_num * topK_num，），数据类型支持INT32，undefined要求为ND，索引取值范围[0，tokens_num * topK_num - 1]。droppad模式要求shape为一个1D的（experts_num * capacity），数据类型支持INT32，undefined要求为ND，索引取值范围[0，tokens_num - 1]。支持undefined。
  • routingMapOptional（aclTensor*，计算输入）：Device侧的aclTensor，可选输入，当输入probsOptional为空指针时不需要此输入，应该传入空指针。计算公式中的routingMapOptional，代表对应位置的Token是否被对应专家处理，要求shape为一个2D的（tokens_num，experts_num），数据类型支持INT8、bool。当数据类型为INT8，取值支持0、1，当数据类型为bool，取值支持true、false，undefined要求为ND。支持undefined。
  • probsOptional（aclTensor*，计算输入）：Device侧的aclTensor，可选输入，当不需要时为空指针。计算公式中的probsOptional，代表对应位置的Token被对应专家处理后的结果在最终结果中的权重，shape与routingMapOptional相同，数据类型与permutedTokens相同，undefined要求为ND。支持undefined。
  • paddedMode（bool, 计算输入）：host侧的BOOL。可选输入，支持取值为false和true。true表示开启paddedMode，false表示关闭paddedMode，开启paddedMode时，输出outIndex、permuteTokenId的shape为（experts_num* capacity，），关闭paddedMode时，每个token固定被topK_num个专家处理，输出outIndex、permuteTokenId的shape为（tokens_num * topK_num，）。
  • restoreShapeOptional（aclIntArray*，计算输入）：host侧的aclIntArray。支持的数据类型为INT32, size大小为2。为unpermutedTokens的shape。
  • unpermutedTokens（aclTensor*，计算输出）：Device侧的aclTensor，正向输出结果，计算公式中的unpermutedTokens，要求为一个维度为2D的Tensor，shape为（tokens_num，hidden_size），数据类型支持BFLOAT16、FLOAT16、FLOAT，undefined要求为ND。支持undefined。
  • outIndex（aclTensor*，计算输出）：Device侧的aclTensor，计算公式中的outIndex，当paddedMode为false时，要求shape为一个1D的（tokens_num * topK_num，），索引取值范围[0，tokens_num * topK_num - 1]。当paddedMode为true时，要求shape为一个1D的（experts_num* capacity，）。索引取值范围[0，experts_num* capacity- 1]。数据类型支持INT32，undefined要求为ND。支持undefined。
  • permuteTokenId（aclTensor*，计算输出）：Device侧的aclTensor，计算公式中的permuteTokenId，当paddedMode为false时，要求shape为一个1D的（tokens_num * topK_num，）。当paddedMode为true时，要求shape为一个1D的（experts_num* capacity，）。索引取值范围[0，tokens_num - 1]。数据类型支持INT32，undefined要求为ND。支持undefined。
  • permuteProbs（aclTensor *，计算输出）：Device侧的aclTensor, 计算公式中的permutedProbs,表示输出经过排序后的probs，shape支持1D维度。数据类型同probsOptional，undefined要求为ND。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

  aclnnStatus：返回状态码，具体参见undefined。

  [object Object]

aclnnMoeTokenUnpermuteWithRoutingMap

• 参数说明：

  • workspace（void*，入参）：在Device侧申请的workspace内存地址。
  • workspaceSize（uint64_t，入参）：在Device侧申请的workspace大小，由第一段接口aclnnMoeTokenUnpermuteWithRoutingMapGetWorkspaceSize获取。
  • executor（aclOpExecutor*，入参）：op执行器，包含了算子计算流程。
  • stream（aclrtStream，入参）：指定执行任务的Stream。

• 返回值：

  返回aclnnStatus状态码，具体参见undefined。

约束说明

topkNum <= 512, pad模式为false时routingMap中每行为1或true的个数固定且小于512。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考undefined。