aclnnMoeGatingTopKSoftmaxV2

产品支持情况

产品	是否支持
[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计算中，如果renorm=0，先对x的输出做Softmax计算，再取topk操作；如果renorm=1，先对x的输出做topk操作，再进行Softmax操作。其中yOut为softmax的topk结果；expertIdxOut为topk的indices结果即对应的专家序号；如果对应的行finished为True，则expert序号直接填num_expert值（即x的最后一个轴大小）。
• 计算公式：

1. renorm = 0,$$softmaxResultOutOptional=softmax(x,axis=-1)$$ $$yOut,expertIdxOut=topK(softmaxResultOutOptional,k=k)$$
2. renorm = 1$$topkOut,expertIdxOut=topK(x, k=k)$$ $$yOut = softmax(topkOut,axis=-1)$$

函数原型

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

• aclnnStatus aclnnMoeGatingTopKSoftmaxV2GetWorkspaceSize(const aclTensor *x, const aclTensor *finishedOptional, int64_t k, int64_t renorm, bool outputSoftmaxResultFlag, const aclTensor *yOut, const aclTensor *expertIdxOut, const aclTensor *softmaxResultOutOptional, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnMoeGatingTopKSoftmaxV2(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnMoeGatingTopKSoftmaxV2GetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：待计算的输入，要求是一个2D/3D的Tensor，数据类型支持FLOAT16、BFLOAT16、FLOAT32，undefined要求为ND，支持undefined。
  • finishedOptional（aclTensor*，可选计算输入）：要求是一个1D/2D的Tensor，数据类型支持bool，shape为x_shape[:-1]，undefined要求为ND，支持undefined。
  • k（int64_t，计算输入）：topk的k值，大小为0 < k <= x的-1轴大小，且k不大于1024。
  • renorm (int64_t, 计算输入)：renorm标记，取值0和1。0表示先计算Softmax，再计算TopK；1表示先计算TopK，再计算Softmax。
  • outputSoftmaxResultFlag (bool, 计算输入)：表示是否输出softmax的结果，取值true和false。当renorm=0时，true表示输出Softmax的结果，false表示不输出；当renorm=1时，该参数不生效，不输出Softmax的结果。
  • yOut（aclTensor*，计算输出）：对x做softmax后取的topk值，要求是一个2D/3D的Tensor，数据类型支持FLOAT16、BFLOAT16、FLOAT32，数据类型与x需要保持一致，其非-1轴要求与x的对应轴大小一致，其-1轴要求其大小同k值，undefined要求为ND，不支持undefined。
  • expertIdxOut（aclTensor*，计算输出）：对x做softmax后取topk值的索引，即专家的序号，shape要求与yOut一致，数据类型支持int32，undefined要求为ND，不支持undefined。
  • softmaxResultOutOptional（aclTensor*，可选输出）：计算过程中Softmax的结果（见示例），shape要求与x一致，数据类型支持FLOAT32，undefined要求为ND，不支持undefined。
  • workspaceSize（uint64_t*，出参）：Device侧的整型，返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：Device侧的aclOpExecutor，返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnMoeGatingTopKSoftmaxV2

• 参数说明：

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

• 返回值：

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

约束说明

k的值不大于1024。 renorm的值只支持0和1。 x和finishedOptional的每一维大小应不大于int32的最大值2147483647。

调用示例

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