CANN MoE Dispatch/Combine 通信算子性能优化实践Kernel 直调与 ACL SHMEM 实现专家分发与聚合【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples导读本篇文章基于 CANN Samples 仓库中Samples/2_Performance/moe_dispatch_and_combine_story的完整实践案例讲解在 MoEMixture-of-Experts大模型训练与推理场景下如何用Ascend C Kernel 直调内核调用符搭配CANN ACL SHMEM对称共享内存实现细粒度的Dispatchtoken 分发与Combine结果聚合通信算子并通过通信前量化与按 token 的流水并行收发打破传统alltoallv的链路长、同步点多等瓶颈。读完本文你将掌握Dispatch/Combine 的数据面/状态面内存划分与计算步骤、MXFP8/INT8 通信前量化的完整实现细节、性能收益实测数据以及从环境准备、构建、数据生成到运行与精度验证的完整复现路径。背景MoE 专家并行的通信瓶颈在大模型训练与推理中MoE 架构凭借动态专家激活带来的计算稀疏性优势以及在千亿参数规模下的高吞吐推理能力已成为超大规模模型的关键技术路线。MoE 主要通过Dispatch与Combine两个核心过程实现输入 token 的动态分发与多专家输出的高效聚合从而在维持海量参数规模的同时获得较高计算效率。但随着专家并行EP规模持续扩大专家之间更频繁的数据交换会带来显著的通信开销逐步演化为影响端到端推理性能的主要瓶颈。在整网推理过程中各层的 TopK 路由结果会动态变化。若采用传统的alltoallv通信流程通常需要先交换各 rank 的发送量信息再通过排序/重排将发往同一专家的数据聚集最后再进行一次alltoallv发送 token 数据。这一流程链路长、同步点多整体效率不高。因此本案例将Dispatch/Combine的数据收发过程进一步细粒度拆分基于共享内存的写入机制按 token 逐个发送与接收通过流水并行pipeline重叠通信与数据整理开销以获得更好的整体性能。方案选型Kernel 直调 ACL SHMEM本案例采用两条关键技术路径的组合整体实现更贴近 Ascend C 工程习惯与开源通信组件生态Kernel 直调核函数由应用在 Device 侧直接启动。在 dispatch_and_combine_final.asc 中DispatchKernel与CombineKernel均以DispatchKernelAIV_CORE_NUM, nullptr, stream(...)的形式直接下发AIV_CORE_NUM 64与aclrtStream、Host/Device 内存管理等 ACL 接口在同一套异步流水线中编排无需额外引入一层算子封装即可验证通信与计算逻辑便于对齐文档、调试与迭代。ACL SHMEM对称共享内存通过对称共享内存窗口与设备侧aclshmem_ptr、aclshmemx_mte_put_nbi等接口提供跨 PE 的单边写与统一寻址抽象便于实现「按 token 写入对端 状态区同步」的细粒度通信。相较于层层alltoallv与全局排序更易与本案例的数据面/状态面划分及流水并行设计结合并复用 CANN 开源 shmem 组件而非自建底层传输。从源码实现看dispatch_and_combine_final.ascHost 侧先通过aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, attributes)完成通信上下文初始化再以aclshmem_align申请 1 GBSHMEM_SPACE_SIZE对称共享内存随后在循环中以交替启动DispatchKernel与CombineKernel共 20 轮loopTimes最后aclshmem_finalize()释放资源。进程编排则由 utils.h 中的MoeDemoForkMain完成——主进程 fork 出rankNum个子进程每个子进程绑定一个rankId对应一张 NPU 设备aclrtSetDevice(deviceId)并以tcp://127.0.0.1:8998作为 PE 间通信地址。Dispatch 功能介绍Dispatch 算子完成以下核心功能完成 token 分发根据 token 选中的专家将 token 发送到对应卡上。为支持后续 FFN 计算将各卡发送过来的 token 连续重排。为支持后续Combine处理统计接收 token 信息支撑Combine将 FFN 计算后的 token 发送回源端。将量化提前到通信之前将量化后的int8 token与fp32 scale合并发送并在接收端重排分离。内存划分PE rank0 PE rank1 PE rankN-1 ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ 数据面 │ │ 数据面 │ │ 数据面 │ │ [0,1022MB) │ │ [0,1022MB) │ │ [0,1022MB) │ │ │ │ │ │ │ │ winOffset_ │ │ winOffset_ │ │ winOffset_ │ │ ▼专家栅格 │ │ ▼专家栅格 │ │ ▼专家栅格 │ │ │ │ │ │ │ ├──────────────┤ ├──────────────┤ ├──────────────┤ │ 状态面 │ │ 状态面 │ │ 状态面 │ │ 1022MB │ │ 1022MB │ │ 1022MB │ │ Dispatch区 │ │ Dispatch区 │ │ Dispatch区 │ │ Combine区 │ │ Combine区 │ │ Combine区 │ └──────────────┘ └──────────────┘ └──────────────┘ │ │ │ └───── mte_put 可以跨 PE 写对方窗口内的偏移 ───────┘ 访问对端aclshmem_ptr(shmemCtx, pe) offset该划分在源码中有精确对应moe_distribute_comm.h中定义了STATUS_REGION_OFFSET 1022UL * 1024UL * 1024UL状态面起始偏移与STATE_OFFSET 32U状态条目间隔Dispatch 侧还定义了WIN_STATE_OFFSET、FLAG_FIELD_OFFSET、CUMSUM_CAL_OFFSET、CUMSUM_FLAG_OFFSET等状态区子域见 moe_distribute_dispatch.hCombine 侧则使用COMBINE_STATE_OFFSET将本卡状态空间放置在 Dispatch 区之后见 moe_distribute_combine.h。跨 PE 访问对端窗口统一通过aclshmem_ptr(shmemContextGM_, rankId) offset计算地址并以aclshmemx_mte_put_nbi发起非阻塞单边写。计算步骤循环发送数据写UB - quant - 写远端内存流水并行从 HBM 搬运待发送数据到 AIV 核内的 UBuffer。对 UBuffer 上的 token 进行量化并将量化后的int8 token与fp32 scale拼接。将源端信息rank_id、bs_id、k_offset即三元组拼接到量化后的 UBuffer。根据expert_ids查找对端 rank 地址并通过AIV UBmem执行写远端发送。发送flag标识与count到所有对端根据expert_id统计发往每个 expert 的 token 数分核处理。统计后执行SyncAll多核同步确认所有核发送完成后将完成flag与count通过AIV UBmem发送到对端。等待对端写入完成并计算偏移分核读取状态区flag直到全部为1表示对应对端已完成发送。读取状态区count计算各 rank 数据搬运到输出中的偏移。执行SyncAll多核同步。本地数据整理分核并行将通信 Shared Memory 中的数据整理到最终输出data、scale、expert token num。输出各 expert 的 token 数量。从源码结构可以印证上述步骤的多核分工moe_distribute_dispatch.h 的Process()将 64 个 AIV 分成两组——前aivUsedAllToAll_个核执行AllToAllDispatch()token 数据的按专家发送后aivUsedCumSum_个核执行CalCumSum()统计发往每个专家的 token 计数、输出epRecvCnt/expertTokenNums最后统一执行LocalWindowCopy()完成本卡数据连续化。其中aivUsedCumSum_按moeExpertNum / 32计算并钳制在[1, min(aivNum_/2, 16)]范围内。发送路径上TokenToExpertInQuant完成「UB 搬运 → MXFP8 量化 → 三元组填充 →aclshmemx_mte_put_nbi写对端」的整条流水接收路径上WaitAndFormatOutput通过CheckDataArriveWithFlag轮询状态区 flag数据到齐后按ep_recv_count前缀和计算输出偏移将通信窗口中的 token 连续化到expand_x并把动态 scaledynamic_scales与三元组assist_info_for_combine即expand_idx分离输出。Combine 功能介绍Combine 算子完成以下核心功能将 FFN 计算后的 token 发送回源端。将选中的KMoE 专家数个 FFN 结果加权求和完成combine计算。内存划分PE rank0 PE rank1 PE rankN-1 ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ 数据面 │ │ 数据面 │ │ 数据面 │ │ [0,1022MB) │ │ [0,1022MB) │ │ [0,1022MB) │ │ │ │ │ │ │ │ winOffset_ │ │ winOffset_ │ │ winOffset_ │ │ ▼专家栅格 │ │ ▼专家栅格 │ │ ▼专家栅格 │ │ │ │ │ │ │ ├──────────────┤ ├──────────────┤ ├──────────────┤ │ 状态面 │ │ 状态面 │ │ 状态面 │ │ 1022MB │ │ 1022MB │ │ 1022MB │ │ Dispatch区 │ │ Dispatch区 │ │ Dispatch区 │ │ Combine区 │ │ Combine区 │ │ Combine区 │ └──────────────┘ └──────────────┘ └──────────────┘ │ │ │ └───── mte_put 可以跨 PE 写对方窗口内的偏移 ───────┘ 访问对端aclshmem_ptr(shmemCtx, pe) offsetCombine 侧的状态区在源码中通过GetShmemSignalAddraclshmem_ptr(shmemSpace, pe) 1022 * 1024 * 1024与winStatusOffset_ COMBINE_STATE_OFFSET dataState_ * WIN_STATE_OFFSET定位dataState_用于双缓冲状态面轮换0/1 翻转避免前后两轮通信的 flag 相互干扰。计算步骤循环发送数据与发送flag按总 token 数分核从recv_count读取总 token 数并平均分配到各核。每个 token 携带三元组信息rank_id、bs_id、k_offset发送时根据三元组计算对端地址偏移。使用AIV UBmem将数据发送到对端。每个 token 发送完成后按 token 发送flag。循环等待并执行combine按 batch size 分核按 batch 循环处理。等待对应 batch 状态K1个状态位全部为1。状态齐备后从通信 Shared Memory 搬运各 FFN 结果。从输入 HBM 搬运expert scale。执行combine对各 MoE FFN 结果乘以expert scale后求和再叠加 shared FFN 结果。源码实现moe_distribute_combine.h中SplitCoreCal()依据selfSendCnt_本卡需要发回的 token 总数把发送任务均分到各 AIVExpertAlltoAllDispatchCopyAdd()按startTokenId_起的每个 token 读取expand_idx三元组rank_id / token_id / topk_id换算目标窗口偏移epOffset tokenId * axisK_ topkId后写入对端并在写完数据后以aclshmemx_mte_put_nbi向对端状态区写 flagLocalWindowCopy()则按 batch 循环WaitDispatch(tokenIndex)轮询K个状态位全部置 1随后用ProcessMoeExpert依次搬运各专家 FFN 结果、乘expertScales权重累加到sumFloatBufLocal_最终Cast回float16写回输出x。性能优化实践通信前量化在数据进入通信链路之前完成量化可显著降低传输字节数并改善算子侧效率。两条路径对比如下Dispatch量化前移到发送侧通信时同步携带动态量化参数如对端解析所需的 scale。Combine发送前按block量化为int8压缩通信量接收后再反量化参与计算。搬运前量化数据搬运GM 输入 → UB并转为float。算 scale取数据最大绝对值 → 按int8表示范围计算 scale。存 scale将 scale 写入 UB 后半段。量化计算广播 scale → 原始数据 ÷ scale → 映射到int8数值范围。转 int8float→fp16→int8写入 UB 前半段。硬件同步保证流水线各阶段顺序正确。通信搬运数据从 UB 搬运量化后的数据到对端缓冲区。搬运后反量化拆输入从同一段缓冲读取int8数据及尾部 scale。类型转换int8→fp16→floatscale 同步转为float。反量化计算广播 scale →int8数据 × scale → 恢复浮点幅值。写回将恢复后的浮点数据写回原int8数据区。两种量化路径的源码细节Dispatch 的 MXFP8 动态量化MoeDistributeDispatchQuant::QuantDynamicMxFp8moe_distribute_dispatch_quant.h依次调用ComputeMaxExp按 32 个元素一个 block 计算最大指数、ComputeScale根据FP8_E5M2_MAX_EXP等常量推导共享指数 scale并以fp8_e8m0_t存储、ComputeDatafp16 → fp32 → fp8_e5m2逐段 Cast 完成量化。量化后的fp8 token与fp8_e8m0 scale拼接存放scale 数量为Align2(Ceil32(axisH_))即每 32 个元素产生一个 scale。这里输入的XType float16_t、输出ExpandXOutType fp8_e5m2_t的模板约束直接对应数据流。Combine 的 INT8 block 量化Int8QuantProcess/Int8DequantProcess按 32 字节 blockelement_per_block 8个 float处理Abs → BlockReduceMax → Muls(1/SCALE_PARAM) → Div → Cast(float→fp16→int8)完成量化接收侧再以int8 → fp16 → float的 Cast 链配合Brcb广播 scale 完成反量化SCALE_PARAM 127.0。性能数据基于 dispatch 和 combine 算子基础版本各自实现一个非量化和量化版本构建低时延和高吞吐场景性能验证用例详细性能数据如下算子名BSHK卡数单卡专家数量化类型算子耗时(us)提升比例dispatch87168884非量化17.1basedispatch87168884mxfp8量化16.61.03dispatch2567168884非量化170basedispatch2567168884mxfp8量化1131.50combine87168884非量化19.4basecombine87168884int8量化18.81.03combine2567168884非量化150.7basecombine2567168884int8量化1081.40从数据可以看出低时延小 batchBS8场景下通信本身开销占比低量化带来的收益有限约 1.03 倍而在高吞吐大 batchBS256场景下通信数据量成为主导因素通信前量化可将dispatch提升至 1.50 倍、combine提升至 1.40 倍。示例演示端到端 Dispatch Combine 流程本示例演示MoEMixture-of-Experts场景下的分布式 Dispatch Combine的端到端流程Dispatch按expert_idsTopK 路由结果将 token 特征从各 rank 发往目标 expert跨 rank all-to-all并输出用于 Combine 的辅助信息如expand_idx、ep_recv_count、expert_token_nums等。Combine根据 Dispatch 生成的assist_info_for_combine/ep_recv_count等信息将 expert 侧处理后的expand_x按 TopK 权重expert_scales汇聚回原 token按 rank 分布式 all-to-all 本地按 TopK sum。数据生成脚本 gen_data.py 在 Python 侧以torch.distributedgloo 后端 all_to_all_single独立实现一遍同语义的 dispatch/combine 参考流程用于生成 goldenDispatch 阶段按expert_ids排序、组包token 动态 scale 源信息三元组、all-to-all 交换、按目标专家前缀和得到ep_recv_countCombine 阶段按源信息回发 FFN 结果按token_id * topk topk_id重排后乘expert_scales加权求和。这一 Python 参考实现与 Device 侧算子的输入输出逐一对应构成精度比对的黄金基准。运行环境与约束硬件NPU 多卡环境。芯片型号通过根工程NPU_ARCH参数指定当前样例支持dav-3510。软件环境需要已安装 Ascend CANN Toolkit并在构建和运行前加载 Toolkit 环境变量。依赖组件本样例通过根工程的 cmake/shmem.cmake 引入cann_samples::shmemACL SHMEM 库构建脚本 CMakeLists.txt 中通过include(${PROJECT_SOURCE_DIR}/cmake/shmem.cmake)完成链接同时依赖ascendcl、runtime、tiling_api等 ACL 运行库。目录结构moe_dispatch_and_combine_story/ ├─ CMakeLists.txt ├─ include/ # device侧核函数 │ ├─ moe_distribute_dispatch.h │ ├─ moe_distribute_dispatch_quant.h │ └─ moe_distribute_combine.h ├─ src/ │ ├─ dispatch_and_combine_final.asc # 终极性能版本 │ └─ utils.h └─ scripts/ ├─ gen_data.py # 生成输入与 golden └─ verify_result.py # 比对输出与 golden说明include/下还包含 moe_distribute_comm.h公共常量/对齐工具与moe_distribute_dispatch_non_quant.hsrc/下另有0_non_quant_naive.asc基础版本dispatch_and_combine_final.asc为含量化与流水并行的终极性能版本。在本示例目录运行本示例运行过程中会在本示例目录生成input/输入 bin按chip_{rankId}分目录golden/golden bin用于精度对比output/算子输出 bin按chip_{rankId}分目录环境准备构建和运行前需要先加载 Ascend Toolkit 环境变量。若使用 root 用户按默认路径安装可执行source /usr/local/Ascend/ascend-toolkit/set_env.sh如果 Toolkit 安装在自定义路径请将上述路径替换为实际安装目录下的set_env.sh。建议在同一个 shell 会话中完成后续的构建、数据生成、运行和校验避免环境变量丢失。此外gen_data.py 依赖numpy、torch需支持torch.float8_e5m2并按数据类型需要ml-dtypesbfloat16 / float8_e5m2、en-dtypesfloat8_e8m0等 Python 包请确保已安装对应依赖仓库根目录 requirements.txt 提供了参考依赖清单。构建在本示例目录执行 CMake 构建目标为moe_dispatch_and_combine_dispatch_and_combine_final见本示例目录下的 CMakeLists.txt。构建完成后可执行文件会生成到build/Samples/2_Performance/moe_dispatch_and_combine_story/目录下。以下命令中的${cann_samples_path}表示用户本地cann-samples仓库所在目录请根据实际路径替换。cd ${cann_samples_path}/Samples/2_Performance/moe_dispatch_and_combine_story cmake -S ../../../ -B ../../../build -DNPU_ARCHdav-3510 cmake --build ../../../build --target moe_dispatch_and_combine_storycann_sample_check_arch(dav-3510)会在配置阶段校验当前支持的 NPU 架构dispatch_and_combine_final.asc以-O3与--npu-arch${NPU_ARCH}编译并链接m、dl、platform、tiling_api、cann_samples::shmem、ascendcl、runtime、stdc库。生成测试数据input golden及 output构建完成后在本示例目录生成一组用于运行和精度校验的测试数据命令行未传入的参数使用默认值python3 ./scripts/gen_data.py --chip-num-per-server 2 --bs 8该脚本会在本示例目录生成input/和golden/。其中input/作为算子运行输入golden/作为后续verify_result.py的精度比对基准。算子运行后会在本示例目录生成output/用于保存实际运行输出。脚本参数均有默认值只需要传入与默认值不同的配置。常用参数如下--bs算子入参 batch size 大小。--h算子入参 hidden size 大小。--k算子入参 topk 大小。--chip-num-per-server生成的 rank 数。--token-dtype-choice0 表示 bfloat161 表示 float16默认 1。--quant-mode0 表示不量化4 表示 MXFP8 动态量化。此外脚本还支持--random-seed随机种子默认 0、--moe-expert-num默认取chip_num_per_server * 4、--expert-recv-info-type0 表示expert_token_nums输出累计前缀和1 表示输出各专家 token 数默认 1。注意--chip-num-per-server与运行时的rankNum、--bs与运行时的bs必须保持一致。运行Dispatch Combine测试数据生成后在本示例目录运行构建产物。命令行参数依次为rankNum和bs需要与生成数据时的--chip-num-per-server、--bs保持一致../../../build/Samples/2_Performance/moe_dispatch_and_combine_story/moe_dispatch_and_combine_dispatch_and_combine_final rankNum bs例如前面生成的是 2 张卡、batch size 为 8 的数据则执行../../../build/Samples/2_Performance/moe_dispatch_and_combine_story/moe_dispatch_and_combine_dispatch_and_combine_final 2 8运行完成后算子输出会写入本示例目录下的output/并按chip_{rankId}分目录保存。程序内部会 fork 出rankNum个进程并行执行各 rank 的 Dispatch Combine循环 20 轮并在每个 rank 结束时打印[SUCCESS] demo run success in relative_pe_id rankId若所有 worker 均成功主进程会输出All workers finished. Status: SUCCESS。精度验证output vs golden算子运行结束后在本示例目录执行精度校验脚本python3 ./scripts/verify_result.py该脚本默认读取本示例目录下的golden/与output/逐个比对golden/**/*.bin与output/**/*.bin按相对路径对应并打印每个 bin 的一致性结果。其比对对象覆盖 Dispatch 输出expand_x、dynamic_scales、assist_info_for_combine、expert_token_nums、ep_recv_count与 Combine 输出x从而端到端验证「分发 → 聚合」全链路的数值正确性。总结本案例展示了在 CANN 平台上实现高性能 MoE 通信算子的完整路径以 Kernel 直调保持与 ACL 生态的无缝衔接以 ACL SHMEM 对称共享内存实现跨 PE 的按 token 细粒度单边写与状态面同步配合通信前量化Dispatch 用 MXFP8、Combine 用 INT8与多核流水并行在 BS256 高吞吐场景下将dispatch提升至 1.50 倍、combine提升至 1.40 倍数据来源见上文性能表。如需深入源码细节可继续阅读 moe_distribute_dispatch.h、moe_distribute_combine.h、moe_distribute_dispatch_quant.h 以及入口实现 dispatch_and_combine_final.asc。【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考