aclnnMoeFusedTopk

产品支持情况

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

  $$sigmoidRes=sigmoid(x)$$

  加上addNum：

  $$normOut = sigmoidRes + addNum$$

  对计算结果按照groupNum进行分组，每组按照topN的sum值对group进行排序，取前groupTopk个组：

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

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

  $$normY,indices=TopK(normOut[groupId, :],k=k)$$

  根据indices从sigmoidRes中选出y:

  $$y = gather(sigmoidRes, indices)$$

  如果isNorm为true，对y按照输入的scale参数进行计算，得到y的结果：

  $$y = y / (ReduceSum(y, dim=-1))*scale$$

  如果enableExpertMapping为true，再将indices中的物理专家按照输入的mappingNum和mappingTable映射到逻辑专家，得到输出的indices。

函数原型

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

• aclnnMoeFusedTopkGetWorkspaceSize(const aclTensor* x, const aclTensor* addNum, const aclTensor* mappingNum, const aclTensor* mappingTable, uint32_t groupNum, uint32_t groupTopk, uint32_t topN, uint32_t topK, uint32_t activateType, bool isNorm, float scale, bool enableExpertMapping, aclTensor* y, aclTensor* indices, uint64_t* workspaceSize, aclOpExecutor** executor)

• aclnnStatus aclnnMoeFusedTopk(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnMoeFusedTopkGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：Device侧的aclTensor，每个token对应各个专家的分数，shape为(numToken, expertNum)，数据类型支持FLOAT16、BFLOAT16、FLOAT32，要求为ND，支持。
  • addNum（aclTensor*，计算输入）：Device侧的aclTensor，与输入x进行计算的偏置值，shape为(expertNum)，数据类型要求与x一致，要求为ND，支持。
  • mappingNum（aclTensor*，计算输入）：Device侧的aclTensor，enableExpertMapping为false时不启用，shape为(expertNum)，每个物理专家被实际映射到的逻辑专家数量，数据类型支持INT32，要求为ND，支持。
  • mappingTable（aclTensor*，计算输入）：Device侧的aclTensor，enableExpertMapping为false时不启用，shape为(expertNum, maxMappingNum)，每个物理专家/逻辑专家映射表，maxMappingNum小于等于128，数据类型支持INT32，要求为ND，支持。
  • groupNum (uint32_t，计算输入)：分组数量，必须大于0。
  • groupTopk (uint32_t，计算输入)：被选择的组的数量，必须大于0。
  • topN (uint32_t，计算输入)：组内选取的用于求和的专家数量，必须大于0。
  • topK (uint32_t，计算输入)：最终选取的专家数量，必须大于0。
  • activateType (uint32_t，计算输入)：激活类型，当前只支持0(ACTIVATION_SIGMOID)。
  • isNorm (bool，计算输入)：是否对输出进行归一化。
  • scale (float，计算输入)：归一化后的系数乘。
  • enableExpertMapping (bool，计算输入)：是否使能物理专家到逻辑专家的映射。
  • y (aclTensor*，计算输出)：Device侧的aclTensor，shape为(numToken, topK)，数据类型支持FLOAT32，要求为ND，支持。
  • indices (aclTensor*，计算输出)：Device侧的aclTensor，shape为(numToken, topK)，数据类型支持INT32，要求为ND，支持。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnMoeFusedTopk

• 参数说明：

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

• 返回值：

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

约束说明

• expertNum必须为groupNum的整数倍。
• groupTopk小于等于groupNum。
• maxMappingNum小于等于128。
• TopK小于等于expertNum。
• TopN小于等于expertNum / groupNum。
• expertNum小于等于1024。
• groupNum小于等于256。

调用示例

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