MoeDistributeCombineV3

产品支持情况

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

功能说明

• 算子功能：当存在TP域通信时，先进行ReduceScatterV通信，再进行AllToAllV通信，最后将接收的数据整合（乘权重再相加）；当不存在TP域通信时，进行AllToAllV通信，最后将接收的数据整合（乘权重再相加）。

  相较于MoeDistributeCombineV2算子，该算子变更如下：

  • 新增context入参，存入通信域相关信息；
  • 新增ccl_buffer_size入参，指定当前通信域大小；
  • 减少group_ep以及group_tp通信域名称入参; 详细说明请参考以下参数说明。

• 计算公式：

  • 不存在TP域通信时：

  $$\begin{gathered} ataOut=AllToAllV(expan{d}_x) \\ {x}_{o}ut=Sum(exper{t}_{s}cales\ast ataOut+exper{t}_{s}cales\ast sharedExpertX) \end{gathered}$$

  • 存在TP域通信时：

  $$\begin{gathered} rsOut=ReduceScatterV(expan{d}_x) \\ ataOut=AllToAllV(rsOut) \\ {x}_{o}ut=Sum(exper{t}_{s}cales\ast ataOut+exper{t}_{s}cales\ast sharedExpertX) \end{gathered}$$

  注意该算子必须与MoeDistributeDispatchV3配套使用，相当于按MoeDistributeDispatchV3算子收集数据的路径原路返还。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
context	输入	本卡通信域信息数据。	FLOAT16、BFLOAT16	ND
expand_x	输入	根据expertIds进行扩展过的token特征。	FLOAT16、BFLOAT16	ND
expert_ids	输入	每个token的topK个专家索引。	INT32	ND
assist_info_for_combine	输入	表示同一专家收到的token个数。	INT32	ND
ep_send_counts	输入	从EP通信域各卡接收的token数。	INT32	ND
expert_scales	输入	每个token的topK个专家的权重。	FLOAT32	ND
tp_send_counts_optional	可选输入	从TP通信域各卡接收的token数，对应MoeDistributeDispatchV3中的tp_recv_counts输出，有TP域通信需传参，无TP域通信传空指针。	INT32	ND
x_active_mask_optional	可选输入	表示token是否参与通信，可传有效数据或空指针；1D时true需排在false前（例：{true, false, true}非法），2D时token对应K个值全为false则不参与通信；默认所有token参与通信；各卡BS不一致时所有token需有效。	BOOL	ND
activation_scale_optional	可选输入	预留参数，当前版本不支持，传空指针即可。	-	ND
weight_scale_optional	可选输入	预留参数，当前版本不支持，传空指针即可。	-	ND
group_list_optional	可选输入	预留参数，当前版本不支持，传空指针即可。	-	ND
expand_scales_optional	可选输入	表示本卡token的权重，对应MoeDistributeDispatchV3中的expand_scales输出。	FLOAT32	ND
shared_expert_x_optional	可选输入	表示共享专家计算后的token，数据类型需与expand_x一致。	FLOAT16、BFLOAT16	ND
elastic_info_optional	可选输入	EP通信域动态缩容信息。	INT32	ND
ori_x_optional	可选输入	表示未经过FFN（Feed-Forward Neural network）的token数据，在使能copyExpert或使能constExpert的场景下需要本输入数据。	FLOAT16、BFLOAT16	ND
const_expert_alpha_1_optional	可选输入	在使能constExpert的场景下需要输入的计算系数。	FLOAT16、BFLOAT16	ND
const_expert_alpha_2_optional	可选输入	在使能constExpert的场景下需要输入的计算系数。	FLOAT16、BFLOAT16	ND
const_expert_v_optional	可选输入	在使能constExpert的场景下需要输入的计算系数。	FLOAT16、BFLOAT16	ND
performance_info_optional	可选输入	表示本卡等待各卡数据的通信时间，单位为us（微秒）。单次算子调用各卡通信耗时会累加到该Tensor上，算子内部不进行自动清零，因此用户每次启用此Tensor开始记录耗时前需对Tensor清零。	INT64	ND
ep_world_size	属性	EP通信域大小。	INT64	ND
ep_rank_id	属性	EP域本卡Id，取值范围[0, ep_world_size)，同一个EP通信域中各卡的ep_rank_id不重复。	INT64	ND
moe_expert_num	属性	MoE专家数量，满足moe_expert_num % (ep_world_size - shared_expert_num) = 0。	INT64	ND
ccl_buffer_size	属性	当前通信域Buffer大小。 默认值为""。	STRING	ND
tp_world_size	可选属性	TP通信域大小。 默认值为0。	INT64	ND
tp_rank_id	可选属性	TP域本卡Id，无TP域通信时传0即可。 默认值为0。	INT64	ND
expert_shard_type	可选属性	表示共享专家卡分布类型，当前仅支持传0，表示共享专家卡排在MoE专家卡前面。 默认值为0。	INT64	ND
shared_expert_num	可选属性	表示共享专家数量（一个共享专家可复制部署到多个卡上）。 默认值为1。	INT64	ND
shared_expert_rank_num	可选属性	表示共享专家卡数量，取值范围[0, ep_world_size)；为0时需满足shared_expert_num为0或1，不为0时需满足shared_expert_rank_num % shared_expert_num = 0。 默认值为0。	INT64	ND
global_bs	可选属性	EP域全局的batch size大小；各rank BS一致时，global_bs = BS * ep_world_size 或 0；各rank BS不一致时，global_bs = max_bs * ep_world_size（max_bs为单卡Bs最大值）。 默认值为0。	INT64	ND
out_dtype	可选属性	用于指定输出x的数据类型，预留参数，当前版本不支持，传0即可。 默认值为0。	INT64	ND
comm_quant_mode	可选属性	通信量化类型，取值范围0或2；0表示通信不量化，2表示通信int8量化。 默认值为0。	INT64	ND
group_list_type	可选属性	group_list格式，预留参数，当前版本不支持，传0即可。 默认值为0。	INT64	ND
comm_alg	可选属性	表示通信亲和内存布局算法。 默认值为""。	STRING	ND
zero_expert_num	可选属性	零专家数量。 默认值为0。	INT64	ND
copy_expert_num	可选属性	copy专家数量。 默认值为0。	INT64	ND
const_expert_num	可选属性	常量专家数量。 默认值为0。	INT64	ND
x_out	输出	表示处理后的token，数据类型、数据格式与expand_x保持一致。	FLOAT16、BFLOAT16	ND

