MoE Dispatch 通信算子基于 PTO-ISA 的 MegaMoE Token 分发内核实现指南Ascend A5【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa导读本文围绕 CANN pto-isa 仓库中kernels/manual/a5/moe_dispatch目录下的独立通信算子展开完整讲解如何用 PTO-ISA 指令集在 Ascend A5 上实现 MegaMoEMixture of Experts的 Dispatch 阶段从远端 rank 的共享内存HCCL window中拉取量化 token把交织的[int8 token | float scale]行拆分为紧凑的gmA与gmPerTokenScale两个独立输出。读完本文你将掌握 Direct / ViaGM / WithSync 三条 kernel 路径的设计差异、事件驱动的 ping-pong 流水与软件 SYNCALL 同步原理、run.sh 的完整构建运行方式以及该算子作为 MegaMoE 融合算子构建模块的接口兼容性。算子背景MegaMoE 中的 Dispatch 阶段在 MoEMixture of Experts模型中每个 token 需要被路由到指定的专家Expert进行计算。MegaMoE 将路由Routing、分发Dispatch、专家前馈GEMM/FFN与合并Combine融合为一个整体流水线InitRouting → [Dispatch] → GEMM (FFN) → Combine ^^^^^^^^ 本文算子WithSync 路径覆盖 InitRouting Dispatchkernels/manual/a5/moe_dispatch目录下的算子即为其中 Dispatch 环节的独立standalonePTO-ISA 实现。它把各 rank 在InitRouting阶段写入共享内存的量化 token 数据按路由表拉取到本地并完成 token 数据与 per-token scale 的分离为后续 GEMM 提供紧凑、无 padding 的输入。三条 kernel 路径总览该算子提供了三条相互独立的 kernel 路径覆盖追求性能与追求兼容两种诉求路径步数数据流定位Direct--mode direct2 步TLOAD 远端 GM → UB → TSTORE 拆分快速路径自适应 UB tilingViaGM--mode viagm4 步TGET 远端 GM → 本地 GM → TLOAD → UB → TSTORE 拆分MegaMoE 兼容路径WithSync--mode sync集成式CrossRankSync → Direct dispatch设备侧路由表计算 dispatch完全自包含三条路径的核心差异在于远端数据如何进入本地以及路由表从哪里来Direct直接用TLOAD从远端 rank 的 HCCL 共享内存GM 地址把交织行读入本地 UB是最短路径ViaGM先用 PTO-ISA 通信指令TGET定义见 include/pto/comm/pto_comm_inst.hpp把远端数据搬到本地一块临时 GM buffer再走与 Direct 相同的TLOAD → UB → TSTORE流程与 MegaMoE 原有实现DispatchCopyPerToken行为一致WithSync在同一个 kernel 内先完成跨 rank 的tokenPerExpert全收集与路由表计算CrossRankSync再调用 Direct dispatch 主循环全程无需 host 预计算路由表。核心算法解析Direct / ViaGM 主循环逐专家、逐远端 rank 拉取两条路径的主循环结构完全一致Direct 实现于 moe_dispatch_kernel.cpp 的MoeDispatchDirectViaGM 对应MoeDispatchViaGMfor each 本地专家 (groupIdx): for each 远端 rank (dstEpIdx, 按 coreIdx 跨步): 1. 计算远端 shmem 中的源地址 2. 计算本地目标偏移 (gmA/gmPerTokenScale) 3. [Direct] TLOAD 交织行到 UBTSTORE 拆分 token 和 scale [ViaGM] TGET 行到本地 GM再 TLOAD → UB → TSTORE 拆分 4. 事件驱动 ping-pongTLOAD(N1) 与 TSTORE(N) 重叠执行 // 跨 rank 连续流水rank 边界无气泡关键细节多核并行外层for (dstEpIdx coreIdx; dstEpIdx EP; dstEpIdx coreNum)把 EP 个远端 rank 按get_block_idx()/get_block_num()跨步分配给多个 AIV 核每个核处理一个或多个远端 rank边界裁剪rowStart依据cumsumMM前缀和表计算超过maxOutputSize直接跳过单次拉取行数rows也会被maxOutputSize - rowStart裁剪逐块搬运每批最多搬运MOVE_NUM行processCount (rows MOVE_NUM - 1) / MOVE_NUM末批处理余数rem。远端地址换算CommRemotePtrDirect/ViaGM 路径都需要把本地视角的 shmem 指针换算成远端 rank 的地址。kernel 中通过CommRemotePtr实现moe_dispatch_kernel.cpptemplate typename T AICORE inline __gm__ T* CommRemotePtr(__gm__ CommDeviceContext* ctx, __gm__ T* localPtr, int pe) { uint64_t localBase ctx-windowsIn[ctx-rankId]; uint64_t offset (uint64_t)localPtr - localBase; return (__gm__ T*)(ctx-windowsIn[pe] offset); }其原理是每个 rank 的 HCCL window 基地址记录在CommDeviceContext::windowsIn[]数组中windowsIn[rankId]为本地 windowwindowsIn[pe]为远端 pe 的 window。先用本地指针减去本地 window 基地址得到相对偏移再以该偏移叠加到目标 rank 的 window 基地址上即得到远端视角的全局地址。CommDeviceContext结构体定义在 comm_context.h其二进制布局与 A5Ascend 910 A5DAV_3510上HcclAllocComResourceByTiling()返回的HcclCombinOpParamA5对齐windowsIn/Out数组长度为HCCL_MAX_RANK_NUM 64A3 平台为 32这是 A5 的一个差异点。WithSync设备侧 CrossRankSync DispatchWithSync 路径把 MegaMoE 的CrossRankSyncAndlocalTokenPerExpertAllGatherAndGetSumPreRankV2的 dispatch 部分整体搬进 kernelMoeDispatchWithSync分三个阶段Phase A — TPE AllGather for each 远端 rank i: TSTORE 本地 tokenPerExpert 到 rank i 的 TPE 交换区附加 DataAsFlag 偏移 for each 远端 rank i: TWAIT 等待 rank i 数据到达轮询 GM 信号非零 Phase B — 路由表计算 B.1TLOAD 每行 TPETADDS 去除 DataAsFlag 偏移TSTORE 回写 SYNCALL基于 GM 轮询的软件同步 B.2向量化前缀和 — TLOAD row[i]TADD 累加器TSTORE → cumsumMM B.3标量循环 — 从 cumsumMM 列累加 preSumBeforeRank[i] Phase C — Dispatch 使用计算好的路由表调用 MoeDispatchDirectPhase ADataAsFlag 机制为避免先置零信号、再写数据两步导致的竞态Phase A 采用DataAsFlag技巧把本地tokenPerExpert行的每个 int32 值统一加上DATA_AS_FLAG_OFFSET 0x800000定义见 moe_dispatch_config.h然后通过TSTORE直接写入远端 rank 的 TPE 交换区。远端 rank 只需用TWAIT(signal, 0, WaitCmp::NE)轮询该地址是否非零即可感知数据已到达——数据本身即同步信号无需额外的 flag 内存。Phase A 的写入路径是TLOAD本地 TPE 行 →set_flag(PIPE_MTE2, PIPE_V)/wait_flag→TADDS叠加偏移 →set_flag(PIPE_V, PIPE_MTE3)→ 对每个远端 rankTSTORE。TWAIT是 PTO-ISA 通信指令集中的信号同步原语include/pto/comm/pto_comm_inst.hpp按WaitCmp比较条件阻塞等待信号。Phase B路由表计算Phase B 只由coreIdx 0的核执行其他核在软件 SYNCALL 处等待B.1 去偏移逐行TLOAD各源 rank 的 TPE 行对非本 rank 的数据用TADDS减去DATA_AS_FLAG_OFFSET写回 workspace 中的wsTPE区域B.2 向量化前缀和维护 UB 中的累加 TileTLOAD row[i]后TADD累加器再TSTORE到wsCumsumMM i * paddedExpNum得到各专家各 rank 的累计 token 数B.3 标量累加preSumBeforeRank[srcRank * expertPerRank g]通过标量循环累加得出期间对 TPE 值使用dcci缓存行失效指令保证读到最新 GM 值。其中paddedExpNum (EP * expertPerRank 7) ~7TPE 行被 padding 到 8 的倍数以满足 DMA 搬运的 32B 对齐要求。workspace 布局在代码注释与SyncWorkspaceBytes中均有明确约定[0 .. EP*paddedExpNum) : cumsumMM完整行 [EP*paddedExpNum .. EP*expertPerRank) : preSumBeforeRank [EP*paddedExpNum EP*expertPerRank .. end) : tokenPerExpertpaddedPhase C软件 SYNCALLPhase C 通过pto::SYNCALLpto::SyncAllMode::Soft(syncGmG)实现跨核屏障SyncAllMode定义于 include/pto/common/type.hpp软同步 workspace 大小为SYNCALL_SOFT_WORKSPACE_INT32 16个 int32即一个缓存行。软 SYNCALL 基于 GM 轮询的到达计数器实现刻意避开对 FFTS 硬件同步的依赖这也是 A5 上的一个工程考量。关键实现机制源码级自适应 MOVE_NUMDispatchTraitsUB 容量有限A5 单核 UB 共 256KBkernel 中按半区UB_HALF_SIZE 128 * 1024规划 ping/pong 双 buffer。MOVE_NUM每批搬运行数由编译期模板DispatchTraitsTILE_COLS自动推导moe_dispatch_config.htemplate int TILE_COLS struct DispatchTraits { static constexpr int32_t MAX_ROWS UB_HALF_SIZE / TILE_COLS; static constexpr int32_t MOVE_NUM (MAX_ROWS UB_MOVE_NUM_MAX) ? UB_MOVE_NUM_MAX : (MAX_ROWS 1 ? MAX_ROWS : 1); };即MOVE_NUM取UB 半区能容纳的行数与上限UB_MOVE_NUM_MAX 16的较小值。当hiddenSizeTILE_COLS增大时MAX_ROWS自动收缩从而在保证不溢出 UB 的前提下尽可能保持搬运吞吐。当前 kernel 入口固定实例化为MoeDispatchDirect128, 160, DispatchTraits160::MOVE_NUM即 TILE_COLS160128 字节 token 32 字节 scale 对齐行。事件驱动 ping-pongMTE2 与 MTE3 重叠Direct/ViaGM 的拆分循环使用两个 UB Tile*Ping/*PongPONG 偏移(MOVE_NUM * TILE_COLS 31) ~31做 32B 对齐轮流作为 TLOADMTE2与 TSTOREMTE3的目标配合set_flag/wait_flag事件实现软件流水TSTORE(pendTokenDstG, prevTokenView); // 存上一批的 token TSTORE(pendScaleDstG, prevScaleView); // 存上一批的 scale TLOAD(loadTile, remoteSrcG); // 同时加载下一批 set_flag(PIPE_MTE3, PIPE_MTE2, prevEvent); set_flag(PIPE_MTE2, PIPE_MTE3, curEvent); wait_flag(PIPE_MTE3, PIPE_MTE2, prevEvent);核心逻辑是TLOAD(N1)与TSTORE(N)使用不同 eventEVENT_ID0/EVENT_ID1分别挂载到不同的 ping/pong 缓冲区互不冲突从而重叠执行。loop 结束后用最后一次wait_flag冲刷尾部。跨 rank 连续流水无气泡与每个远端 rank 内部做 ping-pong、rank 边界刷新的朴素实现不同本算子将 ping-pong 状态hasPending、pendingPP、pendingRows及目标指针提升到dstEpIdx循环之外使上一 rank 的最后一批 TSTORE 与下一 rank 的首批 TLOAD 仍然重叠rank 边界处不产生流水气泡。这是 Direct 路径相对朴素实现的主要性能优化点之一。Token/Scale 分离的数据布局远端共享内存中每一行是交织格式hiddenSize字节的 int8 token 数据 UB_ALIGN 32字节 paddingfloat scale 位于行内偏移hiddenSize处即 padding 区偏移 0。行跨度为ShmemRowStride(hiddenSize) hiddenSize 32moe_dispatch_config.h。kernel 在 UB 中建立三个视图interleaved、tokenView、scaleView三者共享同一块 UB 内存、偏移不同tokenViewPing偏移PING_OFFSET列宽HIDDEN_SIZEscaleViewPing偏移PING_OFFSET HIDDEN_SIZE列宽UB_ALIGN。因此一次TLOAD载入交织行后无需任何向量运算即可通过两次TSTORE分别输出紧凑 tokengmA[maxOutputSize, hiddenSize]无 padding与 scalegmPerTokenScale[maxOutputSize]每行 32 字节、float 在偏移 0 处。拆分本身零计算只有搬运。规格说明项目值数据类型tokenint8_t数据类型scalefloat存储在 32B 对齐行中远端行格式hiddenSize字节 token UB_ALIGN32字节 paddingscale 在偏移 0 处输出 tokengmA[maxOutputSize, hiddenSize]— 紧凑无 padding输出 scalegmPerTokenScale[maxOutputSize]— 每行 32 字节float 在偏移 0 处默认 hiddenSize128支持的 AI 处理器Ascend A5执行模型仅 AIV向量核通过 mpirun 多 rankA5 平台注意事项HCCL 窗口偏移勘误A5 的 MTE2 DMA 引擎读取 HCCL window 基地址偏移 [16..31] 字节时返回全零。因此 host 驱动在 main.cpp 中通过WindowAlloc(ctx.hostCtx.windowsIn[rankId], winOffset, totalShmemAlloc)以winOffset 256跳过缺陷区域保证共享内存数据完整HCCL V2 TilingA5 使用 V2 tiling 初始化路径复用 tests/npu/a5/comm/st/testcase/common.hpp 中的 comm 测试框架TestContext、WindowAlloc、HcclHostBarrier等编译器目标kernel 以--cce-aicore-archdav-c310-vec编译A5 向量核 ISAhost 侧编译器为bisheng-DRUN_MODEnpu选择 NPU 运行时CMakeLists.txt设备上下文差异A5 的CommDeviceContext含 CCU 寄存器字段xnAddr/ckeAddr/msAddr/msSize本 dispatch kernel 不使用这些寄存器仅保留以对齐二进制布局。目录结构与源码导读kernels/manual/a5/moe_dispatch/ ├── moe_dispatch_kernel.cpp # 设备 kernel三路径 dispatchDirect/ViaGM/WithSync 及 __global__ 入口 ├── moe_dispatch_launch.h # Host 侧 launch wrapper 声明LaunchMoeDispatchK128 等 ├── main.cpp # Host 驱动MPI 初始化、路由表生成、数据生成、launch、golden 验证 ├── moe_dispatch_config.h # 形状常量、DispatchTraits、workspace 布局、DataAsFlag 常量 ├── comm_context.h # 设备侧 CommDeviceContext 结构体HCCL window 基地址等 ├── CMakeLists.txt # 构建配置bisheng dav-c310-vec内核与 host 可执行分别构建 ├── run.sh # 构建和运行便捷脚本 ├── README.md # 英文说明 └── README_zh.md # 中文说明值得深入阅读的对应关系三路径算法主体MoeDispatchDirect、MoeDispatchViaGM、MoeDispatchWithSyncmoe_dispatch_kernel.cpp入口为MoeDispatchDirect_K128/MoeDispatchViaGM_K128/MoeDispatchWithSync_K128三个extern C __global__函数编译期形状与 workspace 布局DispatchTraits、TPEAreaBytes、SyncWorkspaceBytes、DATA_AS_FLAG_OFFSETmoe_dispatch_config.hPTO 通信指令集 APITGET、TWAIT、TNOTIFY等定义于 include/pto/comm/pto_comm_inst.hpp软同步模式SyncAllMode::Soft与SYNCALL_SOFT_WORKSPACE_INT32定义于 include/pto/common/type.hpp。构建与运行环境准备# 设置 CANN 环境脚本路径按实际部署调整 source /mnt/data/ntlab/liulei/set_env_new.sh export HCCL_WHITELIST_DISABLE1run.sh会检查ASCEND_HOME_PATH环境变量未设置时报错退出run.sh。构建产物为build/moe_dispatchhost 可执行与moe_dispatch_kernel内核动态库。常用命令# 构建并运行 Direct 路径默认2 卡 bash run.sh all --ep 2 --mode direct # 构建并运行 ViaGM 路径4 卡 bash run.sh all --ep 4 --mode viagm # 构建并运行 WithSync 路径CrossRankSync Dispatch2 卡 bash run.sh all --ep 2 --mode sync # 指定设备从 4 号卡开始 bash run.sh all --ep 4 --first-device 4 --mode direct # 仅构建 bash run.sh build --ep 2 --hidden 128 --debug # 仅运行构建后 bash run.sh run --ep 2 --mode direct # 清理重新构建 bash run.sh all --ep 4 --mode viagm --cleanrun.sh 参数说明参数默认值说明--ep N2rank 数量EP 数--mode direct\|viagm\|syncdirectkernel 路径选择--first-device N0起始 NPU 设备号--hidden N128隐藏层维度K--tokens N64每 rank 最大 token 数--max-output N512最大输出行数--experts N1每 rank 专家数--clean—强制清理重建--debug—启用调试模式参数通过 cmake 变量CONFIG_EP、CONFIG_EXPERT_PER_RANK、CONFIG_HIDDEN_SIZE、CONFIG_MAX_TOKENS_PER_RANK、CONFIG_MAX_OUTPUT_SIZE、CONFIG_FIRST_DEVICE_ID编译期注入CMakeLists.txt 与 moe_dispatch_config.h运行阶段run.sh通过DISPATCH_MODE环境变量把路径模式传给 host 程序main.cpp 中std::getenv(DISPATCH_MODE)解析最终以mpirun -n ${EP} ./moe_dispatch拉起多 rank。CMake 构建要点kernel 目标moe_dispatch_kernel使用CMAKE_CCE_COMPILE_OPTIONS含-xcce、-mllvm系列 AICore 栈与地址变换选项与--cce-aicore-archdav-c310-vec、--cce-fatobj-link编译host 目标moe_dispatch链接moe_dispatch_kernel以及ascendcl、hcomm、tiling_api、register、platform等 CANN 运行时库RUN_MODEnpu时链接runtimeRUN_MODEsim时则链接runtime_camodelinclude 路径中 PTO 头目录仓库根目录include置于最前以覆盖 CANN 自带同名头文件A5 comm 测试框架路径tests/npu/a5/comm/st/testcase亦被加入 include 搜索路径。验证机制golden 对比main.cpp 是一个完整的 MPI HCCL 多 rank 测试驱动其验证流程为以固定 seedseed 42在各 rank 上确定性生成路由表RoutingTables::GeneratetokenPerExpert[src][dst][g]、cumsumMM、preSumBeforeRank以seed rankId * 1000生成各 rank 的合成 shmem 数据GenerateShmemDataint8 token 随机填充 float scale 写入行内偏移 hiddenSize 处在 host 侧用ComputeGolden模拟 Dispatch 主循环得到期望的expectedGmA与expectedGmScale启动目标路径 kernel回读gmA与gmPerTokenScale逐行逐列与 golden 对比输出[PASS] Rank N: verified M rows或逐行[FAIL]报告。值得注意的是RoutingTables::Generate与 kernel 内MoeDispatchDirect的寻址逻辑rowStart、preSumBeforeRank行偏移一一对应是理解路由表语义的最佳阅读入口。与 MegaMoE 的关系该算子定位为 MegaMoE Dispatch 阶段的独立验证实现可直接作为完整 MegaMoE 融合算子的构建模块接口兼容kernel 参数cumsumMM、tokenPerExpert、preSumBeforeRank、shmemBase与 MegaMoE 完全对齐WithSync 路径等价于 MegaMoE 的CrossRankSyncAndlocalTokenPerExpertAllGatherAndGetSumPreRankV2DispatchAndCombine的 dispatch 部分即一个 kernel 内完成路由同步与分发ViaGM 路径功能上等价于 MegaMoE 的DispatchCopyPerToken行为兼容、便于对照验证Direct 路径PTO-ISA 独有优化绕过中间 GM buffer用 TLOAD 直接跨 rank 读取并完成拆分。从源码结构看README 的 Reference 小节其设计参考了 MegaMoE 源码vllm-ascend/csrc/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_kernel.hpp与独立设计文档dispatch_pto_isa_design.mdPTO-ISA 通信指令 API 以 include/pto/comm/pto_comm_inst.hpp 为准。【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 MoE Dispatch 通信算子：基于 PTO-ISA 的 MegaMoE Token 分发内核实现指南（Ascend A5）
MoE Dispatch 通信算子基于 PTO-ISA 的 MegaMoE Token 分发内核实现指南Ascend A5【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa导读本文围绕 CANN pto-isa 仓库中kernels/manual/a5/moe_dispatch目录下的独立通信算子展开完整讲解如何用 PTO-ISA 指令集在 Ascend A5 上实现 MegaMoEMixture of Experts的 Dispatch 阶段从远端 rank 的共享内存HCCL window中拉取量化 token把交织的[int8 token | float scale]行拆分为紧凑的gmA与gmPerTokenScale两个独立输出。读完本文你将掌握 Direct / ViaGM / WithSync 三条 kernel 路径的设计差异、事件驱动的 ping-pong 流水与软件 SYNCALL 同步原理、run.sh 的完整构建运行方式以及该算子作为 MegaMoE 融合算子构建模块的接口兼容性。算子背景MegaMoE 中的 Dispatch 阶段在 MoEMixture of Experts模型中每个 token 需要被路由到指定的专家Expert进行计算。MegaMoE 将路由Routing、分发Dispatch、专家前馈GEMM/FFN与合并Combine融合为一个整体流水线InitRouting → [Dispatch] → GEMM (FFN) → Combine ^^^^^^^^ 本文算子WithSync 路径覆盖 InitRouting Dispatchkernels/manual/a5/moe_dispatch目录下的算子即为其中 Dispatch 环节的独立standalonePTO-ISA 实现。它把各 rank 在InitRouting阶段写入共享内存的量化 token 数据按路由表拉取到本地并完成 token 数据与 per-token scale 的分离为后续 GEMM 提供紧凑、无 padding 的输入。三条 kernel 路径总览该算子提供了三条相互独立的 kernel 路径覆盖追求性能与追求兼容两种诉求路径步数数据流定位Direct--mode direct2 步TLOAD 远端 GM → UB → TSTORE 拆分快速路径自适应 UB tilingViaGM--mode viagm4 步TGET 远端 GM → 本地 GM → TLOAD → UB → TSTORE 拆分MegaMoE 兼容路径WithSync--mode sync集成式CrossRankSync → Direct dispatch设备侧路由表计算 dispatch完全自包含三条路径的核心差异在于远端数据如何进入本地以及路由表从哪里来Direct直接用TLOAD从远端 rank 的 HCCL 共享内存GM 地址把交织行读入本地 UB是最短路径ViaGM先用 PTO-ISA 通信指令TGET定义见 include/pto/comm/pto_comm_inst.hpp把远端数据搬到本地一块临时 GM buffer再走与 Direct 相同的TLOAD → UB → TSTORE流程与 MegaMoE 原有实现DispatchCopyPerToken行为一致WithSync在同一个 kernel 内先完成跨 rank 的tokenPerExpert全收集与路由表计算CrossRankSync再调用 Direct dispatch 主循环全程无需 host 预计算路由表。核心算法解析Direct / ViaGM 主循环逐专家、逐远端 rank 拉取两条路径的主循环结构完全一致Direct 实现于 moe_dispatch_kernel.cpp 的MoeDispatchDirectViaGM 对应MoeDispatchViaGMfor each 本地专家 (groupIdx): for each 远端 rank (dstEpIdx, 按 coreIdx 跨步): 1. 计算远端 shmem 中的源地址 2. 计算本地目标偏移 (gmA/gmPerTokenScale) 3. [Direct] TLOAD 交织行到 UBTSTORE 拆分 token 和 scale [ViaGM] TGET 行到本地 GM再 TLOAD → UB → TSTORE 拆分 4. 事件驱动 ping-pongTLOAD(N1) 与 TSTORE(N) 重叠执行 // 跨 rank 连续流水rank 边界无气泡关键细节多核并行外层for (dstEpIdx coreIdx; dstEpIdx EP; dstEpIdx coreNum)把 EP 个远端 rank 按get_block_idx()/get_block_num()跨步分配给多个 AIV 核每个核处理一个或多个远端 rank边界裁剪rowStart依据cumsumMM前缀和表计算超过maxOutputSize直接跳过单次拉取行数rows也会被maxOutputSize - rowStart裁剪逐块搬运每批最多搬运MOVE_NUM行processCount (rows MOVE_NUM - 1) / MOVE_NUM末批处理余数rem。远端地址换算CommRemotePtrDirect/ViaGM 路径都需要把本地视角的 shmem 指针换算成远端 rank 的地址。kernel 中通过CommRemotePtr实现moe_dispatch_kernel.cpptemplate typename T AICORE inline __gm__ T* CommRemotePtr(__gm__ CommDeviceContext* ctx, __gm__ T* localPtr, int pe) { uint64_t localBase ctx-windowsIn[ctx-rankId]; uint64_t offset (uint64_t)localPtr - localBase; return (__gm__ T*)(ctx-windowsIn[pe] offset); }其原理是每个 rank 的 HCCL window 基地址记录在CommDeviceContext::windowsIn[]数组中windowsIn[rankId]为本地 windowwindowsIn[pe]为远端 pe 的 window。先用本地指针减去本地 window 基地址得到相对偏移再以该偏移叠加到目标 rank 的 window 基地址上即得到远端视角的全局地址。CommDeviceContext结构体定义在 comm_context.h其二进制布局与 A5Ascend 910 A5DAV_3510上HcclAllocComResourceByTiling()返回的HcclCombinOpParamA5对齐windowsIn/Out数组长度为HCCL_MAX_RANK_NUM 64A3 平台为 32这是 A5 的一个差异点。WithSync设备侧 CrossRankSync DispatchWithSync 路径把 MegaMoE 的CrossRankSyncAndlocalTokenPerExpertAllGatherAndGetSumPreRankV2的 dispatch 部分整体搬进 kernelMoeDispatchWithSync分三个阶段Phase A — TPE AllGather for each 远端 rank i: TSTORE 本地 tokenPerExpert 到 rank i 的 TPE 交换区附加 DataAsFlag 偏移 for each 远端 rank i: TWAIT 等待 rank i 数据到达轮询 GM 信号非零 Phase B — 路由表计算 B.1TLOAD 每行 TPETADDS 去除 DataAsFlag 偏移TSTORE 回写 SYNCALL基于 GM 轮询的软件同步 B.2向量化前缀和 — TLOAD row[i]TADD 累加器TSTORE → cumsumMM B.3标量循环 — 从 cumsumMM 列累加 preSumBeforeRank[i] Phase C — Dispatch 使用计算好的路由表调用 MoeDispatchDirectPhase ADataAsFlag 机制为避免先置零信号、再写数据两步导致的竞态Phase A 采用DataAsFlag技巧把本地tokenPerExpert行的每个 int32 值统一加上DATA_AS_FLAG_OFFSET 0x800000定义见 moe_dispatch_config.h然后通过TSTORE直接写入远端 rank 的 TPE 交换区。远端 rank 只需用TWAIT(signal, 0, WaitCmp::NE)轮询该地址是否非零即可感知数据已到达——数据本身即同步信号无需额外的 flag 内存。Phase A 的写入路径是TLOAD本地 TPE 行 →set_flag(PIPE_MTE2, PIPE_V)/wait_flag→TADDS叠加偏移 →set_flag(PIPE_V, PIPE_MTE3)→ 对每个远端 rankTSTORE。TWAIT是 PTO-ISA 通信指令集中的信号同步原语include/pto/comm/pto_comm_inst.hpp按WaitCmp比较条件阻塞等待信号。Phase B路由表计算Phase B 只由coreIdx 0的核执行其他核在软件 SYNCALL 处等待B.1 去偏移逐行TLOAD各源 rank 的 TPE 行对非本 rank 的数据用TADDS减去DATA_AS_FLAG_OFFSET写回 workspace 中的wsTPE区域B.2 向量化前缀和维护 UB 中的累加 TileTLOAD row[i]后TADD累加器再TSTORE到wsCumsumMM i * paddedExpNum得到各专家各 rank 的累计 token 数B.3 标量累加preSumBeforeRank[srcRank * expertPerRank g]通过标量循环累加得出期间对 TPE 值使用dcci缓存行失效指令保证读到最新 GM 值。其中paddedExpNum (EP * expertPerRank 7) ~7TPE 行被 padding 到 8 的倍数以满足 DMA 搬运的 32B 对齐要求。workspace 布局在代码注释与SyncWorkspaceBytes中均有明确约定[0 .. EP*paddedExpNum) : cumsumMM完整行 [EP*paddedExpNum .. EP*expertPerRank) : preSumBeforeRank [EP*paddedExpNum EP*expertPerRank .. end) : tokenPerExpertpaddedPhase C软件 SYNCALLPhase C 通过pto::SYNCALLpto::SyncAllMode::Soft(syncGmG)实现跨核屏障SyncAllMode定义于 include/pto/common/type.hpp软同步 workspace 大小为SYNCALL_SOFT_WORKSPACE_INT32 16个 int32即一个缓存行。软 SYNCALL 基于 GM 轮询的到达计数器实现刻意避开对 FFTS 硬件同步的依赖这也是 A5 上的一个工程考量。关键实现机制源码级自适应 MOVE_NUMDispatchTraitsUB 容量有限A5 单核 UB 共 256KBkernel 中按半区UB_HALF_SIZE 128 * 1024规划 ping/pong 双 buffer。MOVE_NUM每批搬运行数由编译期模板DispatchTraitsTILE_COLS自动推导moe_dispatch_config.htemplate int TILE_COLS struct DispatchTraits { static constexpr int32_t MAX_ROWS UB_HALF_SIZE / TILE_COLS; static constexpr int32_t MOVE_NUM (MAX_ROWS UB_MOVE_NUM_MAX) ? UB_MOVE_NUM_MAX : (MAX_ROWS 1 ? MAX_ROWS : 1); };即MOVE_NUM取UB 半区能容纳的行数与上限UB_MOVE_NUM_MAX 16的较小值。当hiddenSizeTILE_COLS增大时MAX_ROWS自动收缩从而在保证不溢出 UB 的前提下尽可能保持搬运吞吐。当前 kernel 入口固定实例化为MoeDispatchDirect128, 160, DispatchTraits160::MOVE_NUM即 TILE_COLS160128 字节 token 32 字节 scale 对齐行。事件驱动 ping-pongMTE2 与 MTE3 重叠Direct/ViaGM 的拆分循环使用两个 UB Tile*Ping/*PongPONG 偏移(MOVE_NUM * TILE_COLS 31) ~31做 32B 对齐轮流作为 TLOADMTE2与 TSTOREMTE3的目标配合set_flag/wait_flag事件实现软件流水TSTORE(pendTokenDstG, prevTokenView); // 存上一批的 token TSTORE(pendScaleDstG, prevScaleView); // 存上一批的 scale TLOAD(loadTile, remoteSrcG); // 同时加载下一批 set_flag(PIPE_MTE3, PIPE_MTE2, prevEvent); set_flag(PIPE_MTE2, PIPE_MTE3, curEvent); wait_flag(PIPE_MTE3, PIPE_MTE2, prevEvent);核心逻辑是TLOAD(N1)与TSTORE(N)使用不同 eventEVENT_ID0/EVENT_ID1分别挂载到不同的 ping/pong 缓冲区互不冲突从而重叠执行。loop 结束后用最后一次wait_flag冲刷尾部。跨 rank 连续流水无气泡与每个远端 rank 内部做 ping-pong、rank 边界刷新的朴素实现不同本算子将 ping-pong 状态hasPending、pendingPP、pendingRows及目标指针提升到dstEpIdx循环之外使上一 rank 的最后一批 TSTORE 与下一 rank 的首批 TLOAD 仍然重叠rank 边界处不产生流水气泡。这是 Direct 路径相对朴素实现的主要性能优化点之一。Token/Scale 分离的数据布局远端共享内存中每一行是交织格式hiddenSize字节的 int8 token 数据 UB_ALIGN 32字节 paddingfloat scale 位于行内偏移hiddenSize处即 padding 区偏移 0。行跨度为ShmemRowStride(hiddenSize) hiddenSize 32moe_dispatch_config.h。kernel 在 UB 中建立三个视图interleaved、tokenView、scaleView三者共享同一块 UB 内存、偏移不同tokenViewPing偏移PING_OFFSET列宽HIDDEN_SIZEscaleViewPing偏移PING_OFFSET HIDDEN_SIZE列宽UB_ALIGN。因此一次TLOAD载入交织行后无需任何向量运算即可通过两次TSTORE分别输出紧凑 tokengmA[maxOutputSize, hiddenSize]无 padding与 scalegmPerTokenScale[maxOutputSize]每行 32 字节、float 在偏移 0 处。拆分本身零计算只有搬运。规格说明项目值数据类型tokenint8_t数据类型scalefloat存储在 32B 对齐行中远端行格式hiddenSize字节 token UB_ALIGN32字节 paddingscale 在偏移 0 处输出 tokengmA[maxOutputSize, hiddenSize]— 紧凑无 padding输出 scalegmPerTokenScale[maxOutputSize]— 每行 32 字节float 在偏移 0 处默认 hiddenSize128支持的 AI 处理器Ascend A5执行模型仅 AIV向量核通过 mpirun 多 rankA5 平台注意事项HCCL 窗口偏移勘误A5 的 MTE2 DMA 引擎读取 HCCL window 基地址偏移 [16..31] 字节时返回全零。因此 host 驱动在 main.cpp 中通过WindowAlloc(ctx.hostCtx.windowsIn[rankId], winOffset, totalShmemAlloc)以winOffset 256跳过缺陷区域保证共享内存数据完整HCCL V2 TilingA5 使用 V2 tiling 初始化路径复用 tests/npu/a5/comm/st/testcase/common.hpp 中的 comm 测试框架TestContext、WindowAlloc、HcclHostBarrier等编译器目标kernel 以--cce-aicore-archdav-c310-vec编译A5 向量核 ISAhost 侧编译器为bisheng-DRUN_MODEnpu选择 NPU 运行时CMakeLists.txt设备上下文差异A5 的CommDeviceContext含 CCU 寄存器字段xnAddr/ckeAddr/msAddr/msSize本 dispatch kernel 不使用这些寄存器仅保留以对齐二进制布局。目录结构与源码导读kernels/manual/a5/moe_dispatch/ ├── moe_dispatch_kernel.cpp # 设备 kernel三路径 dispatchDirect/ViaGM/WithSync 及 __global__ 入口 ├── moe_dispatch_launch.h # Host 侧 launch wrapper 声明LaunchMoeDispatchK128 等 ├── main.cpp # Host 驱动MPI 初始化、路由表生成、数据生成、launch、golden 验证 ├── moe_dispatch_config.h # 形状常量、DispatchTraits、workspace 布局、DataAsFlag 常量 ├── comm_context.h # 设备侧 CommDeviceContext 结构体HCCL window 基地址等 ├── CMakeLists.txt # 构建配置bisheng dav-c310-vec内核与 host 可执行分别构建 ├── run.sh # 构建和运行便捷脚本 ├── README.md # 英文说明 └── README_zh.md # 中文说明值得深入阅读的对应关系三路径算法主体MoeDispatchDirect、MoeDispatchViaGM、MoeDispatchWithSyncmoe_dispatch_kernel.cpp入口为MoeDispatchDirect_K128/MoeDispatchViaGM_K128/MoeDispatchWithSync_K128三个extern C __global__函数编译期形状与 workspace 布局DispatchTraits、TPEAreaBytes、SyncWorkspaceBytes、DATA_AS_FLAG_OFFSETmoe_dispatch_config.hPTO 通信指令集 APITGET、TWAIT、TNOTIFY等定义于 include/pto/comm/pto_comm_inst.hpp软同步模式SyncAllMode::Soft与SYNCALL_SOFT_WORKSPACE_INT32定义于 include/pto/common/type.hpp。构建与运行环境准备# 设置 CANN 环境脚本路径按实际部署调整 source /mnt/data/ntlab/liulei/set_env_new.sh export HCCL_WHITELIST_DISABLE1run.sh会检查ASCEND_HOME_PATH环境变量未设置时报错退出run.sh。构建产物为build/moe_dispatchhost 可执行与moe_dispatch_kernel内核动态库。常用命令# 构建并运行 Direct 路径默认2 卡 bash run.sh all --ep 2 --mode direct # 构建并运行 ViaGM 路径4 卡 bash run.sh all --ep 4 --mode viagm # 构建并运行 WithSync 路径CrossRankSync Dispatch2 卡 bash run.sh all --ep 2 --mode sync # 指定设备从 4 号卡开始 bash run.sh all --ep 4 --first-device 4 --mode direct # 仅构建 bash run.sh build --ep 2 --hidden 128 --debug # 仅运行构建后 bash run.sh run --ep 2 --mode direct # 清理重新构建 bash run.sh all --ep 4 --mode viagm --cleanrun.sh 参数说明参数默认值说明--ep N2rank 数量EP 数--mode direct\|viagm\|syncdirectkernel 路径选择--first-device N0起始 NPU 设备号--hidden N128隐藏层维度K--tokens N64每 rank 最大 token 数--max-output N512最大输出行数--experts N1每 rank 专家数--clean—强制清理重建--debug—启用调试模式参数通过 cmake 变量CONFIG_EP、CONFIG_EXPERT_PER_RANK、CONFIG_HIDDEN_SIZE、CONFIG_MAX_TOKENS_PER_RANK、CONFIG_MAX_OUTPUT_SIZE、CONFIG_FIRST_DEVICE_ID编译期注入CMakeLists.txt 与 moe_dispatch_config.h运行阶段run.sh通过DISPATCH_MODE环境变量把路径模式传给 host 程序main.cpp 中std::getenv(DISPATCH_MODE)解析最终以mpirun -n ${EP} ./moe_dispatch拉起多 rank。CMake 构建要点kernel 目标moe_dispatch_kernel使用CMAKE_CCE_COMPILE_OPTIONS含-xcce、-mllvm系列 AICore 栈与地址变换选项与--cce-aicore-archdav-c310-vec、--cce-fatobj-link编译host 目标moe_dispatch链接moe_dispatch_kernel以及ascendcl、hcomm、tiling_api、register、platform等 CANN 运行时库RUN_MODEnpu时链接runtimeRUN_MODEsim时则链接runtime_camodelinclude 路径中 PTO 头目录仓库根目录include置于最前以覆盖 CANN 自带同名头文件A5 comm 测试框架路径tests/npu/a5/comm/st/testcase亦被加入 include 搜索路径。验证机制golden 对比main.cpp 是一个完整的 MPI HCCL 多 rank 测试驱动其验证流程为以固定 seedseed 42在各 rank 上确定性生成路由表RoutingTables::GeneratetokenPerExpert[src][dst][g]、cumsumMM、preSumBeforeRank以seed rankId * 1000生成各 rank 的合成 shmem 数据GenerateShmemDataint8 token 随机填充 float scale 写入行内偏移 hiddenSize 处在 host 侧用ComputeGolden模拟 Dispatch 主循环得到期望的expectedGmA与expectedGmScale启动目标路径 kernel回读gmA与gmPerTokenScale逐行逐列与 golden 对比输出[PASS] Rank N: verified M rows或逐行[FAIL]报告。值得注意的是RoutingTables::Generate与 kernel 内MoeDispatchDirect的寻址逻辑rowStart、preSumBeforeRank行偏移一一对应是理解路由表语义的最佳阅读入口。与 MegaMoE 的关系该算子定位为 MegaMoE Dispatch 阶段的独立验证实现可直接作为完整 MegaMoE 融合算子的构建模块接口兼容kernel 参数cumsumMM、tokenPerExpert、preSumBeforeRank、shmemBase与 MegaMoE 完全对齐WithSync 路径等价于 MegaMoE 的CrossRankSyncAndlocalTokenPerExpertAllGatherAndGetSumPreRankV2DispatchAndCombine的 dispatch 部分即一个 kernel 内完成路由同步与分发ViaGM 路径功能上等价于 MegaMoE 的DispatchCopyPerToken行为兼容、便于对照验证Direct 路径PTO-ISA 独有优化绕过中间 GM buffer用 TLOAD 直接跨 rank 读取并完成拆分。从源码结构看README 的 Reference 小节其设计参考了 MegaMoE 源码vllm-ascend/csrc/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_kernel.hpp与独立设计文档dispatch_pto_isa_design.mdPTO-ISA 通信指令 API 以 include/pto/comm/pto_comm_inst.hpp 为准。【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考