MoeInitRoutingV3

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：MoE的routing计算，根据aclnnMoeGatingTopKSoftmaxV2的计算结果做routing处理，支持不量化和动态量化模式。本接口针对V2接口aclnnMoeInitRoutingV2做了如下功能变更，请根据实际情况选择合适的接口：

  1.增加动态量化功能，支持输出expendX的 int8动态量化输出。

  2.增加参数activeExpertRangeOptional，支持筛选有效范围内的expertId。

  3.删除属性expertTokensBeforeCapacityFlag、删除输出expertTokensBeforeCapacityOut（使用expertTokensCountOrCumsumOut进行输出）。

• 计算公式：

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

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

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

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

  3.在drop模式下，对sortedExpertIdx的每个专家统计直方图结果，得出expertTokensCountOrCumsumOutOptional：

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

  4.计算quant结果：

  • 动态quant：

    • 若不输入scale：

      $$dynamicQuantScaleOutOptional=row\_max(abs(x))/127$$

      $$quantResult=round(x/dynamicQuantScaleOutOptional)$$

    • 若输入scale:

      $$dynamicQuantScaleOutOptional=row\_max(abs(x\ast scaleOptional))/127$$

      $$quantResult=round(x/dynamicQuantScaleOutOptional)$$

  5.对quantResult取前NUM_ROWS个sortedRowIdx的对应位置的值，得出expandedXOut：

  $$expandedXOut[i]=quantResult[sortedRowIdx[i]\%NUM\_ROWS]$$

  6.expandedRowIdxOut的有效元素数量availableIdxNum计算方式为，expertIdx中activeExpertRangeOptional范围内的元素的个数

  $$availableIdxNum=|\{x\in expertIdx|expert\_start\le x<expert\_end\}|$$

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	MOE的输入，即token特征输入，对应公式中x。	FLOAT32、FLOAT16、BFLOAT16、INT8、HIFLOAT8、FLOAT4_E2M1、FLOAT8_E4M3FN、FLOAT8_E5M2。	ND
expertIdx	输入	每一行特征对应的K个处理专家，里面元素专家id不能超过专家数。对应公式中 expertIdx。	INT32	ND
scaleOptional	可选输入	表示用于计算quant结果的参数。如果不输入表示计算时不使用scale，对应公式中scale。 • 仅quantMode=-1且x的数据类型为FLOAT4_E2M1、FLOAT8_E4M3FN或FLOAT8_E5M2时，scale数据类型支持FLOAT8_E8M0。	FLOAT32、FLOAT8_E8M0	ND
offsetOptional	可选输入	表示用于计算quant结果的偏移值。在非量化场景下和动态quant场景下不输入，对应公式中offsetOptional。	FLOAT32	ND
activeNum	属性	表示总的最大处理row数，输出expandedXOut只有这么多行是有效的。	INT	-
expertCapacity	属性	表示每个专家能够处理的tokens数，取值范围大于等于0。	INT	-
expertNum	属性	表示专家数，expertTokensNumType为key\_value模式时，取值范围为[0, 5120], 其它模式取值范围[0, 10240]。	INT	-
dropPadMode	属性	表示是否为 DropPad 场景，取值为 0 和 1。 • 0：表示 Dropless 场景，该场景下不校验 expertCapacity。 • 1：表示 DropPad 场景。	INT	-
expertTokensNumType	属性	取值为0、1和2 。 • 0：表示 comsum 模式。 • 1：表示 count 模式，即输出的值为各个专家处理的 token 数量的累计值。 • 2：表示 key\_value 模式，即输出的值为专家和对应专家处理 token 数量的累计值。	INT	-
expertTokensNumFlag	属性	取值为false和true。 • false：表示不输出 expertTokensCountOrCumsumOut。 • true：表示输出 expertTokensCountOrCumsumOut。	Bool	-
quantMode	属性	取值为0、1、-1、2、3、6、7、8、9。 • 0：表示静态quant场景。 • 1：表示动态quant场景。 • -1：表示不量化场景。 • 2：表示MXFP8量化场景，expandedXOut量化到FLOAT8_E5M2。 • 3：表示MXFP8量化场景，expandedXOut量化到FLOAT8_E4M3FN。 • 6：表示HIF8直转量化场景下，expandedXOut量化到HIFLOAT8。 • 7：表示HIF8 PERTENSOR量化场景，expandedXOut量化到HIFLOAT8。 • 8：表示HIF8 PERTOKEN量化场景，expandedXOut量化到HIFLOAT8。 • 9：表示MXFP4量化场景，expandedXOut量化到FLOAT4_E2M1。	INT	-
activeExpertRangeOptional	可选属性	长度为2，数组内的值为[expertStart, expertEnd], 表示活跃的expert范围在expertStart和expertEnd之间，左闭右开。要求值大于等于0，并且expertEnd不大于expertNum。	ListInt	-
rowIdxType	属性	表示expandedRowIdxOut使用的索引类型，取值为0、1。（性能模板仅支持1） • 0：表示gather类型的索引。 • 1：表示scatter类型的索引。	INT	-
expandedXOut	输出	根据expertIdx进行扩展过的特征。非量化场景下数据类型同x，量化场景quantMode为0、1时数据类型支持INT8，quantMode为2、3时数据类型分别支持FLOAT8_E5M2、FLOAT8_E4M3FN，quantMode为6、7、8时数据类型支持HIFLOAT8，quantMode为9时数据类型支持FLOAT4_E2M1。	FLOAT32、FLOAT16、BFLOAT16、INT8、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8、FLOAT4_E2M1	ND
expandedRowIdxOut	输出	expandedXOut和x的索引映射关系，前availableIdxNum\*H个元素为有效数据，其余无效数据，当rowIdxType为0时，无效数据由-1填充；当rowIdxType为1时，无效数据未初始化。	INT32	ND
expertTokensCountOrCumsumOut	输出	• 在expertTokensNumType为1的场景下，表示activeExpertRangeOptional范围内expert对应的处理token的总数。 • 在expertTokensNumType为2的场景下，表示activeExpertRangeOptional范围内token总数为非0的expert，以及对应expert处理token的总数。	INT64	ND
expandedScaleOut	输出	输出量化计算过程中scaleOptional的中间值。 • 非量化场景下为可选输入, 如果输入则要求为1D的Tensor, 类型为FLOAT32。当输入x数据类型为FLOAT4_E2M1、FLOAT8_E4M3FN或FLOAT8_E5M2时, 如果输入则要求3D的Tensor, 类型为FLOAT8_E8M0。 • quantMode为2、3、9时, 数据类型支持FLOAT8_E8M0。 • 其余场景数据类型支持FLOAT32。	FLOAT32、FLOAT8_E8M0	ND

