MoE Dispatch — PTO-ISA 独立通信算子【免费下载链接】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概述基于 PTO-ISA 指令集实现的 MegaMoE Dispatch 独立通信算子。 该算子从远端 rank 的共享内存中拉取量化 token将交织格式的[int8 token | float scale]行拆分为紧凑的独立输出gmA和gmPerTokenScale。提供三条独立的 kernel 路径Direct2步TLOAD 远端 GM → UB → TSTORE 拆分— 自适应 UB tiling 的快速路径ViaGM4步TGET 远端 GM → 本地 GM → TLOAD → UB → TSTORE 拆分— 兼容 MegaMoE 的路径WithSync集成CrossRankSync → Direct dispatch— 设备侧路由表计算 dispatch支持的 AI 处理器Ascend A5数据流Direct 路径modedirect 远端 GM ──TLOAD──▶ UB (ping/pong) ──TSTORE──▶ gmA (token) ──TSTORE──▶ gmPerTokenScale (scale) ViaGM 路径modeviagm 远端 GM ──TGET──▶ 本地临时 GM ──TLOAD──▶ UB (ping/pong) ──TSTORE──▶ gmA ──TSTORE──▶ gmPerTokenScale WithSync 路径modesync Phase ATPE AllGatherTSTORE 远端写 TWAIT 本地 TPE ──TSTOREDataAsFlag──▶ 所有远端 rank ──TWAIT──▶ 接收所有远端 rank 的 TPE Phase B计算路由表设备侧 B.1 去除 DataAsFlag 偏移向量化 TLOAD/TADDS/TSTORE B.2 计算 cumsumMM 前缀和向量化 TLOAD/TADD/TSTORE B.3 计算 preSumBeforeRank标量累加 Phase C使用计算好的路由表调用 MoeDispatchDirect算法 Direct / ViaGM for 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 之间无气泡 WithSync 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 使用计算好的路由表调用 MoeDispatchDirect核心特性三路径设计Direct快速、ViaGM兼容、WithSync自包含集成 CrossRankSyncWithSync 路径在设备侧计算路由表无需 host 预计算向量化 cumsumMMTLOAD/TADD/TSTORE 前缀和事件流水32B 对齐 padding软件 SYNCALL基于 GM 轮询的跨核同步避免 FFTS 硬件依赖自适应 MOVE_NUM编译时DispatchTraitsTILE_COLS根据 UB 容量自动缩减每批处理行数事件驱动 ping-pongset_flag/wait_flag重叠 MTE2TLOAD和 MTE3TSTORE管道跨 rank 连续流水ping-pong 状态跨远端 rank 持续rank 边界无刷新多核并行每个 AIV 核处理一个或多个远端 rank跨步分配Token/Scale 分离远端行[int8×K][float scale 对齐到 32B]→ 紧凑 token scale 输出A5 特殊说明HCCL 窗口偏移勘误A5 MTE2 DMA 引擎读取 HCCL 窗口基地址 [16..31] 字节时返回全零。 Host 驱动设置winOffset256跳过缺陷区域。HCCL V2 TilingA5 使用 V2 tiling 初始化路径tests/npu/a5/下的common.hpp。编译器目标dav-c310-vecA5 向量核 ISA。规格说明项目值数据类型tokenint8_t数据类型scalefloat存储在 32B 对齐行中远端行格式hiddenSize字节 token UB_ALIGN32字节 paddingscale 在偏移 0 处输出 tokengmA[maxOutputSize, hiddenSize]— 紧凑无 padding输出 scalegmPerTokenScale[maxOutputSize]— 每行 32 字节float 在偏移 0 处默认 hiddenSize128执行模型仅 AIV向量核通过 mpirun 多 rank目录结构kernels/manual/a5/moe_dispatch/ ├── moe_dispatch_kernel.cpp # 设备 kernel三路径 dispatch ├── main.cpp # Host 驱动MPI 初始化、数据生成、launch、验证 ├── moe_dispatch_config.h # 形状常量、DispatchTraits、workspace 布局 ├── hccl_context.h # 设备侧 HCCL 上下文结构体 ├── CMakeLists.txt # 构建配置bisheng dav-c310-vec ├── run.sh # 构建和运行便捷脚本 ├── README.md # 英文说明 └── README_zh.md # 本文件构建与运行# 设置环境 source /mnt/data/ntlab/liulei/set_env_new.sh export HCCL_WHITELIST_DISABLE1 # 构建并运行 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—启用调试模式与 MegaMoE 的关系本算子验证了 MegaMoE 的 Dispatch 阶段可直接作为完整 MegaMoE 融合算子的构建模块MegaMoE 完整流水线 InitRouting → [Dispatch] → GEMM (FFN) → Combine ^^^^^^^^ 本算子WithSync 覆盖 InitRouting Dispatch接口兼容参数cumsumMM、tokenPerExpert、preSumBeforeRank、shmemBase与 MegaMoE 完全对齐WithSync 路径等价于 MegaMoE 的CrossRankSyncAndlocalTokenPerExpertAllGatherAndGetSumPreRankV2DispatchAndCombinedispatch 部分ViaGM 路径功能上等价于 MegaMoE 的DispatchCopyPerTokenDirect 路径PTO-ISA 独有优化绕过中间 GM buffer参考资料MegaMoE 源码vllm-ascend/csrc/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_kernel.hpp设计文档/mnt/data/ntlab/liulei/docs/megamoe/dispatch_pto_isa_design.mdPTO-ISA TGET APIinclude/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），仅供参考