CANN SHMEM 经典 MoE Dispatch 示例详解基于对称窗口的 Token 路由与 Combine 辅助信息生成【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本篇文章深入解析 CANN SHMEM 开源仓库中 examples/dispatch/dispatch_classic 这一经典 MoEMixture of ExpertsDispatch 示例。它展示了如何基于 OpenSHMEM 风格的对称数据窗口SHMEM window与 MTE 数据搬运能力将本 PEProcessing Element的 token/topK 按expert_ids路由到目标 expert 所在 PE并同步生成 combine 阶段所需的辅助信息。读完本文你将掌握该算子的输入输出语义、Kernel 三阶段实现原理、构建与运行方法、全部命令行参数含义以及如何用内置的性能测试框架对 dispatch 进行多 shape 扫描与基准测量。示例定位与核心价值该示例实现了非量化 MoE 的经典 dispatch 算子对应设计文档DOC/moe_dispatch_combine_non_quant_architecture.md。在 MoE 前向计算中token 需要根据路由结果expert id被分发到保存对应 expert 权重的设备上完成计算计算完成后再把结果聚合回原设备前者即为 Dispatch后者为 Combine。本示例聚焦 Dispatch 阶段其职责是根据expert_ids将本 PE 的 token/topK 路由到目标 expert 所在 PE并生成 combine 阶段需要的辅助信息。从实现策略上该示例与同目录下的 双平面 MoE Dispatch 示例 形成对照经典版本优先保证逻辑清晰、可验证和输出顺序稳定数据面统一使用 MTE 传输双平面版本则在此基础上同时启用 MTE 与 SDMA 自适应选择传输路径追求通信阶段性能。建议先以经典版本建立正确性基线再用双平面版本做性能对比。功能说明输入与输出语义输入输入说明x本 PE 的 token hidden states形状为(bs, h)expert_ids每个 token/topK 对应的全局 expert id形状为(bs, topk)输出输出说明expand_x按目标 expert 聚合后的 token 数据assist_info_for_combinecombine 回传使用的辅助信息格式为[src_rank_id, src_token_id, src_topk_id]ep_recv_count按(local_expert_id, src_rank_id)排列的累计接收计数expert_token_nums每个本地 expert 实际收到的 token 数输出顺序固定为(local_expert_id, src_rank_id)这一确定性的顺序是后续 combine 阶段与 golden 校验的基础。其中ep_recv_count保存的是累计计数cumulative count即从第一个 segment 累加到当前 segment 的 token 总数等价于每个 segment 在最终expand_x中的前缀偏移边界expert_token_nums则是按本地 expert 维度聚合的段内 token 数。实现逻辑Kernel 三阶段流水Kernel 启动pe_size个 AIV core每个 active core 负责一个目标 rankcore_id dst_rank的 core 负责写入对应目标 PE 的对称窗口。整个流程分为发送、等待与 compact 三个阶段。Stage 1发送阶段路由与写入每个 active core 扫描本 PE 的expert_ids统计发往自己负责目标 PE 的每个 local expert 的 token 数对应代码中的segment_counts与slot_offsets数组。对每个命中的 token/topK通过 MTEput_nbi写 payload 到目标 PE 的 SHMEM window。写入assist_info_for_combine、token ready flag 和 segment count 信号。所有发送 core 完成后aclshmem_quiet()保证写操作对远端可见。Stage 2等待阶段接收端aiv_index 0的 core等待各来源 rank 的 count 信号通过aclshmem_signal_wait_until轮询count_base中的 ready 标志随后累加得到累计接收计数。Stage 3compact 阶段根据 count 构造ep_recv_count和expert_token_nums再按确定顺序(local_expert_id, src_rank_id)将 payload 与辅助信息 compact 到最终输出expand_x、assist_info_for_combine。compact 统一放在 core 0 上执行保证前缀边界与输出写入按唯一确定顺序完成。对称窗口的内存布局从 dispatch_kernel.cpp 可以确认SHMEM window 由四段连续区域组成payloadtoken 数据 - assist辅助信息 - readytoken 级 ready flag - countsegment 计数对应 host 侧 main.cpp 中的DispatchWindowBytes/DispatchCountOffset计算每个 region 都按 32 字节对齐。Kernel 在栈上为每个目标 local expert 分配固定工作区segment_counts与slot_offsets容量上限为DISPATCH_MAX_LOCAL_EXPERT_NUM 1024这是expertPerPe上限的由来。关键 API 调用链aclshmemx_mte_put_nbi非阻塞 MTE 远端写用于 payload 与 ready 信号数据面。aclshmemx_signal_op(..., ACLSHMEM_SIGNAL_SET, dst_rank)向远端写辅助信息字段控制面。aclshmem_signal_wait_until(..., ACLSHMEM_CMP_EQ, ...)接收端等待 ready/count。aclshmem_quiet()同步本地写队列。aclshmemi_sync_core_soft()Kernel 内 core 间软同步。构建在仓库根目录执行bash scripts/build.sh按目标平台追加参数A2/A3 平台bash scripts/build.sh -examplesAscend950 平台bash scripts/build.sh -soc_type Ascend950 -examples构建产物包括build/bin/dispatchhost 可执行程序与链接进可执行程序的 kernel 库。示例通过 examples/CMakeLists.txt 中的aclshmem_add_collective_example(dispatch)宏注册dispatch_kernel.cpp被编译为dispatch_kernel共享库main.cpp编译为dispatch可执行文件两者均链接shmem库并包含examples/utils等头文件目录。运行基础 2 卡测试cd examples/dispatch/dispatch_classic bash scripts/run.sh -pes 2 -bs 8 -h 16 -topk 2 -expertPerPe 2 -type int32_t8 卡、64 expert 测试cd examples/dispatch/dispatch_classic bash scripts/run.sh -pes 8 -bs 8 -h 16 -topk 2 -expertPerPe 8 -type int32_trun.sh的执行流程见 scripts/run.sh调用 data_gen.py 生成输入x、路由矩阵expert_ids与 golden 数据写入golden/shape_bs_h_topk_moe_expert_num_pes/rank_id/并附带meta.json记录 shape 元信息。int32_t使用均匀随机整数float16_t/bfloat16_t使用均匀浮点并依赖ml_dtypes包。为每个 PErank启动一个dispatch进程通过SHMEM_UID_SESSION_ID与自适应端口8766 case_index % 1000区分多个 case 的 bootstrap 会话。各进程从 golden 目录读入输入完成 kernel 执行后把结果写入output/expand_x_pe.bin、assist_info_pe.bin、ep_recv_count_pe.bin、expert_token_nums_pe.bin。调用 check_dispatch.py 对四个输出逐一与 golden 比对整数类型使用严格相等浮点类型使用rtol1e-3, atol1e-3的np.allclose容差校验。host 侧 main.cpp 负责解析参数、aclInit/aclrtSetDevice、通过aclshmemx_init_attr初始化 SHMEM 运行时单实例堆大小 1GB、按-type分发到int32_t/fp16_t/bf16_t模板实例最后依次执行aclshmem_barrier_all、aclshmem_finalize与aclFinalize。参数说明run.sh支持的完整参数如下-pes n PE 数量单机示例要求与 -gnpus 相同默认 2。 -gnpus n 本机启动的 NPU 数量必须与 -pes 相同默认 2。 -bs n 每个 PE 的 token 数默认 8。 -h n token hidden size默认 16。 -topk n 每个 token 路由的 expert 数默认 2。 -expertPerPe n 每个 PE 上的 local expert 数范围为 [1, 1024]默认 2。 -type dtype 数据类型支持 int32_t、float16_t、bfloat16_t默认 int32_t。 -fpe id 首个 PE 编号默认 0。 -fnpu id 起始 NPU id默认 0。 -ipport url SHMEM bootstrap 地址默认 tcp://127.0.0.1:8766。 --perf 性能测试模式保留正确性校验并将性能 CSV 写入 --output-dir 目录。 --warmup n 性能测试的预热迭代次数不计入统计默认 5。 --loops n 性能测试的正式测量迭代次数默认 50。 -pes-list a,b,... 性能测试模式下扫描的 PE 数量列表。 -bs-list a,b,... 性能测试模式下扫描的每个 PE token 数列表。 -h-list a,b,... 性能测试模式下扫描的 hidden size 列表。 --topk-list a,b,... 性能测试模式下扫描的 topk 列表。 --expert-per-pe-list a,b,... 性能测试模式下扫描的 local expert 数列表。 --prof-pe id|all 性能采集的 PE 编号为 all 时轮流采集每个 PE 并汇总默认 0。 --output-dir dir 性能测试 CSV 输出目录默认 output/perf。 -a|--analyse mode 性能结果处理方式plot图形化展示、md生成 Markdown 报告、none不处理默认。expertPerPe上限为 1024Kernel 在 AI core 栈上为每个目标 local expert 分配固定工作区segment_counts/slot_offsets超过该上限会被 host 侧拒绝main.cpp 会打印max supported value is 1024并返回错误。此外run.sh在非 perf 模式下会强制校验-gnpus与-pes相等否则直接退出。几个值得注意的实现细节-fpe在 run.sh 中仍被保留为一个 CLI 槽位以兼容共享示例脚本但 host 侧实际将其忽略见 main.cpp 的注释。设备 id 由args.pe_id % g_npus f_npu推导-fnpu用于支持非 0 起始的 NPU 编号。run.sh中-pes与-gnpus被联动赋值保证单机示例下两者始终一致。性能测试run.sh --perf会在每个 shape 上保留正确性校验并在output/perf/下写入 CSV。单 shape profilingcd examples/dispatch/dispatch_classic bash scripts/run.sh --perf -pes 2 -bs 8 -h 256 -topk 2 -expertPerPe 2 -type int32_t \ --warmup 5 --loops 50多 shape、多卡数 sweepcd examples/dispatch/dispatch_classic bash scripts/run.sh --perf --pes-list 2,4,8 --bs-list 8,16,32 --h-list 64,256,1024 \ --topk-list 2 --expert-per-pe-list 2,8 -type int32_t --prof-pe all \ --warmup 5 --loops 50sweep 模式下run.sh对pes-list × bs-list × h-list × topk-list × expert-per-pe-list做全组合遍历--prof-pe all时对每个 PE 轮流采集并汇总生成dispatch_perf_summary.csv。CSV 指标指标含义full_op完整 dispatch包括通信、元数据构造、compact 和同步comm_onlyStage 1 payload 通信及必要的元数据/status 协议单 rank 文件名为dispatch_perf_rankrank.csv。性能数据的产生机制Kernel 内通过SHMEMI_PROF_START/END见 dispatch_kernel.cpp为full_frame_id0与comm_frame_id1两个计时框架打点其中comm_only框架在 Stage 1 结束后即关闭。Host 侧 moe_perf_host.h 中的MoeAppendPerfCsvRows通过aclshmemx_get_prof读取各 core 的周期计数时间换算系数按 SoC 自适应Ascend950 平台 1000 cycles 对应 1us其余平台 50 cycles 对应 1usMoeGetCycleToUs。每行 CSV 记录DataSize/B、Npus、Blockspe_size、UBsize/KB 190、Bandwidth/GB/s、CoreMaxTime/us、Metric以及BS/H/TopK/ExpertPerPe/Dtype/Warmup/Loops/ProfPe/CaseId等字段并追加每个 active core 的SingleCoreTime/us。运行时会通过环境变量SHMEM_CYCLE_PROF_PE指定被采集的 PEMoeGetProfPe非 perf 模式自动以 PE 0 兜底。结果文件可通过-a plot交给 perf_data_process.py 绘制图表或以-a md生成 Markdown 报告。DISPATCH_UB_SIZE_KB 190对应 Kernel 侧UB_DMA_MAX_SIZE 190 * 1024UB 单次 DMA 最大搬运字节数这是 MTE payload 传输按h个元素分块拷贝时的块大小上限。从经典走向双平面完成本示例后可以继续探索同组的 双平面 MoE Dispatch 示例。双平面版本保持与经典 dispatch 完全一致的输入输出语义与输出顺序因此同一组 golden/check 脚本可以验证两条路径。双平面的差异在于当某个(dst_rank, dst_local_expert)segment 的 payload 字节数大于 2MB 且大于当前 PE 的远端平均 segment 大小判定逻辑使用交叉相乘避免整数截断该大段改用aclshmemx_sdma_put_nbi传输小段与全部控制面信号仍走 MTESDMA 每提交 256 次 issue 即调用aclshmemx_sdma_quiet防止 outstanding 请求积压。需要注意的是SDMA 功能要求 CANN 9.0.0 及以上且当前暂不支持 Ascend950 平台基础安装与独立 SDMA demo 可参考 examples/sdma/README.md。对经典与双平面做性能对比时可对相同 shape 分别执行--perf并比较 CSV 中的comm_only若comm_only降低说明大段 payload 走 SDMA 对通信阶段有效若full_op收益不明显则需要结合 shape、路由倾斜度和后续 compact/同步成本综合判断。进一步阅读双平面版本与经典版本的差异与选型建议dispatch_doubleplane README数据生成与 golden 构造data_gen.py正确性校验脚本check_dispatch.py性能 CSV 生成与单位换算moe_perf_host.h构建系统集成方式examples/CMakeLists.txt【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考