约束说明

• 输入值域限制：

  • activeNum 当前未使用，校验需等于NUM_ROWS*K。
  • expertCapacity 当前未使用，仅校验非空。
  • dropPadMode 当前只支持0，代表 Dropless 场景。
  • expertTokensNumType 当前只支持 1 和 2，分别代表 count 模式和 key_value 模式。
  • expertTokensNumFlag 只支持 true，代表输出 expertTokensCountOrCumsumOut。
  • quantMode:
    • Atlas A2 训练系列产品/Atlas A2 推理系列产品、Atlas A3 训练系列产品/Atlas A3 推理系列产品：支持1、-1，分别代表动态量化场景和不量化场景。
    • Ascend 950PR/Ascend 950DT：支持-1、1、2、3、6、7、8、9，分别表示不量化、动态量化、MXFP8量化到FLOAT8_E5M2、MXFP8量化到FLOAT8_E4M3FN、按直转方式量化到HIFLOAT8、按PERTENSOR模式量化到HIFLOAT8、按PERTOKEN模式量化到HIFLOAT8，MXFP4量化到FLOAT4_E2M1。

• 其他限制：该算子部分产品支持两种性能模板，进入两种性能模板需要分别额外满足以下条件，不满足条件则进入通用模板。

  • 支持性能模板的产品：

    • Atlas A2 训练系列产品/Atlas A2 推理系列产品。
    • Atlas A3 训练系列产品/Atlas A3 推理系列产品。

  • 进入低时延性能模板需要同时满足以下条件：

    • x、expertIdx、scaleOptional 输入 Shape 要求分别为：(1, 7168)、(1, 8)、(256, 7168)。
    • x 数据类型要求：BFLOAT16。
    • 属性要求：activeExpertRangeOptional=[0, 256]、 quantMode=1、expertTokensNumType=2、expertNum=256。

  • 进入大 batch 性能模板需要同时满足以下条件：

    • NUM_ROWS范围为[384, 8192]
    • K=8
    • expertNum=256
    • expertEnd-expertStart<=32
    • quantMode=-1
    • rowIdxType=1
    • expertTokensNumType=1