• Atlas A3 训练系列产品/Atlas A3 推理系列产品：
  • 不支持expand_scales_optional。
  • 不支持comm_alg。

约束说明

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

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

• 调用算子过程中使用的ep_world_size、moe_expert_num、ccl_buffer_size、tp_world_size、expert_shard_type、shared_expert_num、shared_expert_num、global_bs、comm_alg参数取值所有卡需保持一致，网络中不同层中也需保持一致，且和MoeDistributeDispatchV3算子对应参数也保持一致。

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

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

• 参数约束：

  • zero_expert_num：取值范围：[0, MAX_INT32)，MAX_INT32 = 2^31 - 1, 合法的零专家的ID的值是[moe_expert_num, moe_expert_num + zero_expert_num)。
  • copy_expert_num：取值范围：[0, MAX_INT32)，MAX_INT32 = 2^31 - 1, 合法的copy专家的ID的值是[moe_expert_num + zero_expert_num, moe_expert_num + zero_expert_num + copy_expert_num)。
  • const_expert_num：取值范围：[0, MAX_INT32)，MAX_INT32 = 2^31 - 1, 合法的常量专家的ID的值是[moe_expert_num + zero_expert_num + copy_expert_num, moe_expert_num + zero_expert_num + copy_expert_num + const_expert_num)。
  • ori_x_optional：可选择传入有效数据或填空指针，当copy_expert_num不为0或const_expert_num不为0时必须传入有效输入；当传入有效数据时，要求shape为 (BS, H)，数据类型需与expand_x保持一致。
  • const_expert_alpha_1_optional：可选择传入有效数据或填空指针，当const_expert_num不为0时必须传入有效输入；当传入有效数据时，要求shape为(const_expert_num, )，数据类型需与expand_x保持一致。
  • const_expert_alpha_2_optional：可选择传入有效数据或填空指针，当const_expert_num不为0时必须传入有效输入；当传入有效数据时，要求shape为(const_expert_num, )，数据类型需与expand_x保持一致。
  • const_expert_v_optional：可选择传入有效数据或填空指针，当const_expert_num不为0时必须传入有效输入；当传入有效数据时，要求shape为(const_expert_num, H)，数据类型需与expand_x保持一致。

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

• 通信域使用约束：

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

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

  • 该场景下单卡包含双DIE（简称为“晶粒”或“裸片”），因此参数说明里的“本卡”均表示单DIE。
  • 参数说明里shape格式说明：
    • H：表示hidden size隐藏层大小，取值范围[1024, 8192]。
    • BS：表示batch sequence size，即本卡最终输出的token数量，取值范围为[1, 512]。
  • 参数约束：
    • ep_world_size：取值支持8、16、32、64、128、144、256、288。
    • moe_expert_num：取值范围(0, 1024]。
    • tp_world_size：取值范围[0, 2]，0和1表示无TP域通信，有TP域通信时仅支持2。
    • tp_rank_id：取值范围[0, 1]，同一个TP通信域中各卡的tp_rank_id不重复。无TP域通信时，传0即可。
    • shared_expert_num：当前取值范围[0, 4]。
    • comm_quant_mode：int8量化当且仅当tp_world_size < 2时可使能。
    • performance_info_optional：预留参数，当前版本不支持，传空指针即可。
    • ccl_buffer_size：调用get_low_latency_ccl_buffer_size接口(../../torch_extension/npu_ops_transformer/ops/deep_ep.py)。

调用说明

调用方式	样例代码	说明
npu_low_latency_combine接口	deepep.py	通过npu_low_latency_combine接口方式调用moe_distribute_combine_v3算子。