PTO MoE算子实战dispatch/combine如何与计算流水线重叠实现性能倍增【免费下载链接】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-isaMoEMixture of Experts推理的性能瓶颈往往不在算力而在dispatch分发与 combine归并两段跨卡通信token 要先按 expert 排序、AllToAll 发到目标卡计算完再发回来累加每一步都在算等通信、通信等算。PTO MoE 算子MegaMoE 融合算子的思路很直接——把 reorder、Dispatch、GMM1、SwiGLU、GMM2、Combine、Unpermute 七段全部融合进一个 kernel按 expert 粒度让 AICCube 计算和 AIV向量/通信互不等待让 dispatch/combine 被计算流水线吃掉。本文带你读懂这套流水重叠方案的设计逻辑以及它在实测中如何把 M2048 的 MoE FFN 耗时压低约 17%30%。MoE dispatch/combine 为什么慢串行衔接的代价先看传统 MoE FFN 的主链路 它包含 7 个串行阶段本地按 expert 重排 → AllToAll 发送 → 目标卡再次重排 → GMM1/SwiGLU/GMM2 计算 → 按卡重排 → AllToAllV 发回 → topK 加权累加。问题在于两次 permute、两次 AllToAll 都是独立的 kernel 或独立阶段阶段边界处存在全局同步 barrier。AIC 等 Dispatch 搬数据、AIV 等 GMM2 算完才能 Combine——每个 expert 的数据都排队过隧道流水线气泡bubble随 rank 数和 expert 数累积这正是 MoE 通信掩盖优化的核心目标。MegaMoE 融合算子七阶段如何串成一条流水线PTO 的 MegaMoE 方案把上述流程压缩成一个设备端流水线FrontReorder - Dispatch - GMM1 - SwiGLU - GMM2 - Combine - Unpermute核心设计是把通信与计算的粒度拆细到单个 expertAIV 核按Dispatch(i) → SwiGLU(i) → Combine(i)逐 expert 推进AIC 核同时轮转执行各 expert 的 GMM1/GMM2。第 i 个 expert 的 GEMM 可以与第 i-1 个 expert 的 AllToAllV 写回同时进行宏观上变成两条并行的泳道两个精妙的细节让重叠真正落地 ⚡动态 AIC 分组GMM1 启动时占满全部 24 个 AIC算完后释放 8 个GMM2 先由 8 个 AIC 起步GMM1 结束、AIC 释放后在 expert 边界无全局 barrier 地扩展到 24 个。两组 AIC 的交接完全由expert 级到达/就绪标志驱动不需要整卡同步。两段式 SwiGLUSwiGLU 切成 segmentGMM1 一算完某段SwiGLU 立刻消费并量化GMM2 马上跟进——激活函数不再独占一个串行阶段。完整的时间线实况M2048、16 rank如下可以看到 AIV 的 Dispatch/SwiGLU/Combine 与 AIC 的 GMM1/GMM2 全程交错关键优化点count-as-flag 与 PTO tile GMM流水线要跑起来还依赖两个去同步化技巧① count-as-flag消灭 AllToAll count 后的全局 barrier。前重排阶段各 rank 要把每个 expert 收到多少行的 count 发给对端。传统做法是 count 交换后再做一次全同步PTO 直接把 count 行加一个 marker 值写到对端窗口数据本身兼作到达标志对端轮询非零即可省掉一整轮 barrier。② PTO tile 编程模式重做 GMM。GMM1/GMM2 是耗时大头PTO 用输出 tile swizzleN 向 9 列蛇形调度提升 L1 命中、L1→L0 片上多级复用、L1/L0 双缓冲重叠预取与计算平均 tile 计算时间从基线 30.05us 降到 20.93us约1.44x这些细节的完整伪码与接口TSORT32/TMRGSORT排序、TLOAD/TSTORE远端拉取、TMATMUL_ACC累加、CrossCoreSetFlag跨核通知等在 overview.md 和 pseudocode.md 中有逐阶段拆解。实测性能M 越大重叠收益越明显8 卡、K7168/N4096/topK8、每卡 16 个本地 expert 的典型 case 下PTO MegaMoE 与 AscendC 实现对比如下 规律很清晰小 M16128通信占比高但数据量小两者基本持平M512 起PTO 的 GMM tile 优化 通算重叠开始拉开差距领先约 20%30%M20484425us vs 5353us领先约 17%此时 AIC 已接近满载进一步空间在降低 AIC/AIV 对 HBM 带宽的争用。A5 进阶CV 直传让 GMM 输出不落 HBM面向 A5Ascend950PR的 dispatch_mega_combine 版本在 expert 级重叠之上又进一步GMM1 通过 CV FIFO 把 tile 直接交给配对的 AIV 做 SwiGLUGMM2 结果也直传 Combine两层 GMM 输出都不落 GM/HBM 中间结果从根本上减少 AIC 与 AIV 争抢 HBM 带宽配合动态 GMM mailbox 调度和 token-ready Unpermute一个 token 的全部 topK 路结果就绪即可被消费无需等 Combine 整体完成A5 场景下的流水线更加饱满。快速上手运行 PTO MoE 融合算子示例仓库中可直接运行的入口8 卡 A2/A3 主示例kernels/manual/a2a3/dispatch_mega_combine/中文说明见 README_zh.md设备端流水线入口在 op_kernel/dispatch_mega_combine.hA5 版本kernels/manual/a5/dispatch_mega_combine/如果想先分别理解两段通信可以参考独立的 moe_dispatch 与 moe_combine 示例典型用法默认 M2048 casecd ${git_clone_path}/kernels/manual/a2a3/dispatch_mega_combine ./run.sh换其他 M 值直接改参数例如bash run.sh --world-size 8 --m 512 --k 7168 --n 4096 --topk 8 --experts 16 --max-output-size 81940 --reuse-data。小结PTO MoE 算子的性能倍增并非来自单点加速而是三件事的叠加七阶段融合消除阶段间 barrier、expert 粒度通算重叠让 dispatch/combine 藏进 GMM 时间线、PTO tile 模式让 GEMM 本身快 ~30%。对新手来说记住一条原则即可——MoE 算子优化的本质是把通信与计算的同步粒度从整批 token 细化到单个 expert让 AIC 和 AIV 永远都有活干 。【免费下载链接】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），仅供参考