aclnnMoeGatingTopK

产品支持情况

[object Object]undefined

功能说明

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

• 计算公式：

  对输入做Sigmoid或者SoftMax：

  $$if normType==1: normOut=Sigmoid(x) else: normOut=SoftMax(x)$$

  如果bias不为空：

  $$normOut = normOut + bias$$

  对计算结果按照groupCount进行分组，每组按照groupSelectMode取max或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”接口执行计算。

aclnnMoeGatingTopKGetWorkspaceSize

• 参数说明：

  [object Object]

• 返回值：

  [object Object][object Object]

aclnnMoeGatingTopK

• 参数说明：

  [object Object]

• 返回值：

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

约束说明

• 确定性计算：
  • aclnnMoeGatingTopK默认确定性实现。

• 输入shape限制：
  • x最后一维（即专家数）要求不大于2048。
• 输入值域限制：
  • 要求1 <= k <= x_shape[-1] / groupCount * kGroup。
  • 要求1 <= kGroup <= groupCount，并且kGroup * x_shape[-1] / groupCount的值要大于等于k。
  • 要求groupCount > 0，x_shape[-1]能够被groupCount整除且整除后的结果大于groupSelectMode，并且整除的结果按照32个数对齐后乘groupCount的结果不大于2048。
  • renorm仅支持0，表示先进行norm操作，再计算topk。
• 其他限制：
  • groupSelectMode取值0和1，0表示使用最大值对group进行排序, 1表示使用topk2的sum值对group进行排序。
  • normType取值0和1，0表示使用Softmax函数，1表示使用Sigmoid函数。
  • outFlag取值true和false，true表示输出，false表示不输出。

调用示例

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