CANN SHMEM 双平面 MoE Combine 示例详解:MTE+SDMA 自适应分段回传的实现与性能对比

CANN SHMEM 双平面 MoE Combine 示例详解:MTE+SDMA 自适应分段回传的实现与性能对比 CANN SHMEM 双平面 MoE Combine 示例详解MTESDMA 自适应分段回传的实现与性能对比【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本文导读本篇文章以 CANN SHMEM 开源仓库中的examples/combine/combine_doubleplane示例为主线深入讲解非量化 MoE combine 算子的“双平面”版本——如何在 expert 输出回传阶段同时启用 MTEMulti-Task Engine多任务引擎数据搬运路径与 SDMASynchronous DMA两条数据面根据 segment分段大小自适应选择传输路径。读完本文你将掌握双平面 combine 的输入输出语义、SDMA 判定阈值逻辑、Host/Kernel 两级初始化方式、完整构建运行与性能对比方法并能结合仓库源码理解其底层调用链。一、示例定位与经典 Combine 的关系combine_doubleplane位于 examples/combine/combine_doubleplane它实现的是非量化 MoE combine 的双平面版本。它保持与 examples/combine/combine_classic 完全相同的输入输出和校验语义唯一差异在于expert 输出回传阶段的数据搬运方式经典版本combine_classic数据面与控制面全部走 MTE双平面版本combine_doubleplane数据面按 segment 大小在MTE 与 SDMA之间自适应选择控制面status ready 信号仍走 MTE。示例目录结构见 examples/combine/README.mdexamples/combine/ ├── scripts/ │ ├── data_gen.py # classic 与 doubleplane 共用的数据生成脚本 │ └── check_combine.py # classic 与 doubleplane 共用的结果校验脚本 ├── combine_classic/ │ ├── main.cpp │ ├── combine_kernel.cpp │ └── scripts/run.sh └── combine_doubleplane/ ├── main.cpp ├── combine_doubleplane_kernel.cpp ├── combine_kernel.h └── scripts/run.sh注意当前双平面版本暂不支持 Ascend950平台配套编译运行。二、为什么使用双平面经典 combine 中所有 expert 输出回传都走 MTE。对于小 shape 或均匀路由的场景MTE-only 的固定开销较低但当某些来源 rank 对应的 expert 输出较多时回传 payload 会形成大批量远端写此时 MTE 固定开销占比升高。双平面设计的目标可以概括为四点小段继续走 MTE避免 SDMA 的固定开销issue/event/quiet 成本大段走 SDMA提高大 payload 远端写效率status ready 等控制信号仍走 MTE保证同步协议简单可靠且 status 不会早于 payload 可见输出x_out与经典 combine 完全一致便于功能校验和性能对比。适合优先尝试双平面的典型场景bs、h、topk较大combine 回传总数据量较高dispatch 后的ep_recv_count显示部分(local_expert_id, src_rank_id)segment 明显偏大关注通信阶段comm_only性能希望比较 MTE-only 与 MTESDMA 的差异。需要同时注意反向场景如果每个 segment 都较小SDMA 的 issue/event/quiet 成本可能抵消收益此时应回归经典 MTE-only 路径。三、功能说明输入、输出与校验公式3.1 输入输入包括dispatch 阶段的输出输入含义expand_xdispatch 后按 expert 聚合的 token 数据。本示例的数据生成脚本使用 identity expert 计算因此expand_x可直接视为 expert 输出assist_info_for_combine每条 expert 输出对应的来源[src_rank_id, src_token_id, src_topk_id]ep_recv_countdispatch 阶段生成的 segment 累计接收计数以及combine 自身的输入输入含义expert_ids每个 token/topK 对应的 expert idexpert_scales每个 token/topK 的加权系数3.2 输出x_out每个 token 的最终 combine 结果。3.3 校验公式x_out[token] sum(topk_output[token, topk] * expert_scales[token, topk])该公式在数据生成脚本 examples/combine/scripts/data_gen.py 的combine_golden()中实现对每个 rank 的每个 token、每个 topK将 token 的原始向量与对应expert_scales相乘后累加作为 golden 输出写入golden_x_out.bin。数据生成细节同一脚本中dispatch_golden()expert_ids在[0, moe_expert_num)内随机生成moe_expert_num pe_size * expert_per_peexpert_scales在[0.05, 1.0]区间均匀采样int32 时系数固定为 1.0expand_x生成方式与类型相关int32使用[-5, 5]整数随机float16_t使用[-1.0, 1.0]均匀浮点随机。四、方案设计双平面的核心决策逻辑4.1 Host 侧初始化Host 初始化 SHMEM 时同时启用 MTE 与 SDMA 两种数据通路ACLSHMEM_DATA_OP_MTE | ACLSHMEM_DATA_OP_SDMA在 main.cpp 中可以看到完整初始化序列先通过test_set_attr填充aclshmemx_init_attr_t属性再设置data_op_engine_type为 MTE 与 SDMA 的组合然后调用aclshmemx_set_qp_num(ACLSHMEM_DATA_OP_SDMA, pe_size)为每个 AIV 显式配置 SDMA QP 数量数量等于 PE 数最后执行aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, attributes)。4.2 Kernel 侧 SDMA 判定逻辑Kernel 仍然启动pe_size个 AIV core每个 active core 负责一个来源 rank。对每个(src_rank, local_expert)segment先通过ep_recv_count得到该 segment 的 token 数再判断是否启用 SDMA。SDMA 判定逻辑可以理解为segment_bytes token_count * h * sizeof(T) use_sdma src_rank ! my_rank token_count 0 segment_bytes 2MB segment_bytes 当前 PE 的远端平均 segment 字节数代码中为了避免整数除法截断使用交叉相乘避免除法截断误差segment_bytes * threshold_den threshold_num其中threshold_num remote_tokens * h * sizeof(T) threshold_den max(pe_size - 1, 1) * local_expert_num也就是说只有“大于 2MB 且大于当前 PE 远端平均段大小”的远端大段才走 SDMA本地段、空段、小段和普通段仍走 MTE。以上逻辑对应 kernel 中的实现 combine_doubleplane_kernel.cppCombineUseSdma()依次检查peer_rank my_rank、token_count 0、threshold_den 0直接返回 false随后计算bytes token_count * h * elem_size若bytes MIN_SDMA_BYTES常量定义为2 * 1024 * 1024即 2MB返回 false最后返回bytes * threshold_den threshold_num。CombineThresholdNum()遍历所有非本地 segment累加remote_tokens再乘以h * sizeof(T)得到阈值分子CombineThresholdDen()计算(pe_size 1 ? pe_size - 1 : 1) * local_expert_num得到阈值分母。五、实现逻辑MTE 与 SDMA 的分工双平面 combine 的主要阶段如下与 combine_doubleplane_kernel.cpp 的 Stage 1 实现一一对应遍历 segment每个 active core 根据ep_recv_count遍历自己负责来源 rank 的所有 local expert segmentsegment 索引为local_expert * pe_size src_ranksegment 的 token 区间由ep_recv_count的相邻累计值界定。计算阈值并判定计算当前 PE 的远端平均 segment 字节数并对每个 segment 判断use_sdma。先提交 SDMA payload对大段使用aclshmemx_sdma_qp_put_nbi非阻塞写回来源 PE 的 combine windowkernel 中以data_base slot * data_stride为远端目标地址expand_x i * h为本地源地址经sdma_tmpUB 缓冲区中转。限制 outstanding 数量每提交256 次SDMA issue 后调用aclshmemx_sdma_qp_quiet避免 outstanding 请求无限积压常量SDMA_ISSUE_LIMIT 256。小段走 MTE 直连对小段使用aclshmemx_mte_put_nbi写 payload并立即写 status ready。SDMA 段统一完成后再写 statusSDMA payload 阶段完成后统一sdma_quiet再用 MTE 为 SDMA 段写 status ready。来源 PE 等待 status每个 token 的全部 topK status ready 后继续。加权归约与收尾按expert_scales做加权求和写出x_out再清理 status写回 0并通过aclshmemi_sync_core_soft()同步退出。关于 Status 与数据可见性的关键设计kernel 注释明确说明SDMA 只负责大段 payload 数据面status ready 控制面仍由 MTE 写入。这样可以保证 status 不会早于 payload 可见避免来源 PE 读到未写完的数据。其他值得注意的 kernel 常量combine_doubleplane_kernel.cppUB_DMA_MAX_SIZE 190 * 1024MTE 直连路径的单次最大搬运字节数SDMA_UB_SIZE 64、SDMA_UB_OFFSET 191 * 1024SDMA 的 UB 中转缓冲区大小与偏移COMBINE_ASSIST_FIELDS 3assist 元数据字段数对应[src_rank_id, src_token_id, src_topk_id]COMBINE_ALIGN_BYTES 32payload 行对齐字节数COMBINE_STATUS_READY 1status ready 标记值。六、与经典 Combine 的关系可对比基线双平面版本与经典版本保持输入输出完全一致combine_doubleplane/README.md消费相同的expand_x、assist_info_for_combine、ep_recv_count使用相同的expert_ids和expert_scales产生相同语义的x_out使用相同的 golden/check 脚本验证正确性data_gen.py 与 check_combine.py。因此推荐的工作流是先用经典 combine 建立正确性基线再用双平面对相同 shape 做性能对比。经典版本的完整说明见 combine_classic/README.md。七、环境要求与平台限制双平面路径依赖 SHMEM 的 SDMA 能力。结合 examples/sdma/README.md 可以确认SDMA 功能要求 CANN 9.0.0 及以上并需要安装匹配硬件平台的 toolkit 和 ops-legacy 软件包SDMA put/get 接口在 Atlas 200I A2/A3、Atlas 300T A2/A3 等 A2/A3 平台可用Ascend950 仅支持 SDMA get不支持 SDMA put这也是双平面示例当前不支持 Ascend950 的原因之一基础安装与独立 SDMA demo 可参考 examples/sdma/README.md。八、构建在仓库根目录执行bash scripts/build.sh -examples若需要先编译并安装 SHMEM 软件包SDMA 依赖可参考 examples/sdma/README.md 的步骤bash scripts/build.sh -package ./install/*/SHMEM_1.0.0_linux-*.run --install bash scripts/build.sh -examples九、运行9.1 基础 2 卡测试cd examples/combine/combine_doubleplane bash scripts/run.sh -pes 2 -bs 8 -h 16 -topk 2 -expertPerPe 2 -type int32_t9.2 8 卡、64 expert 测试cd examples/combine/combine_doubleplane bash scripts/run.sh -pes 8 -bs 8 -h 16 -topk 2 -expertPerPe 8 -type int32_t脚本run.sh会自动调用父目录公共脚本../scripts/data_gen.py生成 combine 输入和 golden 输出启动每个 PE 对应的combine_doubleplane进程通过build/bin/combine_doubleplane二进制按pe_id绑定 NPU输出写入output/x_out_rank.bin调用父目录公共脚本../scripts/check_combine.py校验结果。校验逻辑check_combine.pyint32_t使用精确相等比较浮点类型使用rtol1e-2, atol1e-2的np.allclose比较。main.cpp 侧的CheckArray同样对fp16_t使用1e-2绝对误差阈值main.cpp。9.3 常用参数-pes n PE/NPU 数量单机示例要求 -gnpus 与 -pes 相同。 -bs n 每个 PE 的 token 数。 -h n token hidden size默认 7168。 -topk n 每个 token 路由的 expert 数默认 8。 -expertPerPe n 每个 PE 上的 local expert 数。 -type dtype 数据类型支持 int32_t、float16_t。 -fnpu id 起始 NPU id默认 0。 -ipport url SHMEM bootstrap 地址默认 tcp://127.0.0.1:8767。bfloat16_t当前未在 combine doubleplane 示例中实例化。原因是 CANN 9.0 beta 后端不支持 combine 累加路径需要的标量 bf16 castrun.sh会在启动前主动拒绝-type bfloat16_t见 run.sh。9.4 如何选择使用推荐用同一组 shape对比经典版和双平面cd examples/combine/combine_classic bash scripts/run.sh --perf -pes 8 -bs 32 -h 7168 -topk 8 -expertPerPe 8 -type int32_t --warmup 5 --loops 50 cd ../combine_doubleplane bash scripts/run.sh --perf -pes 8 -bs 32 -h 7168 -topk 8 -expertPerPe 8 -type int32_t --warmup 5 --loops 50优先观察comm_only。如果comm_only降低说明大段回传走 SDMA 对通信阶段有效如果full_op收益较小需要结合 status wait、加权归约和同步开销一起分析。十、性能测试10.1 单 shape profilingcd examples/combine/combine_doubleplane bash scripts/run.sh --perf -pes 2 -bs 8 -h 256 -topk 2 -expertPerPe 2 -type int32_t \ --warmup 5 --loops 5010.2 多 shape、多卡数 sweepcd examples/combine/combine_doubleplane 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 50run.sh对 list 参数会逐组合并遍历执行--pes-list、--bs-list、--h-list、--topk-list、--expert-per-pe-list并为每个 case 分配独立的 bootstrap/session 端口避免冲突。10.3 CSV 指标说明每个 case 会写出两个指标行见 main.cpp 的AppendPerfCsvRowsfull_op完整 combine包括回传通信、status wait、加权归约、清理和同步对应FULL_FRAME_ID 0打点区间comm_onlyStage 1 回传通信及必要的完成/status 协议对应COMM_FRAME_ID 1打点区间。单 rank 文件名为combine_doubleplane_perf_rankrank.csv。使用--prof-pe all时脚本会轮流 profile 每个 PE最终调用 examples/utils/summarize_moe_perf.py 生成combine_doubleplane_perf_summary.csv。CSV 列结构如下DataSize/B,Npus,Blocks,UBsize/KB,Bandwidth/GB/s,CoreMaxTime/us,Metric,GlobalDataSize/B, PerPeBandwidth/GB/s,BS,H,TopK,ExpertPerPe,Dtype,Warmup,Loops,ProfPe,CaseId,SingleCoreTime/us其中前六列与 examples/utils/perf_data_process.py 兼容耗时统计基于SHMEM_CYCLE_PROF_PE环境变量控制的打点数据周期到微秒的换算系数由GetCycleToUs()决定Ascend950 为 1000其余平台为 50见 main.cpp。性能测试结束后还可使用--analyse plot或--analyse md调用统一性能报告脚本生成图表或 Markdown 报告。十一、卡数限制与注意事项AIV/QP 上限示例按pe_size启动相同数量的 AIV 并配置同数 SDMA QP底层最多支持 72 个 AIV/QPACLSHMEM_MAX_AIV_PER_NPU因此pe_size不应超过 72见 examples/combine/README.md。性能采集时的块槽位限制开启性能采集设置环境变量SHMEM_CYCLE_PROF_PE时SHMEMI_PROF打点数组仅有 64 个 block 槽位ACLSHMEM_CYCLE_PROF_MAX_BLOCK因此请确保pe_size不超过 64超出部分的 AIV 打点数据无法记录可能导致性能数据不完整。不开启性能采集时不受此限制main.cpp 的GetCoreTimesUs中也以min(block_num, ACLSHMEM_CYCLE_PROF_MAX_BLOCK)截断读取与之一致。SDMA QP 配置需要多核并发时应在aclshmemx_init_attr之前显式调用aclshmemx_set_qp_num(ACLSHMEM_DATA_OP_SDMA, qp_num)未调用时仅创建一个 SDMA stream/QP。平台能力双平面依赖 SDMA put当前不支持仅支持 SDMA get 的 Ascend950 平台运行时需保证 toolkit 与 ops-legacy 版本匹配硬件平台。十二、总结双平面 combine 是 MoE 大规模路由场景下优化 expert 输出回传带宽的一个务实方案它没有改变 combine 的数学语义和同步协议而是把数据面拆成 MTE 与 SDMA 两条路径用“2MB 且超过远端平均段大小”的简单阈值完成自适应分流并把 status 控制面始终保留在 MTE 上以保证数据可见性。结合 combine_classic 基线做同 shape 对比通过comm_only/full_op两个指标即可快速判断 SDMA 分流在目标 workload 上的收益为后续算子级调优提供了可复现的工程范式。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考