开发者
资源
[object Object]

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

  • 计算公式:

    Step 1: 归一化

    根据 normType 对输入 x 做归一化:

    normOut={SoftMax(x),normType=0Sigmoid(x),normType=1normOut = \begin{cases} SoftMax(x), & normType = 0 \\ Sigmoid(x), & normType = 1 \end{cases}

    Step 2: 加偏置

    若 bias 不为空,加偏置得到用于选择的值:

    normValue=normOut+biasnormValue = normOut + bias

    否则 normValue=normOutnormValue = normOut

    Step 3: 分组筛选(仅 groupCount > 1 时执行)

    将 normValue 按 groupCount 分组,根据 groupSelectMode 计算每组得分:

    groupedValue=Reshape(normValue, [batch, groupCount, 1])groupedValue = Reshape(normValue,\ [batch,\ groupCount,\ -1]) groupScore={ReduceMax(groupedValue, dim=1),groupSelectMode=0ReduceSum(TopK(groupedValue, k=2, dim=1), dim=1),groupSelectMode=1groupScore = \begin{cases} ReduceMax(groupedValue,\ dim=-1), & groupSelectMode = 0 \\ ReduceSum(TopK(groupedValue,\ k=2,\ dim=-1),\ dim=-1), & groupSelectMode = 1 \end{cases}

    选取得分最高的 kGroup 个组,将未选中组的对应位置置为 -\infty

    groupIdx=TopK(groupScore, k=kGroup).indicesgroupIdx = TopK(groupScore,\ k=kGroup).indices normValue=Mask(groupedValue, groupIdx, fillValue=)normValue = Mask(groupedValue,\ groupIdx,\ fillValue=-\infty)

    Step 4: Top-K 专家选择

    对 normValue 取 Top-K 得到专家索引,这里只需要expertIdxOut:

    y,expertIdxOut=TopK(normValue[groupIdx,:], k=k)y, expertIdxOut = TopK(normValue[groupIdx, :],\ k=k)

    Step 5: Renorm 与缩放

    normType=1 时做归一化;normType=0 时,renorm 参数生效,renorm=1 时做renorm:

    if (normType=1) or (normType=0 and renorm=1):if\ (normType = 1)\ or\ (normType = 0\ and\ renorm = 1): yOut=normOutReduceSum(normOut, dim=1)+eps\quad yOut = \frac{normOut}{ReduceSum(normOut,\ dim=-1) + eps}

    最终输出:

    yOut=yOut×routedScalingFactoryOut = yOut \times routedScalingFactor

    Step 6: 可选输出

    若 outFlag 为 True,第三个输出为 normOut;否则为空。

[object Object]

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

[object Object]
[object Object]
[object Object]
  • 参数说明:

    [object Object]
  • 返回值:

    aclnnStatus:返回状态码,具体参见

    第一段接口完成入参校验,出现以下场景时报错:

    [object Object][object Object]
[object Object]
  • 参数说明:

    [object Object]
  • 返回值:

    返回aclnnStatus状态码,具体参见

[object Object]
  • 确定性计算:
    • 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和1,表示先进行norm操作,再计算topk。
  • 其他限制:
    • groupSelectMode取值0和1,0表示使用最大值对group进行排序, 1表示使用topk2的sum值对group进行排序。
    • normType取值0和1,0表示使用Softmax函数,1表示使用Sigmoid函数。
    • outFlag取值true和false,true表示输出,false表示不输出。
[object Object]

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

[object Object]