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$$

函数原型

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

• 返回值：

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

  [object Object]

aclnnMoeGatingTopK

• 参数说明：

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

• 返回值：

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

约束说明

• 专家个数不超过2048个。

调用示例

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