CANN ops-transformer 算子解析：MoeDistributeDispatchV3 的 EP 域 Token 分发与动态量化通信实现

CANN ops-transformer 算子解析MoeDistributeDispatchV3 的 EP 域 Token 分发与动态量化通信实现【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本文基于 CANN ops-transformer 仓库中mc2/moe_distribute_dispatch_v3算子的官方文档与源码系统讲解 MoEMixture of Experts并行推理与训练场景下MoeDistributeDispatchV3算子如何完成Token 量化可选 EPExpert Parallelism专家并行域 AllToAllV 通信分发这一关键链路并给出完整的参数说明、约束条件、low_latency_dispatch高阶封装接口的使用方式与可复现的调用示例。读完本文你将掌握该算子在 Atlas A3 与 Ascend 950 系列产品上的适用场景、context/ccl_buffer_size等新增入参的含义以及它与MoeDistributeCombineV3系列算子配套使用的正确姿势。产品支持情况MoeDistributeDispatchV3在当前仓库中的产品适配情况如下表所示产品是否支持Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从算子定义文件 moe_distribute_dispatch_v3_def.cpp 可以看到该算子为两类硬件平台分别注册了独立的 AICore 内核入口ascend910_93Atlas A3对应arch22/..._a3.cppascend950Ascend 950对应arch35/..._apt.cpp与上表的产品支持矩阵一一对应。功能说明MoeDistributeDispatchV3的核心职责是对 token 数据进行量化可选随后执行 EP 域的 AllToAllV 通信把本卡持有的 token 按照expert_ids指定的 TopK 专家索引分发到对应的专家所在卡。TPTensor Parallelism张量并行域通信在当前版本不再支持tp_world_size、tp_rank_id等 TP 相关参数为预留参数。情形 1非量化场景quant_mode0直接对输入X执行 EP 域的 AllToAllV 通信通信结果同时作为allToAllXOut与expand_x_out输出$$ allToAllXOut AllToAllV(X) $$$$ expand_x_out AllToAllV(X) $$情形 2pertoken 动态量化场景quant_mode2先对 token 做逐 tokenpertoken动态量化再执行 AllToAllV 通信量化系数同样参与通信供下游反量化使用$$ xFp32 CastToFp32(X) \times scales $$$$ dynamicScales \frac{dstTypeMax}{Max(Abs(xFp32))} $$$$ quantOut CastToInt8(xFp32 \times dynamicScales) $$$$ allToAllXOut AllToAllV(quantOut) $$$$ allToAllDynamicScalesOut AllToAllV(1.0 / dynamicScales) $$$$ expand_x_out AllToAllV(quantOut) $$$$ dynamic_scales_out allToAllDynamicScalesOut $$其中dstTypeMax表示量化目标类型如 INT8所能表示的最大值scales为每个专家的量化平滑参数可选输入。值得注意的是emax表示该类型最大正规数对应的指数部分的值在 mx 量化quant_mode4/5等更细粒度的量化公式中会用到详见后文low_latency_dispatch接口的量化模式说明。与 CombineV3 系列算子的配套关系在 Atlas A3 产品上MoeDistributeDispatchV3必须与MoeDistributeCombineV3一起使用二者分别完成 MoE 并行中的分发dispatch与回收combine两个阶段。相较于此前的MoeDistributeDispatch算子MoeDistributeDispatchV3的变更点如下新增context入参存放本卡通信域相关信息由MoeDistributeBuffer封装类在构造时统一创建新增ccl_buffer_size入参指定当前通信域 Buffer 大小需调用get_low_latency_ccl_buffer_size接口计算减少group_ep与group_tp通信域名称入参不再需要显式传入通信域名称字符串。从原型定义 moe_distribute_dispatch_v3_proto.h 可以看到context被声明为DT_INT32类型的必选输入ep_world_size、ep_rank_id、moe_expert_num、ccl_buffer_size为四个REQUIRED_ATTR必选属性其余参数均为带默认值的可选属性.ATTR(...)。参数说明下表完整列出MoeDistributeDispatchV3的全部参数输入、属性、输出数据格式统一为 ND参数名输入/输出/属性描述数据类型数据格式context输入本卡通信域信息数据FLOAT16、BFLOAT16NDx输入本卡发送的 token 数据FLOAT16、BFLOAT16NDexpert_ids输入每个 token 的 topK 个专家索引INT32NDscales_optional可选输入每个专家的量化平滑参数非量化场景传空指针动态量化可传有效数据或空指针FLOAT32NDx_active_mask_optional可选输入表示 token 是否参与通信可传有效数据或空指针1D 时 true 需排在 false 前例{true, false, true}非法2D 时 token 对应 K 个值全为 false 则不参与通信默认所有 token 参与通信各卡 BS 不一致时所有 token 需有效BOOLNDexpert_scales_optional可选输入每个 token 的 topK 个专家权重FLOAT32NDelastic_info_optional可选输入EP 通信域动态缩容信息FLOAT32NDperformance_info_optional可选输入表示本卡等待各卡数据的通信时间单位为 us微秒。单次算子调用各卡通信耗时会累加到该 Tensor 上算子内部不进行自动清零因此用户每次启用此 Tensor 开始记录耗时前需对 Tensor 清零INT64NDep_world_size属性EP 通信域大小INT64NDep_rank_id属性EP 域本卡 Id取值范围 [0, ep_world_size)同一个 EP 通信域中各卡的 ep_rank_id 不重复INT64NDmoe_expert_num属性MoE 专家数量满足moe_expert_num % (ep_world_size - shared_expert_num) 0INT64NDccl_buffer_size属性当前通信域 Buffer 大小默认值为 STRINGNDtp_world_size可选属性TP 通信域大小预留参数当前版本不支持仅支持传 0 或 1默认值为 0INT64NDtp_rank_id可选属性TP 域本卡 Id预留参数当前版本不支持传 0 即可默认值为 0INT64NDexpert_shard_type可选属性表示共享专家卡分布类型当前仅支持传 0表示共享专家卡排在 MoE 专家卡前面默认值为 0INT64NDshared_expert_num可选属性表示共享专家数量一个共享专家可复制部署到多个卡上默认值为 1INT64NDshared_expert_rank_num可选属性表示共享专家卡数量取值范围 [0, ep_world_size)为 0 时需满足 shared_expert_num 为 0 或 1不为 0 时需满足shared_expert_rank_num % shared_expert_num 0默认值为 0INT64NDquant_mode可选属性表示量化模式支持 0非量化2动态量化默认值为 0INT64NDglobal_bs可选属性EP 域全局的 batch size 大小各 rank BS 一致时global_bs BS * ep_world_size或 0各 rank BS 不一致时global_bs max_bs * ep_world_sizemax_bs 为单卡 BS 最大值默认值为 0INT64NDexpert_token_nums_type可选属性输出 expert_token_nums 中值的语义类型支持 0输出为每个专家处理的 token 数的前缀和1输出为每个专家处理的 token 数量默认值为 1INT64NDcomm_alg可选属性表示通信亲和内存布局算法默认值为 STRINGNDzero_expert_num可选属性零专家数量默认值为 0INT64NDcopy_expert_num可选属性copy 专家数量默认值为 0INT64NDconst_expert_num可选属性常量专家数量默认值为 0INT64NDexpand_x_out输出根据 expert_ids 进行扩展过的 token 特征FLOAT16、BFLOAT16、INT8NDdynamic_scales_out输出量化场景下表示本卡输出 Token 的量化系数仅 quant_mode2 时有该输出FLOAT32NDassist_info_for_combine_out输出表示给同一专家发送的 token 个数INT32NDexpert_token_nums_out输出表示每个专家收到的 token 个数INT64NDep_recv_count_out输出从 EP 通信域各卡接收的 token 数INT32NDtp_recv_count_out输出从 TP 通信域各卡接收的 token 数预留输出当前版本不支持传空指针即可无 TP 域通信则无该输出INT32NDexpand_scales_out输出表示本卡输出 token 的权重FLOAT32ND产品差异说明在 Atlas A3 训练/推理系列产品上不支持expand_scales_out即不支持 TopK 专家权重随 token 一起通信。上述参数在算子原型 moe_distribute_dispatch_v3_proto.h 与宿主侧定义 moe_distribute_dispatch_v3_def.cpp 中均有完整声明例如context的 shape 约束为(2052,)、assist_info_for_combine的 shape 为(A * 128,)与 README 中的说明一致其中expand_x的数据类型列表含DT_HIFLOAT8、DT_FLOAT8_E5M2、DT_FLOAT8_E4M3FN、DT_FLOAT4_E2M1、DT_FLOAT4_E1M2表明该算子为低精度输出预留了能力具体行为受产品与y_dtype约束。约束说明配套使用约束MoeDistributeDispatchV3与 CombineV3 系列算子必须配套使用具体参考调用示例。在不同产品型号、不同通信算法或不同版本中Tensor 输出assist_info_for_combine_out、ep_recv_count_out、tp_recv_count_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_rank_num、global_bs、comm_alg参数取值所有卡需保持一致网络中不同层中也需保持一致且与 CombineV3 系列算子对应参数保持一致。Shape 格式说明A本卡可能接收的最大 token 数量取值范围如下对于共享专家需满足A BS * ep_world_size * shared_expert_num / shared_expert_rank_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_rank_num)且local_expert_num 1时不支持 TP 域通信。特殊专家属性约束zero_expert_num取值范围 [0, MAX_INT32)MAX_INT32 2^31 - 1合法的零专家的 ID 取值是 [moe_expert_num,moe_expert_numzero_expert_num)。copy_expert_num取值范围 [0, MAX_INT32)合法的 copy 专家的 ID 取值是 [moe_expert_numzero_expert_num,moe_expert_numzero_expert_numcopy_expert_num)。const_expert_num取值范围 [0, MAX_INT32)合法的常量专家的 ID 取值是 [moe_expert_numzero_expert_numcopy_expert_num,moe_expert_numzero_expert_numcopy_expert_numconst_expert_num)。通信域与通信方式约束一个模型中的 CombineV3 系列算子和MoeDistributeDispatchV3仅支持相同 EP 通信域且该通信域中不允许有其他算子。当前不支持 TP 域通信。Ascend 950DT 产品仅支持 UB Memory 通信。本文公式中的 / 表示整除。Atlas A3 产品的专属约束Atlas A3 场景下单卡包含双 DIE简称为晶粒或裸片因此参数说明里的本卡均表示单 DIE。该产品的参数约束如下elastic_info_optional当前版本不支持传空指针即可。ep_world_size取值范围 [2, 768]。moe_expert_num取值范围 (0, 1024]。shared_expert_num取值支持 [0, 4]。comm_alg当前版本仅支持、fullmesh_v1、fullmesh_v2三种输入方式。默认值开启 fullmesh_v1 模板fullmesh_v1开启 fullmesh_v1 模板fullmesh_v2开启 fullmesh_v2 模板其中comm_alg仅在tp_world_size取值为 1 时生效且不支持在各卡 BS 不一致、输入xActiveMask和特殊专家场景下开启。ep_recv_count_out要求 shape 为 (ep_world_size*local_expert_num,)。performance_info_optional预留参数当前版本不支持传空指针即可。ccl_buffer_size调用get_low_latency_ccl_buffer_size接口获取详见下文。Shape 格式说明H表示 hidden size 隐藏层大小取值范围 [1024, 8192]BS表示 batch sequence size即本卡最终输出的 token 数量取值范围为 [1, 512]。这些取值范围与 Python 封装 moe_distribute_buffer.py 中get_low_latency_ccl_buffer_size静态方法内的torch._check校验完全一致world_size ∈ [2, 768]、hidden ∈ [1024, 8192]、num_max_dispatch_tokens_per_rank ∈ [1, 512]、num_moe_expert ∈ [1, 1024]、topk ∈ [1, 16]、num_shared_expert ∈ [0, 4]说明 README 约束在接口层同样会被强制校验。调用说明官方推荐的调用方式是low_latency_dispatch接口该接口定义于 moe_distribute_buffer.py 的MoeDistributeBuffer类中完整接口文档见 torchapi_low_latency_dispatch.md。调用方式样例代码说明low_latency_dispatch 接口moe_distribute_buffer.py通过 low_latency_dispatch 接口方式调用 moe_distribute_dispatch_v3 算子通信域 context 与 Buffer 的构造MoeDistributeBuffer在构造时完成通信域初始化from cann_ops_transformer.ops import MoeDistributeBuffer distribute_buffer MoeDistributeBuffer(ep_group)构造过程中MoeDistributeBuffer通过CommContextManager后端按产品映射为kfc或channel创建context张量并从ccl_buffer_size参数创建通信 Buffer。该封装屏蔽了context、group_ep/group_tp通信域名称等底层细节这正是 V3 版本相对旧版算子新增 context、减少通信域名称入参这一变更在用户侧的体现。ccl_buffer_size 的计算方法MoeDistributeBuffer.get_low_latency_ccl_buffer_size是一个静态方法可依据模型规模直接算出当前通信域需要的 Buffer 大小单位 MBccl_buffer_size MoeDistributeBuffer.get_low_latency_ccl_buffer_size( world_sizeworld_size, num_max_dispatch_tokens_per_rankglobalBS, hiddenh, num_moe_expertnum_experts, topkk, num_shared_expertshared_expert_num, num_shared_expert_ranksshared_expert_rank_num, comm_algcomm_alg, )从源码看其内部按dispatch 方向 combine 方向双份内存进行估算minimum_buffer_size 2 * (dispatch 侧 token 数 * token_need_size_dispatch * world_size * local_moe_expert_num combine 侧 token 数 * token_need_size_combine * (topk shared_expert_num)) 1MB再按 32Bub_align、480Bfull_mesh_data_align仅 fullmesh_v2、512Bwin_addr_align等对齐规则向上取整后折算为 MB 并返回。因此ccl_buffer_size实际是依据模型配置动态计算的而不是 README 中 STRING 类型默认值所暗示的固定值——在使用low_latency_dispatch接口时该值会自动传入底层算子。low_latency_dispatch 函数原型MoeDistributeBuffer.low_latency_dispatch(x, topk_idx, num_experts, *, quant_mode0, comm_alg, x_smooth_scaleNone, x_active_maskNone, topk_weightsNone, zero_expert_num0, copy_expert_num0, const_expert_num0, elastic_infoNone, expert_shard_type0, shared_expert_num1, shared_expert_rank_num0, expert_token_nums_type1, num_max_dispatch_tokens_per_rank0, y_dtypeNone, x_dtypeNone, x_smooth_scales_dtypeNone) - (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)返回 6 个值expand_x、dynamic_scales、assist_info_for_combine、expert_token_nums、ep_recv_counts、expand_scales。其中assist_info_for_combine与ep_recv_counts需要原样传给low_latency_combine的assist_info_for_combine与ep_send_counts参数从 moe_distribute_buffer.py 中low_latency_combine的调用可见模型业务逻辑不应解析或依赖其中的具体元素值。该接口底层的算子调用为torch.ops.cann_ops_transformer.npu_moe_distribute_dispatch( contextself.context, xx, expert_idstopk_idx, ep_world_sizeself.world_size, ep_rank_idself.rank_id, moe_expert_numnum_experts, ccl_buffer_sizeself.ccl_buffer_size, ... global_bsnum_max_dispatch_tokens_per_rank * self.world_size, ... )即global_bs在接口层被自动换算为num_max_dispatch_tokens_per_rank * world_size与 README 中global_bs BS * ep_world_size的约定一致。Python 层的npu_moe_distribute_dispatch再经由 aclnn 单算子接口aclnnMoeDistributeDispatchV5下发到内核入口见 aclnn_moe_distribute_dispatch_v5.cpp其GetWorkspaceSize/执行两段式接口内部复用aclnnInnerMoeDistributeDispatchV3完成统一实现。接口级量化模式扩展quant_mode 0~5需要说明的是low_latency_dispatch接口在 README 描述的 quant_mode 0/2 基础上扩展支持了更丰富的量化模式详见 torchapi_low_latency_dispatch.mdquant_mode0非量化quant_mode1静态量化quant_out Cast(CastToFp32(x) × scales, dstType)quant_mode2pertoken 动态量化README 公式对应此模式quant_mode3pergroup 动态量化dynamic_scalesshape 为(A, Ceil(H, 128))quant_mode4mx 动态量化shared_exp floor(log2(max(x))) - emaxdynamic_scales_value 2^shared_expquant_mode5mx clip 动态量化先对max_abs做max(max_abs, 10^-4)clamp 再求shared_exp。其中 quant_mode 4/5 的公式使用了 README 中提到的emax该类型最大正规数对应的指数部分值。宿主侧 shape 推导 moe_distribute_dispatch_v3_infershape.cpp 中同样以QuantMode枚举形式定义了 0~5 六种模式并有PER_GROUP_SIZE 128、MX_QUANT_SIZE 32、ASSIST_INFO_NUM_PER_A 128等常量与接口文档的 shape 规则对应进一步印证了多量化模式在算子层面的落地。此外low_latency_dispatch还支持y_dtype/x_dtype/x_smooth_scales_dtype三个逻辑数据类型参数可在低精度场景hifloat8、float8_e5m2、float8_e4m3fn、float4_e2m1、float4_e1m2、float8_e8m0 等下显式指定输出与输入的 ACL 数据类型适合在 Ascend 950 上追求更低通信带宽的量化推理场景。通信 Buffer 大小与环境变量调用low_latency_dispatch前需检查HCCL_BUFFSIZE环境变量取值是否合理该变量表示单个通信域占用内存大小单位 MB不配置时默认为 200MB通信域开设大小可通过MoeDistributeBuffer.get_low_latency_ccl_buffer_size接口计算。EP 通信域内不同comm_alg的最小 Buffer 要求如下comm_alg为fullmesh_v1或要求 2 * (local_expert_num * max_bs * ep_world_size * Align512(Align32(2*H) 64) (K shared_expert_num) * max_bs * Align512(2*H))comm_alg为fullmesh_v2要求 2 * (local_expert_num * max_bs * ep_world_size * 480Align512(Align32(2*H) 64) (K shared_expert_num) * max_bs * Align512(2*H))其中480Align512(x) ((x480-1)/480)*512Align512(x) ((x512-1)/512)*512Align32(x) ((x32-1)/32)*32。调用示例以下为单算子模式下的完整调用流程示例取自 torchapi_low_latency_dispatch.md量化模式为 pertoken 动态量化并演示了特殊专家zero_expert_num/copy_expert_num的用法import torch import torch_npu from torch.multiprocessing import Process import torch.distributed as dist from cann_ops_transformer.ops import MoeDistributeBuffer # 控制模式 quant_mode 2 # 2 为动态量化 input_dtype torch.bfloat16 world_size 16 # EP 通信域大小例16 个 DIE shared_expert_rank_num 0 num_experts 32 # MoE 专家数 bs 8 # 本卡 token 数量 h 7168 # hidden size k 8 # topK zero_expert_num 1 # 零专家 copy_expert_num 1 # copy 专家 const_expert_num 0 # 常量专家 globalBS bs * world_size def run_npu_process(rank): torch_npu.npu.set_device(rank) dist.init_process_group(backendhccl, rankrank, world_sizeworld_size, init_methodftcp://127.0.0.1:50001) ep_group dist.new_group(backendhccl, rankslist(range(world_size))) x torch.randn(bs, h, dtypeinput_dtype).npu() topk_idx torch.randint(0, num_experts, (bs, k), dtypetorch.int32).npu() topk_weights torch.randn(bs, k, dtypetorch.float32).npu() scales torch.randn(num_experts, h, dtypetorch.float32).npu() # 构造通信域 buffer distribute_buffer MoeDistributeBuffer(ep_group) # dispatch量化 EP 域 alltoallv 分发 expand_x, dynamic_scales, assist_info_for_combine, expert_token_nums, ep_recv_counts, _ \ distribute_buffer.low_latency_dispatch( xx, topk_idxtopk_idx, shared_expert_num0, shared_expert_rank_numshared_expert_rank_num, num_expertsnum_experts, quant_modequant_mode, num_max_dispatch_tokens_per_rankglobalBS, zero_expert_numzero_expert_num, copy_expert_numcopy_expert_num, const_expert_numconst_expert_num) expand_x expand_x.to(input_dtype) # combine将 dispatch 输出配套传回完成 token 回收 x distribute_buffer.low_latency_combine( xexpand_x, topk_idxtopk_idx, topk_weightstopk_weights, assist_info_for_combineassist_info_for_combine, ep_send_countsep_recv_counts, shared_expert_num0, shared_expert_rank_numshared_expert_rank_num, num_expertsnum_experts, num_max_dispatch_tokens_per_rankglobalBS, ori_xx, zero_expert_numzero_expert_num, copy_expert_numcopy_expert_num, const_expert_numconst_expert_num) if __name__ __main__: ps [Process(targetrun_npu_process, args(rank,)) for rank in range(world_size)] for p in ps: p.start() for p in ps: p.join() print(run npu success.)调用前需满足前置校验shared_expert_rank_num不能大于ep_world_size当shared_expert_rank_num 0时ep_world_size必须是其整数倍num_experts必须是moe_rank_num ep_world_size - shared_expert_rank_num的整数倍。low_latency_dispatch与low_latency_combine均支持单算子模式与 torchair 图模式torch.compile(model, backendnpu_backend)调用图模式下的完整示例同样可在 torchapi_low_latency_dispatch.md 中找到。源码与测试验证若希望深入理解算子实现与验证行为可在当前仓库中重点关注以下文件算子原型声明op_graph/moe_distribute_dispatch_v3_proto.h含各输入/属性/输出的数据类型与取值范围注释宿主算子定义与内核注册op_host/moe_distribute_dispatch_v3_def.cppA3 与 950 分别注册_a3/_apt内核入口Shape 推导op_host/moe_distribute_dispatch_v3_infershape.cppQuantMode 0~5、A/K/local_expert_num 等 shape 规则的落地实现aclnn 单算子接口op_api/aclnn_moe_distribute_dispatch_v5.cppPython 高阶封装mc2/common/torch_extension/moe_distribute_buffer.pyMoeDistributeBuffer、get_low_latency_ccl_buffer_size、low_latency_dispatch/low_latency_combine接口文档mc2/common/docs/torchapi_low_latency_dispatch.md单元测试tests/ut/op_api/test_aclnn_moe_distribute_dispatch_v5.cpp、tests/ut/op_kernel/test_moe_distribute_dispatch_v3.cpp、tests/ut/op_host/test_moe_distribute_dispatch_v3_infershape.cpp以及 arch22/arch35 两套 tiling 单测覆盖了算子在不同架构下的 Tiling 逻辑。综合来看MoeDistributeDispatchV3是 MoE 大模型在 NPU 上实现专家并行分发链路的关键算子它以可选量化 EP 域 AllToAllV为核心通过context/ccl_buffer_size实现低延迟通信域管理并以low_latency_dispatch接口与MoeDistributeCombineV3系列算子形成闭环支撑了 Atlas A3 与 Ascend 950 系列产品上的 MoE 训练与推理。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考