MoeDistributeDispatchV2

产品支持情况

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

功能说明

算子功能：对token数据进行量化（可选），当存在TP域通信时，先进行EP（Expert Parallelism）域的AllToAllV通信，再进行TP（Tensor Parallelism）域的AllGatherV通信；当不存在TP域通信时，进行EP（Expert Parallelism）域的AllToAllV通信。

• 情形1：如果quantMode=0（非量化场景）：

$$\begin{gathered} allToAllXOut=AllToAllV(X) \\ expandXOut=\{\begin{array}{ll}AllToAllV(X),& 无TP通信域\\ AllGatherV(allToAllXOut),& 有TP通信域\end{array} \end{gathered}$$

• 情形2：如果quantMode=1（静态量化场景）：

$$\begin{gathered} xFp32=CastToFp32(X)\times scales \\ quantOut=Cast(xFp32,dstType) \\ allToAllXOut=AllToAllV(quantOut) \\ expandXOut=\{\begin{array}{ll}AllToAllV(quantOut),& 无TP通信域\\ AllGatherV(allToAllXOut),& 有TP通信域\end{array} \end{gathered}$$

• 情形3：如果quantMode=2（pertoken动态量化场景）：

$$\begin{gathered} xFp32=CastToFp32(X)\times scales \\ dynamicScales=dstTypeMax/Max(Abs(xFp32)) \\ quantOut=CastToInt8(xFp32\times dynamicScales) \\ allToAllXOut=AllToAllV(quantOut) \\ allToAllDynamicScalesOut=AllToAllV(1.0/dynamicScales) \\ expandXOut=\{\begin{array}{ll}AllToAllV(quantOut),& 无TP通信域\\ AllGatherV(allToAllXOut),& 有TP通信域\end{array} \\ dynamicScalesOut=\{\begin{array}{ll}allToAllDynamicScalesOut,& 无TP通信域\\ AllGatherV(allToAllDynamicScalesOut),& 有TP通信域\end{array} \end{gathered}$$

• 情形4：如果quantMode=3（pertile动态量化场景）：

$$\begin{gathered} xFp32=CastToFp32(X)\times scales \\ dynamicScales=dstTypeMax/Max(Abs(xFp32)) \\ quantOut=CastToInt8(xFp32\times dynamicScales) \\ allToAllXOut=AllToAllV(quantOut) \\ allToAllDynamicScalesOut=AllToAllV(1.0/dynamicScales) \\ expandXOut=\{\begin{array}{ll}AllToAllV(quantOut),& 无TP通信域\\ AllGatherV(allToAllXOut),& 有TP通信域\end{array} \\ dynamicScalesOut=\{\begin{array}{ll}allToAllDynamicScalesOut,& 无TP通信域\\ AllGatherV(allToAllDynamicScalesOut),& 有TP通信域\end{array} \end{gathered}$$

• 情形5：如果quantMode=4（mx量化场景）：

$$\begin{gathered} sharedExp=Floor(lo{g}_2(max(x)))-emax \\ dynamicScales=2^{sharedExp} \\ quantOut=CastToFp8(X/dynamicScales) \\ allToAllXOut=AllToAllV(quantOut) \\ allToAllDynamicScalesOut=AllToAllV(1.0/dynamicScales) \\ expandXOut=\{\begin{array}{ll}AllToAllV(quantOut),& 无TP通信域\\ AllGatherV(allToAllXOut),& 有TP通信域\end{array} \\ dynamicScalesOut=\{\begin{array}{ll}allToAllDynamicScalesOut,& 无TP通信域\\ AllGatherV(allToAllDynamicScalesOut),& 有TP通信域\end{array} \end{gathered}$$

其中，$emax$表示该类型最大正规数对应的指数部分的值。

• Atlas A2 训练系列产品/Atlas A2 推理系列产品：该算子必须与MoeDistributeCombineV2配套使用。
• Atlas A3 训练系列产品/Atlas A3 推理系列产品/Ascend 950PR/Ascend 950DT：该算子必须与MoeDistributeCombineV2或MoeDistributeCombineAddRmsNorm配套使用。

说明：MoeDistributeCombineV2、MoeDistributeCombineAddRmsNorm算子在后续文档中统称为CombineV2系列算子。 ・ 相较于MoeDistributeDispatch算子，该算子变更如下：

• 输出了更详细的token信息辅助CombineV2系列算子高效地进行全卡同步，因此原算子中shape为(BS * K,)的expandIdx出参替换为shape为(A * 128,)的assistInfoForCombineOut参数；
• 新增commAlg入参，代替HCCL_INTRA_PCIE_ENABLE和HCCL_INTRA_ROCE_ENABLE环境变量。

详细说明请参考以下参数说明。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	本卡发送的token数据。	FLOAT16、BFLOAT16、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8、FLOAT4_E2M1、FLOAT4_E1M2	ND
expertIds	输入	每个token的topK个专家索引。	INT32	ND
scalesOptional	可选输入	每个专家的量化平滑参数。	FLOAT32、FLOAT8_E8M0	ND
xActiveMaskOptional	可选输入	表示token是否参与通信，可传有效数据或空指针；1D时true需排在false前（例：{true, false, true}非法），2D时token对应K个值全为false则不参与通信；默认所有token参与通信；各卡BS不一致时所有token需有效。	BOOL	ND
expertScalesOptional	可选输入	每个token的topK个专家权重。	FLOAT32	ND
elasticInfoOptional	可选输入	EP通信域动态缩容信息。当某些通信卡因异常而从通信域中剔除，实际参与通信的卡数可从本参数中获取。	FLOAT32	ND
performanceInfoOptional	可选输入	表示本卡等待各卡数据的通信时间，单位为us（微秒）。单次算子调用各卡通信耗时会累加到该Tensor上，算子内部不进行自动清零，因此用户每次启用此Tensor开始记录耗时前需对Tensor清零。	INT64	ND
groupEp	属性	EP通信域名称（专家并行通信域），字符串长度范围为[1, 128)，不能和groupTp相同。	STRING	ND
epWorldSize	属性	EP通信域大小。	INT64	ND
epRankId	属性	EP域本卡ID，取值范围[0, epWorldSize)，同一个EP通信域中各卡的epRankId不重复。	INT64	ND
moeExpertNum	属性	MoE专家数量，满足moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0。	INT64	ND
groupTp	可选属性	TP通信域名称（数据并行通信域）。 默认值为""。	STRING	ND
tpWorldSize	可选属性	TP通信域大小，取值范围[0, 2]，0和1表示无TP域通信，有TP域通信时仅支持2。 默认值为0。	INT64	ND
tpRankId	可选属性	TP域本卡ID，取值范围[0, 1]，同一个TP通信域中各卡的tpRankId不重复；无TP域通信时传0即可。 默认值为0。	INT64	ND
expertShardType	可选属性	表示共享专家卡分布类型，当前仅支持传0，表示共享专家卡排在MoE专家卡前面。 默认值为0。	INT64	ND
sharedExpertNum	可选属性	表示共享专家数量（一个共享专家可复制部署到多个卡上）。 默认值为1。	INT64	ND
sharedExpertRankNum	可选属性	表示共享专家卡数量，取值范围[0, epWorldSize)；为0时需满足sharedExpertNum为0或1，不为0时需满足sharedExpertRankNum % sharedExpertNum = 0。 默认值为0。	INT64	ND
quantMode	可选属性	表示量化模式。默认值为0。	INT64	ND
globalBS	可选属性	EP域全局的batch size大小；各rank BS一致时，globalBS = BS * epWorldSize 或 0；各rank BS不一致时，globalBS = maxBS * epWorldSize（maxBS为单卡BS最大值）。 默认值为0。	INT64	ND
expertTokenNumsType	可选属性	输出expertTokenNums中值的语义类型，支持0：expertTokenNums中的输出为每个专家处理的token数的前缀和，1：expertTokenNums中的输出为每个专家处理的token数量。 默认值为1。	INT64	ND
commAlg	可选属性	表示通信亲和内存布局算法。 默认值为""。	STRING	ND
zeroExpertNum	可选属性	零专家数量。 默认值为0。	INT64	ND
copyExpertNum	可选属性	copy专家数量。 默认值为0。	INT64	ND
constExpertNum	可选属性	常量专家数量。 默认值为0。	INT64	ND
expandXOut	输出	根据expertIds进行扩展过的token特征。	FLOAT16、BFLOAT16、INT8、FLOAT8_E4M3FN、FLOAT8_E5M2、HIFLOAT8、FLOAT4_E2M1、FLOAT4_E1M2	ND
dynamicScalesOut	输出	量化场景下，表示本卡输出Token的量化系数。	FLOAT32、FLOAT8_E8M0	ND
assistInfoForCombineOut	输出	表示给同一专家发送的token个数（对应CombineV2系列算子中的assistInfoForCombine）。	INT32	ND
expertTokenNumsOut	输出	表示每个专家收到的token个数。	INT64	ND
epRecvCountsOut	输出	从EP通信域各卡接收的token数（对应CombineV2系列算子中的epSendCounts）。	INT32	ND
tpRecvCountsOut	输出	从TP通信域各卡接收的token数（对应CombineV2系列算子中的tpSendCountsOptional），有TP域通信则有该输出，无TP域通信则无该输出。	INT32	ND
expandScalesOut	输出	表示本卡输出token的权重（对应CombineV2系列算子中的expandScalesOptional）。	FLOAT32	ND

• Atlas A2 训练系列产品/Atlas A2 推理系列产品：
  • 不支持共享专家场景，不支持expertShardType、sharedExpertNum、sharedExpertRankNum属性。
  • 仅支持EP域，无TP域，不支持groupTp、tpWorldSize、tpRankId属性，且tpRecvCounts输出无有效内容。
  • 不支持elasticInfoOptional。
  • 当commAlg = "hierarchy"，expandScalesOut内容有效。
  • 不支持常量专家场景，不支持constExpertNum，使用默认值即可。
• Atlas A3 训练系列产品/Atlas A3 推理系列产品：
  • commAlg 支持""，"fullmesh_v1"，"fullmesh_v2", "hierarchy"三种输入方式。""：默认值，不使能fullmesh_v2模板；"fullmesh_v1"：不使能fullmesh_v2模板；"fullmesh_v2"：使能fullmesh_v2模板，该模板仅支持tpWorldSize为1场景；"hierarchy": 使能跨超模板，该模板仅支持tpWorldSize为1、共享专家为0的场景，且不支持可变BS、二维mask、特殊专家、performanceInfo场景。
  • expertScalesOptional 当commAlg="hierarchy"场景时，要求为2D Tensor，shape为(BS, K)；当commAlg=""，"fullmesh_v1"，
  • epWorldSize 取值范围[2, 768]；当commAlg="hierarchy"场景时，取值范围为[16, 256]，且为16的整数倍。
  • moeExpertNum 取值范围(0, 1024]；当commAlg="hierarchy"场景时，取值范围为(0, 512]。
  • expandScalesOut 当commAlg="hierarchy"场景时，要求为1D Tensor，shape为(A,)；当commAlg=""，"fullmesh_v1"，"fullmesh_v2"场景时，暂不支持该输出。
• Ascend 950PR/Ascend 950DT：
  • 仅支持EP域，无TP域，不支持groupTp、tpWorldSize、tpRankId属性，且tpRecvCounts输出无有效内容。
  • 不支持expandScalesOut。

约束说明

• MoeDistributeDispatchV2与CombineV2系列算子必须配套使用，具体参考调用示例。

• 算子通信域各节点的驱动版本应当相同。

• 在不同产品型号、不同通信算法或不同版本中，MoeDistributeDispatchV2的Tensor输出assistInfoForCombineOut、epRecvCountsOut、tpRecvCountsOut、expandScalesOut中的元素值可能不同，使用时直接将上述Tensor传给CombineV2系列算子对应参数即可，模型其他业务逻辑不应对其存在依赖。

• 调用算子过程中使用的groupEp、epWorldSize、moeExpertNum、groupTp、tpWorldSize、expertShardType、sharedExpertNum、sharedExpertRankNum、globalBS、commAlg参数，HCCL_BUFFSIZE取值所有卡需保持一致，网络中不同层中也需保持一致，且和CombineV2系列算子对应参数也保持一致。

• 参数说明里shape格式说明：

  • A：表示本卡可能接收的最大token数量，取值范围如下：
    • 对于共享专家，要满足A = BS * epWorldSize * sharedExpertNum / sharedExpertRankNum。
    • 对于MoE专家，当globalBS为0时，要满足A >= BS * epWorldSize * min(localExpertNum, K)；当globalBS非0时，要满足A >= globalBS * min(localExpertNum, K)。
  • K：表示选取topK个专家，取值范围为0 < K ≤ 16同时满足0 < K ≤ moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum。
  • localExpertNum：表示本卡专家数量。
    • 对于共享专家卡，localExpertNum = 1
    • 对于MoE专家卡，localExpertNum = moeExpertNum / (epWorldSize - sharedExpertRankNum)，localExpertNum > 1时，不支持TP域通信。

• 属性约束：

  • zeroExpertNum：取值范围：[0, MAX_INT32)，MAX_INT32 = 2^31 - 1, 合法的零专家的ID的值是[moeExpertNum, moeExpertNum + zeroExpertNum)。
  • copyExpertNum：取值范围：[0, MAX_INT32)，MAX_INT32 = 2^31 - 1, 合法的copy专家的ID的值是[moeExpertNum + zeroExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum)。
  • constExpertNum：取值范围：[0, MAX_INT32)，MAX_INT32 = 2^31 - 1, 合法的常量专家的ID的值是[moeExpertNum + zeroExpertNum + copyExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum)。

• 本文公式中的"/"表示整除。

• 通信域使用约束：

  • 一个模型中的CombineV2系列算子和MoeDistributeDispatchV2仅支持相同EP通信域，且该通信域中不允许有其他算子。
  • 一个模型中的CombineV2系列算子和MoeDistributeDispatchV2仅支持相同TP通信域或都不支持TP通信域，有TP通信域时该通信域中不允许有其他算子。

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

  • 参数约束：
    • commAlg：当前版本支持nullptr， ""， "fullmesh"， "hierarchy"四种输入方式，若配置"hierarchy"，建议搭配搭配25.0.RC1.1及以上版本驱动使用。
      • nullptr和""：仅在此场景下，HCCL_INTRA_PCIE_ENABLE和HCCL_INTRA_ROCE_ENABLE配置生效。当HCCL_INTRA_PCIE_ENABLE=1&&HCCL_INTRA_ROCE_ENABLE=0时，调用"hierarchy"算法，否则调用"fullmesh"算法。不推荐使用该方式。
      • "fullmesh"：token数据直接通过RDMA方式发往topk个目标专家所在的卡。
      • "hierarchy"：token数据经过跨机、机内两次发送，仅不同server同号卡之间使用RDMA通信，server内使用HCCS通信。
    • epWorldSize：依commAlg取值，"fullmesh"支持2、3、4、5、6、7、8、16、32、64、128、192、256、384；"hierarchy"支持16、32、64。
    • moeExpertNum：依commAlg取值，"fullmesh"支持(0, 1024]，"hierarchy"支持(0, 512]。
    • epRecvCountsOut：要求shape为 (moeExpertNum + 2 * globalBS * K * serverNum, )，前moeExpertNum个数表示从EP通信域各卡接收的token数，2 * globalBS * K * serverNum存储了机间机内做通信前combine可以提前做reduce的token个数和token在通信区中的偏移，globalBS传入0时在此处应当按照BS * epWorldSize计算。
    • performanceInfoOptional：可选择传入有效数据或填空指针，传入空指针时表示不使能记录通信耗时功能；当传入有效数据时，要求是一个1D的Tensor，shape为(ep_world_size,)，数据类型支持int64；数据格式要求为ND。
  • HCCL_INTRA_PCIE_ENABLE和HCCL_INTRA_ROCE_ENABLE：不推荐使用该环境变量控制通信算法，原HCCL_INTRA_PCIE_ENABLE=1&&HCCL_INTRA_ROCE_ENABLE=0场景，下文均通过commAlg = "hierarchy"替代，默认场景使用commAlg = "fullmesh"替代。
  • commAlg配置"hierarchy"时，不支持scalesOptional、xActiveMaskOptional、oriXOptional、zeroExpertNum、copyExpertNum。
  • quantMode支持0（非量化）、2（pertoken动态量化）。
  • 参数说明里shape格式说明：
    • H：表示hidden size隐藏层大小，取值范围(0, 10240]且为32的整数倍。
    • BS：表示batch sequence size，即本卡最终输出的token数量。
      • commAlg = "fullmesh"：取值范围(0, 256]。
      • commAlg = "hierarchy"：取值范围(0, 512]。
  • HCCL_BUFFSIZE：调用本算子前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB。
    • commAlg = "fullmesh"：要求 >= (BS * epWorldSize * min(localExpertNum, K) * H * 4B + 4MB)。
    • commAlg = "hierarchy"：要求 >= (moeExpertNum + epWorldSize / 4) * Align512(maxBS * (H * 2 + 16 * Align8(K))) * 1B + 8MB，其中Align8(x) = ((x + 8 - 1) / 8) * 8，Align512(x) = ((x + 512 - 1) / 512) * 512。
  • 组网约束：多机场景仅支持交换机组网，不支持双机直连组网。

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

  • 该场景下单卡包含双DIE（简称为“晶粒”或“裸片”），因此参数说明里的“本卡”均表示单DIE。
  • 参数约束：
    • elasticInfoOptional：当前版本不支持，传空指针即可。
    • epWorldSize：取值范围[2, 768]。
    • moeExpertNum：取值范围(0, 1024]。
    • groupTp：字符串长度范围为[1, 128)，不能和groupEp相同。
    • sharedExpertNum：取值支持[0, 4]。
    • commAlg：当前版本仅支持""，"fullmesh_v1"，"fullmesh_v2"，"hierarchy"三种输入方式。
      • ""：默认值，使能fullmesh_v1模板。
      • "fullmesh_v1"：使能fullmesh_v1模板。
      • "fullmesh_v2"：使能fullmesh_v2模板，其中commAlg仅在tpWorldSize取值为1时生效，且不支持在各卡BS不一致、输入xActiveMask和特殊专家场景下使能。
      • "hierarchy": 使能ROCE分层直驱能力，需要根据不同的逻辑超节点设置环境变量HCCL_LOGIC_SUPERPOD_ID，例如两机分别设为export HCCL_LOGIC_SUPERPOD_ID=0和export HCCL_LOGIC_SUPERPOD_ID=1。
    • epRecvCountsOut：要求shape为 (epWorldSize * max(tpWorldSize, 1) * localExpertNum, )。
    • performanceInfoOptional：预留参数，当前版本不支持，传空指针即可。
  • 参数说明里shape格式说明：
    • H：表示hidden size隐藏层大小，取值范围[1024, 8192]。
    • BS：表示batch sequence size，即本卡最终输出的token数量，取值范围为[1, 512]。
  • HCCL_BUFFSIZE：调用本算子前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB。要求 >= 2且满足>= 2 * (localExpertNum * maxBS * epWorldSize * Align512(Align32(2 * H) + 64) + (K + sharedExpertNum) * maxBS * Align512(2 * H))，localExpertNum需使用MoE专家卡的本卡专家数，其中Align512(x) = ((x + 512 - 1) / 512) * 512，Align32(x) = ((x + 32 - 1) / 32) * 32。

• Ascend 950PR/Ascend 950DT：

  • 参数约束：
    • elasticInfoOptional：当前版本不支持，传空指针即可。
    • epWorldSize：取值范围[2, 768]。
    • moeExpertNum：取值范围(0, 1024]。
    • sharedExpertNum：取值支持[0, 4]。
    • commAlg：当前版本仅支持""，"fullmesh_v1"，"fullmesh_v2"三种输入方式。
      • ""：默认值，使能fullmesh_v1模板。
      • "fullmesh_v1"：使能fullmesh_v1模板。
      • "fullmesh_v2"：使能fullmesh_v2模板，其中commAlg仅在tpWorldSize取值为1时生效，且不支持在各卡BS不一致、输入xActiveMask和特殊专家场景下使能。
    • epRecvCountsOut：要求shape为 (epWorldSize * max(tpWorldSize, 1) * localExpertNum, )。
    • performanceInfoOptional：预留参数，当前版本不支持，传空指针即可。
    • expertShardType当前仅支持传0，表示共享专家卡排在MoE专家卡前面。
    • quantMode支持0（非量化）、1（静态量化）、2（pertoken动态量化）、3（pergroup动态量化）、4（mx动态量化）。
  • 参数说明里shape格式说明：
    • H：表示hidden size隐藏层大小，取值范围[1024, 8192]。
    • BS：表示batch sequence size，即本卡最终输出的token数量，依commAlg取值，"fullmesh_v2"和"hierarchy"取值范围为 (0 < BS ≤ 256), "fullmesh_v1"和""取值范围为 (0 < BS ≤ 512)。
  • HCCL_BUFFSIZE：调用本算子前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB。要求 >= 2且满足>= 2 * (localExpertNum * maxBS * epWorldSize * Align512(Align32(2 * H) + 64) + (K + sharedExpertNum) * maxBS * Align512(2 * H))，localExpertNum需使用MoE专家卡的本卡专家数，其中Align512(x) = ((x + 512 - 1) / 512) * 512，Align32(x) = ((x + 32 - 1) / 32) * 32。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_moe_distribute_dispatch_v2.cpp	通过aclnnMoeDistributeDispatchV2接口方式调用moe_distribute_dispatch_v2算子。