算子功能:MoE计算中,对输入x做Sigmoid或者SoftMax计算,对计算结果分组进行排序,最后根据分组排序的结果选取前k个专家。
计算公式:
Step 1: 归一化
根据 normType 对输入 x 做归一化:
Step 2: 加偏置
若 bias 不为空,加偏置得到用于选择的值:
否则 。
Step 3: 分组筛选(仅 groupCount > 1 时执行)
将 normValue 按 groupCount 分组,根据 groupSelectMode 计算每组得分:
选取得分最高的 kGroup 个组,将未选中组的对应位置置为 :
Step 4: Top-K 专家选择
对 normValue 取 Top-K 得到专家索引,这里只需要expertIdxOut:
Step 5: Renorm 与缩放
normType=1 时做归一化;normType=0 时,renorm 参数生效,renorm=1 时做renorm:
最终输出:
Step 6: 可选输出
若 outFlag 为 True,第三个输出为 normOut;否则为空。
每个算子分为两段式接口,必须先调用“aclnnMoeGatingTopKGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用“aclnnMoeGatingTopK”接口执行计算。
参数说明:
[object Object]返回值:
第一段接口完成入参校验,出现以下场景时报错:
[object Object][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表示不输出。