aclnnMoeInitRoutingV2

产品支持情况

产品	是否支持
[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（Mixture of Experts，混合专家模型）中的Routing计算，以MoeGatingTopKSoftmax算子的输出x和expert_idx作为输入，并输出Routing矩阵expanded_x等结果供后续计算使用。本接口针对V1接口（MoeInitRouting，源码未开放）做了如下功能变更，请根据实际情况选择合适的接口：

  • 新增Drop模式，在该模式下输出内容会将每个专家需要处理的Token个数对齐为expertCapacity个，超过expertCapacity个的Token会被Drop，不足的会用0填充。
  • 新增Dropless模式下expertTokensCountOrCumsumOut可选输出，输出每个专家需要处理的累积Token个数（Cumsum），或每个专家需要处理的Token数（Count）。
  • 新增Drop模式下expertTokensBeforeCapacityOut可选输出，输出每个专家在Drop前应处理的Token个数。
  • 删除rowIdx输入。

  说明： Routing计算是MoE模型中的一个环节。MoE模型主要由一组专家模型和一个门控模型组成，在计算时，输入的数据会先根据门控网络（Gating Network，包含MoeGatingTopKSoftmax算子）计算出每个数据元素对应权重最高的k个专家，然后该结果会输入MoeInitRouting算子，生成Routing矩阵。在后续，模型中的每个专家会根据Routing矩阵处理其应处理的数据，产生相应的输出。各专家的输出最后与权重加权求和，形成最终的预测结果。

• 计算公式：

  1.对输入expertIdx做排序，得出排序后的结果sortedExpertIdx和对应的序号sortedRowIdx：

  $$sortedExpertIdx, sortedRowIdx=keyValueSort(expertIdx)$$

  2.以sortedRowIdx做位置映射得出expandedRowIdxOut：

  $$expandedRowIdxOut[sortedRowIdx[i]]=i$$

  3.对x取前numRows个sortedRowIdx的对应位置的值，得出expandedXOut：

  $$expandedXOut[i]=x[sortedRowIdx[i]\%numRows]$$

  4.对sortedExpertIdx的每个专家统计直方图结果，再进行Cumsum，得出expertTokensCountOrCumsumOut：

  $$expertTokensCountOrCumsumOut[i]=Cumsum(Histogram(sortedExpertIdx))$$

  5.对sortedExpertIdx的每个专家统计直方图结果，得出expertTokensBeforeCapacityOut：

  $$expertTokensBeforeCapacityOut[i]=Histogram(sortedExpertIdx)$$

函数原型

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

• aclnnStatus aclnnMoeInitRoutingV2GetWorkspaceSize(const aclTensor *x, const aclTensor *expertIdx, int64_t activeNum, int64_t expertCapacity, int64_t expertNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag, bool expertTokensBeforeCapacityFlag, aclTensor *expandedXOut, aclTensor *expandedRowIdxOut, aclTensor *expertTokensCountOrCumsumOut, aclTensor *expertTokensBeforeCapacityOut, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnMoeInitRoutingV2(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnMoeInitRoutingV2GetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：为MOE的输入，即token特征输入，要求为一个2D的Tensor，shape为[numRows, h]，numRows代表Token个数，h代表每个Token的长度，数据类型支持FLOAT16、BFLOAT16、FLOAT32，undefined要求为ND，支持undefined。
  • expertIdx （aclTensor*，计算输入）：为每个Token对应的k个处理专家的序号，一般为aclnnMoeGatingTopKSoftmaxV2接口的输出。undefined要求为ND，支持undefined。在Drop/Pad场景下或者非Drop/Pad场景下且需要输出expertTokensCountOrCumsumOut时，要求值域范围是[0, expertNum - 1]， 其他场景要求大于等于0。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持INT32，要求是一个2D的shape [numRows, k]。
  • activeNum（int64_t，计算输入）：表示是否为Active场景，该属性在dropPadMode为0时生效，值范围大于等于0；0表示Dropless场景，大于0时表示Active场景，约束所有专家共同处理tokens总量
  • expertCapacity（int64_t， 计算输入）：表示每个专家能够处理的tokens数，值范围大于等于0；Drop/Pad场景下值域范围(0, numRows]，此时各专家将超过capacity的tokens drop掉，不够capacity阈值时则pad全0 tokens；其他场景不关心该属性值。
  • expertNum（int64_t， 计算输入）：表示专家数，值范围大于等于0；Drop/Pad场景下或者expertTokensCountOrCumsumFlag大于0需要输出expertTokensCountOrCumsumOut时，expertNum需大于0。
  • dropPadMode（int64_t， 计算输入）：表示是否为Drop/Pad场景，取值为0和1。
    • 0：表示非Drop/Pad场景，该场景下不校验expertCapacity。
    • 1：表示Drop/Pad场景，需要校验expertNum和expertCapacity，对于每个专家处理的超过和不足expertCapacity的值会做相应的处理。
  • expertTokensCountOrCumsumFlag（int64_t， 计算输入）：取值为0、1和2。
    • 0：表示不输出expertTokensCountOrCumsumOut。
    • 1：表示输出的值为各个专家处理的token数量的累计值。
    • 2：表示输出的值为各个专家处理的token数量。
  • expertTokensBeforeCapacityFlag（bool，计算输入）：取值为false和true。
    • false：表示不输出expertTokensBeforeCapacityOut。
    • true：表示输出的值为在drop之前各个专家处理的token数量。
  • expandedXOut（aclTensor*，计算输出）：根据expertIdx进行扩展过的特征，在Dropless/Active场景下要求是一个2D的Tensor，Dropless场景shape为[numRows * k, h]，Active场景shape为[min(activeNum, numRows * k), h]，在Drop/Pad场景下要求是一个3D的Tensor，shape为[expertNum, expertCapacity, h]。数据类型同x，支持FLOAT16、BFLOAT16、FLOAT32，undefined要求为ND，不支持undefined。
  • expandedRowIdxOut（aclTensor*，计算输出）：expandedXOut和x的索引映射关系， 要求是一个1D的Tensor，Shape为[numRows*k, ]，数据类型支持int32，undefined要求为ND，不支持undefined。
  • expertTokensCountOrCumsumOut（aclTensor*，计算输出）：输出每个专家处理的token数量的统计结果及累加值，通过expertTokensCountOrCumsumFlag参数控制是否输出，该值仅在非Drop/Pad场景下输出，要求是一个1D的Tensor，Shape为[expertNum, ]，数据类型支持int32，undefined要求为ND，不支持undefined。
  • expertTokensBeforeCapacityOut（aclTensor*，计算输出）：输出drop之前每个专家处理的token数量的统计结果，通过expertTokensBeforeCapacityFlag参数控制是否输出，该值仅在Drop/Pad场景下输出，要求是一个1D的Tensor，Shape为[expertNum, ]，数据类型支持int32，undefined要求为ND，不支持undefined。
  • workspaceSize（uint64_t*，出参）：返回用户需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值

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

  [object Object]

aclnnMoeInitRoutingV2

• 参数说明：

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

• 返回值：

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

约束说明

• 输入x都必须为2维，且x的numRows等于expertIdx的numRows。
• dropPadMode为1时，expertNum和expertCapacity必须大于0。
• x的numRows轴必须大于expertCapacity。

调用示例

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