aclnnMoeGatingTopK

产品支持情况

产品	是否支持
[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]	×

功能说明

• 算子功能：MoE计算中，对输入x做Sigmoid计算，对计算结果分组进行排序，最后根据分组排序的结果选取前k个专家。

• 计算公式：

  对输入做sigmoid：

  $$normOut=sigmoid(x)$$

  如果bias不为空：

  $$normOut = normOut + bias$$

  对计算结果按照groupCount进行分组，每组按照topk2的sum值对group进行排序，取前kGroup个组：

  $$groupOut, groupId = TopK(ReduceSum(TopK(Split(normOut, groupCount), k=2, dim=-1), dim=-1),k=kGroup)$$

  根据上一步的groupId获取normOut中对应的元素，将数据再做TopK，得到expertIdxOut的结果：

  $$y,expertIdxOut=TopK(normOut[groupId, :],k=k)$$

  对y按照输入的routedScalingFactor和eps参数进行计算，得到yOut的结果：

  $$yOut = y / (ReduceSum(y, dim=-1)+eps)*routedScalingFactor$$

函数原型

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

• aclnnStatus aclnnMoeGatingTopKGetWorkspaceSize(const aclTensor *x, const aclTensor *biasOptional, int64_t k, int64_t kGroup, int64_t groupCount, int64_t groupSelectMode, int64_t renorm, int64_t normType, bool outFlag, double routedScalingFactor, double eps, const aclTensor *yOut, const aclTensor *expertIdxOut, const aclTensor *outOut, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnMoeGatingTopK(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnMoeGatingTopKGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：Device侧的aclTensor，待计算的输入，要求是一个2D的Tensor，数据类型支持FLOAT16、BFLOAT16、FLOAT32，undefined要求为ND，支持undefined。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：x最后一维的大小（即专家数）当前只支持256。
  • biasOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，与输入x进行计算的bias值，要求是一个1D的Tensor，数据类型支持FLOAT16、BFLOAT16、FLOAT32，数据类型与x需要保持一致，undefined要求为ND，支持undefined，要求biasOptional的shape值与x的最后一维相等。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：必须不为None。
  • k（int64_t，计算输入）：topk的k值。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值范围为[1,32]。
  • kGroup（int64_t，计算输入）：分组排序后取的group个数。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当前仅支持取4。
  • groupCount（int64_t，计算输入）：分组的总个数。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当前仅支持取8。
  • groupSelectMode（int64_t，计算输入）：分组排序的方式，当前仅支持1，表示使用topk2的sum值对group进行排序。
  • renorm（int64_t，计算输入）：renorm标记，当前仅支持0，表示先进行norm操作，再计算topk。
  • normType（int64_t，计算输入）：norm函数类型，当前仅支持1，表示使用Sigmoid函数。
  • outFlag（bool，计算输入）：表示是否输出norm操作的结果，当前仅支持false，表示不输出。
  • routedScalingFactor（double，计算输入）：用于计算yOut使用的routedScalingFactor系数。
  • eps（double，计算输入）：用于计算yOut使用的eps系数。
  • yOut（aclTensor*，计算输出）：Device侧的aclTensor，对x做norm操作，分组排序topk后计算的结果，要求是一个2D的Tensor，数据类型支持FLOAT16、BFLOAT16、FLOAT32，数据类型与x需要保持一致，undefined要求为ND，不支持undefined，第一维的大小要求与x的第一维相同，最后一维的大小与k相同。
  • expertIdxOut（aclTensor*，计算输出）：Device侧的aclTensor，对x做norm操作，分组排序topk后的索引，即专家的序号，shape要求与yOut一致，数据类型支持INT32，undefined要求为ND，不支持undefined。
  • outOut（aclTensor*，计算输出）：Device侧的aclTensor，norm计算的输出结果，通过outFlag参数控制是否生效，shape要求与x保持一致，数据类型支持FLOAT32，undefined要求为ND，不支持undefined。
  • workspaceSize（uint64_t*，出参）：Device侧的整型，返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：Device侧的aclOpExecutor，返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnMoeGatingTopK

• 参数说明：

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

• 返回值：

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

约束说明

• [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：专家个数为256。

调用示例

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