aclnnMoeGatingTopKSoftmax

支持的产品型号

• [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品[object Object]

接口原型

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

• aclnnStatus aclnnMoeGatingTopKSoftmaxGetWorkspaceSize(const aclTensor *x, const aclTensor *finishedOptional, int64_t k, const aclTensor *yOut, const aclTensor *expertIdxOut, const aclTensor *rowIdxOut, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnMoeGatingTopKSoftmax(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

功能描述

• 算子功能：MoE计算中，对x的输出做Softmax计算，取topk操作。其中yOut为softmax的topk结果；expertIdxOut为topk的indices结果即对应的专家序号；rowIdxOut为与expertIdxOut相同shape的列取值结果。如果对应的行finished为True，则expert序号直接填num_expert值（即x的最后一个轴大小）。
• 计算公式：$$softmaxOut=softmax(x,axis=-1)$$ $$yOut,expertIdxOut=topK(softmaxOut,k=k)$$ $$rowIdxRange=arange(expertIdxOut.shape[0]*expertIdxOut.shape[1])$$ $$rowIdxOut=rowIdxRange.reshape([expertIdxOut.shape[1],expertIdxOut.shape[0]]).transpose(1,0)$$

aclnnMoeGatingTopKSoftmaxGetWorkspaceSize

• 参数说明：

  • 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。
  • 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。
  • rowIdxOut（aclTensor*，计算输出）：指示每个位置对应的原始行位置（见示例），shape要求与yOut一致，数据类型支持int32，undefined要求为ND，不支持undefined。
  • workspaceSize（uint64_t*，出参）：Device侧的整型，返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：Device侧的aclOpExecutor，返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnMoeGatingTopKSoftmax

• 参数说明：

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

• 返回值：

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

约束与限制

k的值不大于1024。 x和finishedOptional的每一维大小应不大于int32的最大值2147483647。

调用示例

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