拓冰网站定制 
CANN pto-isa 实战基于 A5 PTO 的 MoE Combine Kernel 实现与调优指南【免费下载链接】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_combine 示例如何在 Ascend950 A5 上使用 PTO 虚拟指令集实现 MoEMixture of Expertsdispatch-compute-combine 流水中的 combine 半段。通过阅读本文你将掌握 combine kernel 的完整数据流变长 return → 跨 rank 同步 → 加权还原、显式低层路由账本routeMeta的布局契约、基于 HCCL RDMA peer window 的TPUT/TNOTIFY/TWAIT通信模式以及从 CMake 构建到 CPU golden 验证的完整工程化实践。背景与定位combine 在 MoE 流水中的角色在 MoE 推理/训练中token 经过 gating 网络后按topK路由到不同的 expert形成 dispatch-compute-combine 三段流水dispatch 阶段把 token 按路由账本分发给 expert 所在的 rank本示例不包含该 kernelcompute 阶段各 rank 上的 expertFFN/GMM对到达本地的 token 行做计算本示例不包含combine 阶段本示例主体本地 expert 完成计算后把 expert 输出行按路由账本返还给原 token 所在 rank并使用 gate 权重probs还原每个 token 的最终输出。本 kernel 是独立 combine kernel直接消费显式低层路由账本routeMeta。expert_ids、assist_info_for_combine、ep_send_counts等路由信息由上游或 host 侧按本算子的账本布局整理后显式写入routeMeta传入整体数据流如下expertOutput[local expert rows, K] - 通过 HCCL peerWindow.ptrD 做变长 return - 通过 TNOTIFY/TWAIT 做跨 rank 完成同步 - 加权还原: outputC[token, :] sum(topK probs * returned rows)支持的 AI 处理器Ascend950A5 架构目录结构一个完整的算子工程kernels/manual/a5/moe_combine 是一个麻雀虽小、五脏俱全的 A5 算子工程源码文件职责如下kernels/manual/a5/moe_combine/ ├── CMakeLists.txt # Bisheng CCE host 构建配置 ├── run.sh # 一键构建和运行脚本 ├── common.h # 共享 ABI: shape, routeMeta layout, peerWindow layout, HCCL context ├── layout.h # Host 侧 layout 计算和 HCCL_BUFFSIZE 估算 ├── kernel_launchers.h # Host 侧 kernel launcher 声明 ├── moe_combine_kernel.cpp # PTO AIV kernel: return wait weighted restore ├── main.cpp # Host 编排: MPI, ACL, HCCL window, fixture, verify, profile ├── golden.h # CPU golden 数据结构和公开接口声明 ├── golden.cpp # CPU golden 路由构造和输出校验实现 ├── hccl_context.h # A5 HCCL direct context 与 peer window 初始化 ├── comm_mpi.h # MPI 动态加载封装 ├── DESIGN.md # A5 翻译和 HCCL window padding 设计 ├── README.md # 英文 README └── README_zh.md # 中文 README工程按 kernel设备侧 host编排/校验 双层组织设备侧只有 moe_combine_kernel.cpp 一个文件全部 PTO 逻辑集中于此host 侧负责参数解析、ACL/HCCL 初始化、fixture 构造、性能统计与 golden 校验。算子说明计算功能对每个 rank本算子消费已经按本地 expert 和来源 rank 排布好的 expert 输出。kernel 内部流程是读取routeMeta得到每个 source rank 给每个 expert 的行数以及这些行在expertOutput中的位置使用 PTOTPUT将每行 expert 输出返还到 token owner rank 的 HCCL peer window使用TNOTIFY/TWAIT等待所有 peer 完成 return 写入读取routeMeta.expandedRowIdx和probs还原outputC[M, K]。对本 rank 的第t个 tokendispatch 阶段会产生topK条 expert route。combine return 完成后这些 route 对应的 expert 输出行已经写回本 rank 的peerWindow.ptrD。expandedRowIdx[t * topK slot]记录第slot条 route 在ptrD中的行号probs[t * topK slot]是这条 route 的 gate 权重。因此对输出的每一列c还原逻辑是outputC[t, c] 0 for slot in 0..topK-1: row expandedRowIdx[t * topK slot] if row 0: outputC[t, c] probs[t * topK slot] * peerWindow.ptrD[row, c]也就是把同一个 token 的topK路 expert 输出按 gate 权重加权求和得到最终的outputC[t, :]。该还原逻辑在 golden.h 的RestoreOutputC函数对应 CPU golden 侧和 moe_combine_kernel.cpp 的RestoreOutputRows/AccumulateRestoreTile中一一对应。覆盖范围包含不包含EP 域内基于 HCCL window 的 combine returnDispatch pack/gather kernel使用TPUT实现变长 all-to-all-like returnHCCL collectiveAllToAllVAPI使用probs做加权还原Expert FFN/GMM 计算显式低层routeMeta契约量化、TP ReduceScatterV、shared/copy/const expertA5 HCCL window 头部 guard 处理上层公共 ABI 适配层从覆盖范围可以看出本示例刻意保持独立 combine kernel的纯粹性它不依赖 HCCL 的AllToAllVcollective API而是直接操作底层 RDMA window 实现变长 return为上层在 EP 域内的定制化通信留出空间。入口契约Kernel Launcher ABIkernel_launchers.h 声明的 host 侧 launcher 签名如下moe_combine_kernel.cpp 末尾给出了其实现内部以MoeCombineKernellaunchBlockCount, nullptr, stream的形式把参数直通传给__global__kernelvoid LaunchMoeCombineKernel(MoeCombineShape shape, uint32_t myRank, uint8_t *expertOutput, uint8_t *probs, uint8_t *outputC, uint8_t *routeMeta, uint8_t *peerWindow, uint8_t *hcclCtx, uint8_t *workspace, void *stream, uint32_t launchBlockCount);运行时输入参数方向存储含义shape输入值传递静态 shape 和 AIV block 数如ep,m,k,topK,expertPerRank,aivBlocksmyRank输入值传递EP 域内 rank idexpertOutput输入aclrtMallocGM本地 expert 输出行形状[maxOutputSize, K]fp16probs输入aclrtMallocGMgate 权重形状[M, topK]fp32outputC输出aclrtMallocGM还原后的 token 输出形状[M, K]fp16routeMeta输入aclrtMallocGM显式 combine 路由账本peerWindow输入/输出HCCL RDMA window远端可见的ptrDreturn buffer 和 signalhcclCtx输入aclrtMallocGM设备侧 HCCL window 地址和 A5 direct-context 字段workspace临时aclrtMallocGM本地 AIV soft sync 区stream输入ACL streamkernel launch streamlaunchBlockCount输入值传递kernel 使用的 AIV block 数这些 buffer 在 main.cpp 的AllocateLocalBuffers中按layout.h计算的字节数通过aclrtMalloc(..., ACL_MEM_MALLOC_HUGE_FIRST)分配。peerWindow内容localWindowBase是 HCCL window 的原始起始地址。A5 在 window 头部保留4096B head guard传给 kernel 的peerWindow指向 live payload即localWindowBase 4096该常量定义于 layout.h 的kMoeCombineWindowHeadGuardByteshost 清零时也会覆盖 head guard live payload 的完整区间见 main.cpp 的ClearDeviceState。A5 localWindowBase [4096B head guard] peerWindow live payload: ptrD countReadySignal[ep] combineDoneSignal[ep]字段位置内容ptrDHCCL window live payloadreturn 目标行被远端TPUT写入countReadySignal[ep]HCCL window live payloadper-rank ready 计数区combineDoneSignal[ep]HCCL window live payloadper-rank 完成计数器远端 rank 完成写入本 rankptrD后TNOTIFY对应槽位peerWindow的 live payload 布局由 layout.h 的ComputePeerWindowLayout计算先放ptrDexpandedRows * K * 2字节fp16再放countReadySignal[ep]与combineDoneSignal[ep]各ep * 4字节int32每个字段按 64 字节对齐。MoeCombineShape该结构体定义于 common.h是 host/device 两侧共享的 ABI字段含义epEP rank 数m每 rank token 数khidden sizetopK每 token 的 expert 路由数expertPerRank每 rank 本地 expert 数expertNum全局 expert 数通常为ep * expertPerRankmaxOutputSize每 rank expert 输出最大行容量aivBlocks逻辑 AIV block 数A5 默认40可传参覆盖routeMeta布局routeMeta是显式低层 combine 路由账本。它是本地 GM不属于 HCCL window。其字节布局由ComputeCombineRouteMetaLayout计算每个字段 64 字节对齐字段含义如下字段形状含义peerTokenPerExpert[ep, expertNumPadded]int32每个 source rank 到每个 global expert 的行数expandedRowIdx[M * topK]int32token route 到peerWindow.ptrD的行映射-1表示无效 routecumsumPerExpert[ep, expertNumPadded]int32每个 source rank 内按 global expert 的 inclusive prefixcumsum[src,e] sum(peerTokenPerExpert[src,0..e])dispatchOffset[expertPerRank]int32每个本地 expert 在expertOutput中的基地址行prevSumBeforeRank[ep, expertPerRank]int32某 source rank 在本地 expert 行段中的前缀偏移其中expertNumPadded是expertNum按16kMoeCombineMetadataPad向上对齐的结果。dispatchOffset与prevSumBeforeRank的组合语义是srcStart dispatchOffset[localExpert] prevSumBeforeRank[src, localExpert]即每个 source rank 的行段在expertOutput中的起始行号这正是 return 阶段定位源数据的关键索引。Kernel 三阶段实现剖析MoeCombineKernel的入口见 moe_combine_kernel.cpp在完成 shape 校验后依次执行三个设备侧阶段中间以SoftSyncAiv软件同步分隔ReturnExpertRowsToOwners - WaitCombinePhase - RestoreOutputRows阶段 1: ReturnExpertRowsToOwners变长 returnkernel 遍历所有本地 expert segmentsegment src_rank * expertPerRank localExpert globalExpert myRank * expertPerRank localExpert rows routeMeta.peerTokenPerExpert[src_rank, globalExpert]对每个非空 segmentsrcStart由dispatchOffset[localExpert] prevSumBeforeRank[src_rank, localExpert]计算dstStart由cumsumPerExpert[src_rank, globalExpert - 1]计算globalExpert 0时为0如果src_rank myRank行被本地复制到本 rank 的peerWindow.ptrDCopyLocalRowsToPeerWindow逐行逐 tile 用TLOAD/TSTORE双缓冲搬运否则PTOTPUT把行 chunk 写入 source rank 的远端 peer windowPutRemoteRowsToOwner。TPUT是 PTO 的通信指令这里体现为pto::comm::TPUT(remoteDst, localSrc, ping, pong)remoteDst通过RemotePtr把本地地址换算为相对windowsIn[myRank]的偏移、再叠加到windowsIn[peerRank]得到远端地址ping/pong两块 UB tile 构成双缓冲让 MTE2 load 与 MTE3 store 流水化。return 阶段还采用了chunk 化分片每个 segment 按固定kMoeCombineRowChunk 8行切成 chunk并以(chunkBase chunk) % blockNum blockId的规则把 chunk 轮流分配给各 AIV block实现负载均衡。阶段 2: WaitCombinePhase跨 rank 完成同步return 写完后每个 rank 通知所有 token-owner rankTNOTIFY(remotePeer.combineDoneSignal[myRank], AtomicAdd) TWAIT(localPeer.combineDoneSignal[peer] 1)对应源码是NotifyCombineOwners对所有src槽位做TNOTIFY(sig, 1, NotifyOp::AtomicAdd)前置pipe_barrier(PIPE_ALL)保证 TPUT 全部落窗与WaitCombinePhase以blockId为起点按blockNum步进遍历所有 peerTWAIT(sig, value, WaitCmp::GE)。Host 会在每轮迭代前清零包含 head guard 和 live payload 的 HCCL window 区域因此 kernel 固定等待每个 peer 的一次 notify前置MPIBarrier只用于保证所有 rank 完成清理后再 launch kernel不计入combine_e2e。阶段 3: RestoreOutputRows加权还原每个 AIV block 负责一段连续 tokenTokenShardBegin/TokenShardEnd均匀切分。对每个 token 和每个列 tile使用TEXPANDS把输出 tile 清零对每个有效 route加载ptrD[expandedRowIdx]使用TAXPY(outTile, ptrTile, prob)累加将 fp16 tile 写回outputC。还原阶段的内部循环在AccumulateRestoreTile中实现使用pto::Eventpto::Op::TAXPY, pto::Op::TLOAD把上一次TAXPY与下一次TLOAD串联成 event chain保证 topK loop 顺序的同时避免每个 route 后都插入完整 pipe barrier。优化说明从源码看设计取舍该 kernel 是 AIV-only combine kernel。对于K7168这类 hidden size一行 fp16 数据是 14 KiB整体主要受 GM/HCCL window 搬运带宽影响。优化目标是让数据搬运尽量流式化同时降低控制面元数据开销。以下优化点均有对应源码佐证显式 routeMeta路由元数据作为独立 GM buffer 传入。peerWindow只保留远端可见 return 数据和信号workspace只保留本地 AIV soft sync 区——职责分离避免 RDMA window 被元数据污染。chunk 化 return 分片return 阶段遍历src_rank x local_expertsegment并按chunkBase % blockNum把行 chunk 分给 AIV blockReturnExpertRowsToOwners中的分配循环。PTOTPUTping/pong 路径远端 return 使用TPUT(remoteDst, localSrc, ping, pong)ping/pong 两块 UB tile 地址分别为0x0与0x1000通过 UB 双缓冲让 MTE2 load 和 MTE3 store 形成流水。Restore route cache当topK 16kRouteCacheMax时每个 token 的 route row 和 prob 会缓存到标量数组RestoreRouteCache{rows[16], probs[16], count}减少内层 restore loop 对 route metadata 的重复读取。DCCI 批量 acquire每个 token 在消费返回的ptrD行前先刷新对应 GM rangeDcciGmRangeNoFence按 64B cacheline 逐个dcci然后对本轮 cached routes 做一次dsb(DSB_DDR)避免逐行刷新的同步开销。A5 restore event chainA5 restore loop 使用 PTO event 串联TAXPY - next TLOAD减少 topK loop 内不必要的PIPE_ALLbarrier。Soft AIV sync同一个 kernel 内用pto::SYNCALLSyncAllMode::Soft分隔 return、wait、restore 阶段SoftSyncAivworkspace 中的localSync区大小按aivBlocks * (8 expertNumPadded)且不小于 64 个 int32 计算。Tiling 与默认参数参数默认值说明PES/ep2EP rank 数M64每 rank token 数K7168hidden sizetopK8每 token expert 路由数expertPerPe2每 rank 本地 expert 数expertNum4PES * expertPerPemaxOutputSizePES * M * topK默认容量默认 shape 下为1024aivBlocks40A5 resource model 默认值可通过--aiv-blocks N覆盖内部 Vector tile 列宽1024示例实现固定值kMoeCombineTileCols内部 return chunk8 rows固定的 return 阶段行 chunkkMoeCombineRowChunk内部 metadata pad16expert metadata 对齐粒度kMoeCombineMetadataPad常用调试命令--aiv-blocks 24并使用PES2, M64, K7168, topK8, expertPerPe2时各布局大小为Layout字节数workspace2304routeMeta2432peerWindowlive payload7340160A5 HCCL head guard4096这些数字与 run.sh 中 host 脚本侧的算术SYNC_SLOTS aivBlocks * (8 expertNumPadded)、EXPANDED_ROWS M * topK、PEER_WINDOW_LIVE_BYTES align(EXPANDED_ROWS * K * 2 PES*4 PES*4)等以及 layout.h 的设备侧/主机侧布局计算完全一致且 kernel 内部也会用MakeWorkspaceLayout/MakeCombineRouteMetaLayout/MakePeerWindowLayout在设备侧重新计算同一布局保证 host/device 两侧 ABI 对齐common.h注释明确要求两测字段名与结构保持一致。另外 layout.h 还提供EstimateHcclBuffSizeMb在head guard peerWindow totalBytes基础上加 64 MiB 安全余量并向上对齐到 MiB估算HCCL_BUFFSIZE。run.sh 中同样实现了这一估算逻辑AUTO_HCCL_BUFFSIZE并在未显式指定--hccl-buffsize-mb时以环境变量HCCL_BUFFSIZE导出。整体架构Host: ParseArgs - ComputeWorkspaceLayout / ComputeCombineRouteMetaLayout / ComputePeerWindowLayout - PrepareHostData and CPU golden - Init HCCL A5 direct context - AllocateLocalBuffers(routeMeta/workspace/expertOutput/probs/outputC) - loop(warmup measured): ClearDeviceState PrepareCombineFixture - 写入 routeMeta expertOutput LaunchMoeCombineKernel Verify outputC Device: ReturnExpertRowsToOwners - WaitCombinePhase - RestoreOutputRowsReturn phase: routeMeta(peerToken/cumsum/offset) expertOutput - local or remote peerWindow.ptrD - TNOTIFY peer combineDoneSignal[myRank] Restore phase: routeMeta.expandedRowIdx probs peerWindow.ptrD - outputCHost 侧的完整编排在 main.cpp 中体现InitRankInfoMPI rank/device 绑定→PrepareHostData确定性数据生成 CPU golden→InitHcclroot info 广播 InitHcclWindowContext→AllocateLocalBuffers→ 迭代循环ClearDeviceState/ fixture / launch / verify。构建与运行环境准备source /usr/local/Ascend/cann/set_env.sh执行run.sh前需要先在 shell 中加载 CANN 环境CMake 要求环境变量ASCEND_HOME_PATH已设置否则报Cannot find ASCEND_HOME_PATH。如果 shell 中没有mpirun请先配置 MPI 环境run.sh 启动时会检查mpirun是否可用。run.sh 还会在每次运行前清理/dev/shm/sem.hccl*与 IPC 残留可用--keep-hccl-shm 1关闭。仅编译手动 CMakecmake -S kernels/manual/a5/moe_combine \ -B /tmp/moe_combine_a5_build \ -DRUN_MODEnpu \ -DSOC_VERSIONAscend950PR_958b cmake --build /tmp/moe_combine_a5_build --target moe_combine -j8构建细节可参考 CMakeLists.txtkernel 以 SHARED 库moe_combine_kernel形式用 Bisheng CCE 编译--cce-aicore-archdav-c310-vec、--cce-fatobj-link、定义PTO_NPU_ARCH_A5host 可执行文件链接runtime ascendcl hcomm等库。A5 快速验证小 shapecd kernels/manual/a5/moe_combine bash run.sh -pes 2 -M 8 -K 64 -topK 2 -expertPerPe 1 --aiv-blocks 24默认 shapecd kernels/manual/a5/moe_combine bash run.sh -pes 2 -M 64 -K 7168 -topK 8 -expertPerPe 2 --aiv-blocks 24主要命令行参数参数默认值含义-pes2rank 数-M64每 rank token 数-K7168hidden size-topK8每 token route 数-expertPerPe2每 rank expert 数--max-output-sizePES * M * topKexpert output 行容量--aiv-blocks0 - 40逻辑 AIV block 数用于匹配不同硬件资源规划--device-base0rank 到 device 映射使用的起始 device id--ndevicesPES示例 launcher 使用的可见 device 数此外 run.sh 还支持数据与调试类参数-iters默认 5、-warmup默认 3、--seed默认 1234、--gen-data 0|1、--verify 0|1、--rtol/--atol默认 1e-2、-debug 0|1|2、--skip-build/--clean-build、--hccl-buffsize-mb、--rank-from-mpi、--rank等可通过bash run.sh --help查看完整说明。注意本项目不支持--case预设必须显式传 shape 参数脚本会校验 shape 字段非零、deviceBase pes ndevices、maxOutputSize EP * M * topK不支持容量不足时的 drop 语义。验证与性能Host 会构造确定性的 CPU golden 路由账本GenerateDeterministicInputstoken 输入按seed rank*131 token*17 col*3生成expert 路由按(flat rank) % expertNum分配prob 归一化将其写入routeMeta拷贝expertOutput启动 kernel并将outputC与 CPU golden 输出对比fp16 往返 rtol/atol容差比较见 golden.h 的CompareOutputs。默认开启验证。CPU golden 会完整模拟 dispatch 侧打包packedA、本地 dispatch 布局dispatchedA/expertOutput、远端ptrD内容BuildPtrD与最终加权还原RestoreOutputC是理解 routeMeta 各字段语义的最佳参考。开启-debug后还会把localTokenPerExpert、peerTokenPerExpert、cumsumPerExpert、expandedRowIdx、packedA_head、ptrD_head等中间张量以二进制文件写到--data-dir默认out/便于逐字段排查。预期成功输出verifyPASS脚本输出的 profile 形式如下A5 runtime 数字需要在 A5 机器上采集当前开发环境可编译 A5 工程[PROFILE] CombineTile M64 K7168 ranks2 topK8 expertPerPe2 warmup3 measured5 samples5 prepare_fixture: avg... us max... us combine_e2e: avg... us max... us verifyPASS关键指标含义指标含义combine_e2ecombine kernel launch 到 stream sync不包含 clear、fixture、verify也不包含 kernel launch 窗口之外的 MPI barrierverifyPASSdeviceoutputC与 CPU golden 一致profile 的统计逻辑在 main.cpp 的PrintProfileSummary中各 rank 的每次迭代计时通过 MPI gather 汇总到 rank 0每个 measured sample 取所有 rank 的最大值再输出 avg/max并附带logical work所有 rank 的 input tokens 与 routed tokens 数便于核算带宽与路由膨胀比。注意事项与限制本示例是独立 combine kernel参考实现不包含 dispatch、expert FFN/GMM、量化、TP ReduceScatterV 以及 shared/copy/const expert 等能力接入上层需自行适配公共 ABI不依赖 HCCLAllToAllVcollective API直接使用底层 window 原语TPUT/TNOTIFY/TWAIT对 EP 域内通信语义有完全控制但也意味着需要自行保证信号清零、barrier 时序与 buffer 容量maxOutputSize小于EP * M * topK时脚本直接报错本实现不支持行丢弃drop语义数值校验采用 fp16 输出 rtol/atol 1e-2容差适合验证路由与加权逻辑正确性不代表 fp16 累加精度上限。【免费下载链接】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），仅供参考 CANN pto-isa 实战：基于 A5 PTO 的 MoE Combine Kernel 实现与调优指南
CANN pto-isa 实战基于 A5 PTO 的 MoE Combine Kernel 实现与调优指南【免费下载链接】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_combine 示例如何在 Ascend950 A5 上使用 PTO 虚拟指令集实现 MoEMixture of Expertsdispatch-compute-combine 流水中的 combine 半段。通过阅读本文你将掌握 combine kernel 的完整数据流变长 return → 跨 rank 同步 → 加权还原、显式低层路由账本routeMeta的布局契约、基于 HCCL RDMA peer window 的TPUT/TNOTIFY/TWAIT通信模式以及从 CMake 构建到 CPU golden 验证的完整工程化实践。背景与定位combine 在 MoE 流水中的角色在 MoE 推理/训练中token 经过 gating 网络后按topK路由到不同的 expert形成 dispatch-compute-combine 三段流水dispatch 阶段把 token 按路由账本分发给 expert 所在的 rank本示例不包含该 kernelcompute 阶段各 rank 上的 expertFFN/GMM对到达本地的 token 行做计算本示例不包含combine 阶段本示例主体本地 expert 完成计算后把 expert 输出行按路由账本返还给原 token 所在 rank并使用 gate 权重probs还原每个 token 的最终输出。本 kernel 是独立 combine kernel直接消费显式低层路由账本routeMeta。expert_ids、assist_info_for_combine、ep_send_counts等路由信息由上游或 host 侧按本算子的账本布局整理后显式写入routeMeta传入整体数据流如下expertOutput[local expert rows, K] - 通过 HCCL peerWindow.ptrD 做变长 return - 通过 TNOTIFY/TWAIT 做跨 rank 完成同步 - 加权还原: outputC[token, :] sum(topK probs * returned rows)支持的 AI 处理器Ascend950A5 架构目录结构一个完整的算子工程kernels/manual/a5/moe_combine 是一个麻雀虽小、五脏俱全的 A5 算子工程源码文件职责如下kernels/manual/a5/moe_combine/ ├── CMakeLists.txt # Bisheng CCE host 构建配置 ├── run.sh # 一键构建和运行脚本 ├── common.h # 共享 ABI: shape, routeMeta layout, peerWindow layout, HCCL context ├── layout.h # Host 侧 layout 计算和 HCCL_BUFFSIZE 估算 ├── kernel_launchers.h # Host 侧 kernel launcher 声明 ├── moe_combine_kernel.cpp # PTO AIV kernel: return wait weighted restore ├── main.cpp # Host 编排: MPI, ACL, HCCL window, fixture, verify, profile ├── golden.h # CPU golden 数据结构和公开接口声明 ├── golden.cpp # CPU golden 路由构造和输出校验实现 ├── hccl_context.h # A5 HCCL direct context 与 peer window 初始化 ├── comm_mpi.h # MPI 动态加载封装 ├── DESIGN.md # A5 翻译和 HCCL window padding 设计 ├── README.md # 英文 README └── README_zh.md # 中文 README工程按 kernel设备侧 host编排/校验 双层组织设备侧只有 moe_combine_kernel.cpp 一个文件全部 PTO 逻辑集中于此host 侧负责参数解析、ACL/HCCL 初始化、fixture 构造、性能统计与 golden 校验。算子说明计算功能对每个 rank本算子消费已经按本地 expert 和来源 rank 排布好的 expert 输出。kernel 内部流程是读取routeMeta得到每个 source rank 给每个 expert 的行数以及这些行在expertOutput中的位置使用 PTOTPUT将每行 expert 输出返还到 token owner rank 的 HCCL peer window使用TNOTIFY/TWAIT等待所有 peer 完成 return 写入读取routeMeta.expandedRowIdx和probs还原outputC[M, K]。对本 rank 的第t个 tokendispatch 阶段会产生topK条 expert route。combine return 完成后这些 route 对应的 expert 输出行已经写回本 rank 的peerWindow.ptrD。expandedRowIdx[t * topK slot]记录第slot条 route 在ptrD中的行号probs[t * topK slot]是这条 route 的 gate 权重。因此对输出的每一列c还原逻辑是outputC[t, c] 0 for slot in 0..topK-1: row expandedRowIdx[t * topK slot] if row 0: outputC[t, c] probs[t * topK slot] * peerWindow.ptrD[row, c]也就是把同一个 token 的topK路 expert 输出按 gate 权重加权求和得到最终的outputC[t, :]。该还原逻辑在 golden.h 的RestoreOutputC函数对应 CPU golden 侧和 moe_combine_kernel.cpp 的RestoreOutputRows/AccumulateRestoreTile中一一对应。覆盖范围包含不包含EP 域内基于 HCCL window 的 combine returnDispatch pack/gather kernel使用TPUT实现变长 all-to-all-like returnHCCL collectiveAllToAllVAPI使用probs做加权还原Expert FFN/GMM 计算显式低层routeMeta契约量化、TP ReduceScatterV、shared/copy/const expertA5 HCCL window 头部 guard 处理上层公共 ABI 适配层从覆盖范围可以看出本示例刻意保持独立 combine kernel的纯粹性它不依赖 HCCL 的AllToAllVcollective API而是直接操作底层 RDMA window 实现变长 return为上层在 EP 域内的定制化通信留出空间。入口契约Kernel Launcher ABIkernel_launchers.h 声明的 host 侧 launcher 签名如下moe_combine_kernel.cpp 末尾给出了其实现内部以MoeCombineKernellaunchBlockCount, nullptr, stream的形式把参数直通传给__global__kernelvoid LaunchMoeCombineKernel(MoeCombineShape shape, uint32_t myRank, uint8_t *expertOutput, uint8_t *probs, uint8_t *outputC, uint8_t *routeMeta, uint8_t *peerWindow, uint8_t *hcclCtx, uint8_t *workspace, void *stream, uint32_t launchBlockCount);运行时输入参数方向存储含义shape输入值传递静态 shape 和 AIV block 数如ep,m,k,topK,expertPerRank,aivBlocksmyRank输入值传递EP 域内 rank idexpertOutput输入aclrtMallocGM本地 expert 输出行形状[maxOutputSize, K]fp16probs输入aclrtMallocGMgate 权重形状[M, topK]fp32outputC输出aclrtMallocGM还原后的 token 输出形状[M, K]fp16routeMeta输入aclrtMallocGM显式 combine 路由账本peerWindow输入/输出HCCL RDMA window远端可见的ptrDreturn buffer 和 signalhcclCtx输入aclrtMallocGM设备侧 HCCL window 地址和 A5 direct-context 字段workspace临时aclrtMallocGM本地 AIV soft sync 区stream输入ACL streamkernel launch streamlaunchBlockCount输入值传递kernel 使用的 AIV block 数这些 buffer 在 main.cpp 的AllocateLocalBuffers中按layout.h计算的字节数通过aclrtMalloc(..., ACL_MEM_MALLOC_HUGE_FIRST)分配。peerWindow内容localWindowBase是 HCCL window 的原始起始地址。A5 在 window 头部保留4096B head guard传给 kernel 的peerWindow指向 live payload即localWindowBase 4096该常量定义于 layout.h 的kMoeCombineWindowHeadGuardByteshost 清零时也会覆盖 head guard live payload 的完整区间见 main.cpp 的ClearDeviceState。A5 localWindowBase [4096B head guard] peerWindow live payload: ptrD countReadySignal[ep] combineDoneSignal[ep]字段位置内容ptrDHCCL window live payloadreturn 目标行被远端TPUT写入countReadySignal[ep]HCCL window live payloadper-rank ready 计数区combineDoneSignal[ep]HCCL window live payloadper-rank 完成计数器远端 rank 完成写入本 rankptrD后TNOTIFY对应槽位peerWindow的 live payload 布局由 layout.h 的ComputePeerWindowLayout计算先放ptrDexpandedRows * K * 2字节fp16再放countReadySignal[ep]与combineDoneSignal[ep]各ep * 4字节int32每个字段按 64 字节对齐。MoeCombineShape该结构体定义于 common.h是 host/device 两侧共享的 ABI字段含义epEP rank 数m每 rank token 数khidden sizetopK每 token 的 expert 路由数expertPerRank每 rank 本地 expert 数expertNum全局 expert 数通常为ep * expertPerRankmaxOutputSize每 rank expert 输出最大行容量aivBlocks逻辑 AIV block 数A5 默认40可传参覆盖routeMeta布局routeMeta是显式低层 combine 路由账本。它是本地 GM不属于 HCCL window。其字节布局由ComputeCombineRouteMetaLayout计算每个字段 64 字节对齐字段含义如下字段形状含义peerTokenPerExpert[ep, expertNumPadded]int32每个 source rank 到每个 global expert 的行数expandedRowIdx[M * topK]int32token route 到peerWindow.ptrD的行映射-1表示无效 routecumsumPerExpert[ep, expertNumPadded]int32每个 source rank 内按 global expert 的 inclusive prefixcumsum[src,e] sum(peerTokenPerExpert[src,0..e])dispatchOffset[expertPerRank]int32每个本地 expert 在expertOutput中的基地址行prevSumBeforeRank[ep, expertPerRank]int32某 source rank 在本地 expert 行段中的前缀偏移其中expertNumPadded是expertNum按16kMoeCombineMetadataPad向上对齐的结果。dispatchOffset与prevSumBeforeRank的组合语义是srcStart dispatchOffset[localExpert] prevSumBeforeRank[src, localExpert]即每个 source rank 的行段在expertOutput中的起始行号这正是 return 阶段定位源数据的关键索引。Kernel 三阶段实现剖析MoeCombineKernel的入口见 moe_combine_kernel.cpp在完成 shape 校验后依次执行三个设备侧阶段中间以SoftSyncAiv软件同步分隔ReturnExpertRowsToOwners - WaitCombinePhase - RestoreOutputRows阶段 1: ReturnExpertRowsToOwners变长 returnkernel 遍历所有本地 expert segmentsegment src_rank * expertPerRank localExpert globalExpert myRank * expertPerRank localExpert rows routeMeta.peerTokenPerExpert[src_rank, globalExpert]对每个非空 segmentsrcStart由dispatchOffset[localExpert] prevSumBeforeRank[src_rank, localExpert]计算dstStart由cumsumPerExpert[src_rank, globalExpert - 1]计算globalExpert 0时为0如果src_rank myRank行被本地复制到本 rank 的peerWindow.ptrDCopyLocalRowsToPeerWindow逐行逐 tile 用TLOAD/TSTORE双缓冲搬运否则PTOTPUT把行 chunk 写入 source rank 的远端 peer windowPutRemoteRowsToOwner。TPUT是 PTO 的通信指令这里体现为pto::comm::TPUT(remoteDst, localSrc, ping, pong)remoteDst通过RemotePtr把本地地址换算为相对windowsIn[myRank]的偏移、再叠加到windowsIn[peerRank]得到远端地址ping/pong两块 UB tile 构成双缓冲让 MTE2 load 与 MTE3 store 流水化。return 阶段还采用了chunk 化分片每个 segment 按固定kMoeCombineRowChunk 8行切成 chunk并以(chunkBase chunk) % blockNum blockId的规则把 chunk 轮流分配给各 AIV block实现负载均衡。阶段 2: WaitCombinePhase跨 rank 完成同步return 写完后每个 rank 通知所有 token-owner rankTNOTIFY(remotePeer.combineDoneSignal[myRank], AtomicAdd) TWAIT(localPeer.combineDoneSignal[peer] 1)对应源码是NotifyCombineOwners对所有src槽位做TNOTIFY(sig, 1, NotifyOp::AtomicAdd)前置pipe_barrier(PIPE_ALL)保证 TPUT 全部落窗与WaitCombinePhase以blockId为起点按blockNum步进遍历所有 peerTWAIT(sig, value, WaitCmp::GE)。Host 会在每轮迭代前清零包含 head guard 和 live payload 的 HCCL window 区域因此 kernel 固定等待每个 peer 的一次 notify前置MPIBarrier只用于保证所有 rank 完成清理后再 launch kernel不计入combine_e2e。阶段 3: RestoreOutputRows加权还原每个 AIV block 负责一段连续 tokenTokenShardBegin/TokenShardEnd均匀切分。对每个 token 和每个列 tile使用TEXPANDS把输出 tile 清零对每个有效 route加载ptrD[expandedRowIdx]使用TAXPY(outTile, ptrTile, prob)累加将 fp16 tile 写回outputC。还原阶段的内部循环在AccumulateRestoreTile中实现使用pto::Eventpto::Op::TAXPY, pto::Op::TLOAD把上一次TAXPY与下一次TLOAD串联成 event chain保证 topK loop 顺序的同时避免每个 route 后都插入完整 pipe barrier。优化说明从源码看设计取舍该 kernel 是 AIV-only combine kernel。对于K7168这类 hidden size一行 fp16 数据是 14 KiB整体主要受 GM/HCCL window 搬运带宽影响。优化目标是让数据搬运尽量流式化同时降低控制面元数据开销。以下优化点均有对应源码佐证显式 routeMeta路由元数据作为独立 GM buffer 传入。peerWindow只保留远端可见 return 数据和信号workspace只保留本地 AIV soft sync 区——职责分离避免 RDMA window 被元数据污染。chunk 化 return 分片return 阶段遍历src_rank x local_expertsegment并按chunkBase % blockNum把行 chunk 分给 AIV blockReturnExpertRowsToOwners中的分配循环。PTOTPUTping/pong 路径远端 return 使用TPUT(remoteDst, localSrc, ping, pong)ping/pong 两块 UB tile 地址分别为0x0与0x1000通过 UB 双缓冲让 MTE2 load 和 MTE3 store 形成流水。Restore route cache当topK 16kRouteCacheMax时每个 token 的 route row 和 prob 会缓存到标量数组RestoreRouteCache{rows[16], probs[16], count}减少内层 restore loop 对 route metadata 的重复读取。DCCI 批量 acquire每个 token 在消费返回的ptrD行前先刷新对应 GM rangeDcciGmRangeNoFence按 64B cacheline 逐个dcci然后对本轮 cached routes 做一次dsb(DSB_DDR)避免逐行刷新的同步开销。A5 restore event chainA5 restore loop 使用 PTO event 串联TAXPY - next TLOAD减少 topK loop 内不必要的PIPE_ALLbarrier。Soft AIV sync同一个 kernel 内用pto::SYNCALLSyncAllMode::Soft分隔 return、wait、restore 阶段SoftSyncAivworkspace 中的localSync区大小按aivBlocks * (8 expertNumPadded)且不小于 64 个 int32 计算。Tiling 与默认参数参数默认值说明PES/ep2EP rank 数M64每 rank token 数K7168hidden sizetopK8每 token expert 路由数expertPerPe2每 rank 本地 expert 数expertNum4PES * expertPerPemaxOutputSizePES * M * topK默认容量默认 shape 下为1024aivBlocks40A5 resource model 默认值可通过--aiv-blocks N覆盖内部 Vector tile 列宽1024示例实现固定值kMoeCombineTileCols内部 return chunk8 rows固定的 return 阶段行 chunkkMoeCombineRowChunk内部 metadata pad16expert metadata 对齐粒度kMoeCombineMetadataPad常用调试命令--aiv-blocks 24并使用PES2, M64, K7168, topK8, expertPerPe2时各布局大小为Layout字节数workspace2304routeMeta2432peerWindowlive payload7340160A5 HCCL head guard4096这些数字与 run.sh 中 host 脚本侧的算术SYNC_SLOTS aivBlocks * (8 expertNumPadded)、EXPANDED_ROWS M * topK、PEER_WINDOW_LIVE_BYTES align(EXPANDED_ROWS * K * 2 PES*4 PES*4)等以及 layout.h 的设备侧/主机侧布局计算完全一致且 kernel 内部也会用MakeWorkspaceLayout/MakeCombineRouteMetaLayout/MakePeerWindowLayout在设备侧重新计算同一布局保证 host/device 两侧 ABI 对齐common.h注释明确要求两测字段名与结构保持一致。另外 layout.h 还提供EstimateHcclBuffSizeMb在head guard peerWindow totalBytes基础上加 64 MiB 安全余量并向上对齐到 MiB估算HCCL_BUFFSIZE。run.sh 中同样实现了这一估算逻辑AUTO_HCCL_BUFFSIZE并在未显式指定--hccl-buffsize-mb时以环境变量HCCL_BUFFSIZE导出。整体架构Host: ParseArgs - ComputeWorkspaceLayout / ComputeCombineRouteMetaLayout / ComputePeerWindowLayout - PrepareHostData and CPU golden - Init HCCL A5 direct context - AllocateLocalBuffers(routeMeta/workspace/expertOutput/probs/outputC) - loop(warmup measured): ClearDeviceState PrepareCombineFixture - 写入 routeMeta expertOutput LaunchMoeCombineKernel Verify outputC Device: ReturnExpertRowsToOwners - WaitCombinePhase - RestoreOutputRowsReturn phase: routeMeta(peerToken/cumsum/offset) expertOutput - local or remote peerWindow.ptrD - TNOTIFY peer combineDoneSignal[myRank] Restore phase: routeMeta.expandedRowIdx probs peerWindow.ptrD - outputCHost 侧的完整编排在 main.cpp 中体现InitRankInfoMPI rank/device 绑定→PrepareHostData确定性数据生成 CPU golden→InitHcclroot info 广播 InitHcclWindowContext→AllocateLocalBuffers→ 迭代循环ClearDeviceState/ fixture / launch / verify。构建与运行环境准备source /usr/local/Ascend/cann/set_env.sh执行run.sh前需要先在 shell 中加载 CANN 环境CMake 要求环境变量ASCEND_HOME_PATH已设置否则报Cannot find ASCEND_HOME_PATH。如果 shell 中没有mpirun请先配置 MPI 环境run.sh 启动时会检查mpirun是否可用。run.sh 还会在每次运行前清理/dev/shm/sem.hccl*与 IPC 残留可用--keep-hccl-shm 1关闭。仅编译手动 CMakecmake -S kernels/manual/a5/moe_combine \ -B /tmp/moe_combine_a5_build \ -DRUN_MODEnpu \ -DSOC_VERSIONAscend950PR_958b cmake --build /tmp/moe_combine_a5_build --target moe_combine -j8构建细节可参考 CMakeLists.txtkernel 以 SHARED 库moe_combine_kernel形式用 Bisheng CCE 编译--cce-aicore-archdav-c310-vec、--cce-fatobj-link、定义PTO_NPU_ARCH_A5host 可执行文件链接runtime ascendcl hcomm等库。A5 快速验证小 shapecd kernels/manual/a5/moe_combine bash run.sh -pes 2 -M 8 -K 64 -topK 2 -expertPerPe 1 --aiv-blocks 24默认 shapecd kernels/manual/a5/moe_combine bash run.sh -pes 2 -M 64 -K 7168 -topK 8 -expertPerPe 2 --aiv-blocks 24主要命令行参数参数默认值含义-pes2rank 数-M64每 rank token 数-K7168hidden size-topK8每 token route 数-expertPerPe2每 rank expert 数--max-output-sizePES * M * topKexpert output 行容量--aiv-blocks0 - 40逻辑 AIV block 数用于匹配不同硬件资源规划--device-base0rank 到 device 映射使用的起始 device id--ndevicesPES示例 launcher 使用的可见 device 数此外 run.sh 还支持数据与调试类参数-iters默认 5、-warmup默认 3、--seed默认 1234、--gen-data 0|1、--verify 0|1、--rtol/--atol默认 1e-2、-debug 0|1|2、--skip-build/--clean-build、--hccl-buffsize-mb、--rank-from-mpi、--rank等可通过bash run.sh --help查看完整说明。注意本项目不支持--case预设必须显式传 shape 参数脚本会校验 shape 字段非零、deviceBase pes ndevices、maxOutputSize EP * M * topK不支持容量不足时的 drop 语义。验证与性能Host 会构造确定性的 CPU golden 路由账本GenerateDeterministicInputstoken 输入按seed rank*131 token*17 col*3生成expert 路由按(flat rank) % expertNum分配prob 归一化将其写入routeMeta拷贝expertOutput启动 kernel并将outputC与 CPU golden 输出对比fp16 往返 rtol/atol容差比较见 golden.h 的CompareOutputs。默认开启验证。CPU golden 会完整模拟 dispatch 侧打包packedA、本地 dispatch 布局dispatchedA/expertOutput、远端ptrD内容BuildPtrD与最终加权还原RestoreOutputC是理解 routeMeta 各字段语义的最佳参考。开启-debug后还会把localTokenPerExpert、peerTokenPerExpert、cumsumPerExpert、expandedRowIdx、packedA_head、ptrD_head等中间张量以二进制文件写到--data-dir默认out/便于逐字段排查。预期成功输出verifyPASS脚本输出的 profile 形式如下A5 runtime 数字需要在 A5 机器上采集当前开发环境可编译 A5 工程[PROFILE] CombineTile M64 K7168 ranks2 topK8 expertPerPe2 warmup3 measured5 samples5 prepare_fixture: avg... us max... us combine_e2e: avg... us max... us verifyPASS关键指标含义指标含义combine_e2ecombine kernel launch 到 stream sync不包含 clear、fixture、verify也不包含 kernel launch 窗口之外的 MPI barrierverifyPASSdeviceoutputC与 CPU golden 一致profile 的统计逻辑在 main.cpp 的PrintProfileSummary中各 rank 的每次迭代计时通过 MPI gather 汇总到 rank 0每个 measured sample 取所有 rank 的最大值再输出 avg/max并附带logical work所有 rank 的 input tokens 与 routed tokens 数便于核算带宽与路由膨胀比。注意事项与限制本示例是独立 combine kernel参考实现不包含 dispatch、expert FFN/GMM、量化、TP ReduceScatterV 以及 shared/copy/const expert 等能力接入上层需自行适配公共 ABI不依赖 HCCLAllToAllVcollective API直接使用底层 window 原语TPUT/TNOTIFY/TWAIT对 EP 域内通信语义有完全控制但也意味着需要自行保证信号清零、barrier 时序与 buffer 容量maxOutputSize小于EP * M * topK时脚本直接报错本实现不支持行丢弃drop语义数值校验采用 fp16 输出 rtol/atol 1e-2容差适合验证路由与加权逻辑正确性不代表 fp16 累加精度上限。【免费下载链接】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），仅供参考