企业官网建设流程全解析

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 虚拟指令集实现 MoE（Mixture of Experts）dispatch-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 三段流水：

1. dispatch 阶段：把 token 按路由账本分发给 expert 所在的 rank（本示例不包含该 kernel）；
2. compute 阶段：各 rank 上的 expert（FFN/GMM）对到达本地的 token 行做计算（本示例不包含）；
3. 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 处理器

• Ascend950（A5 架构）

目录结构：一个完整的算子工程

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 内部流程是：

1. 读取routeMeta，得到每个 source rank 给每个 expert 的行数，以及这些行在expertOutput中的位置；
2. 使用 PTOTPUT将每行 expert 输出返还到 token owner rank 的 HCCL peer window；
3. 使用TNOTIFY/TWAIT等待所有 peer 完成 return 写入；
4. 读取routeMeta.expandedRowIdx和probs，还原outputC[M, K]。

对本 rank 的第t个 token，dispatch 阶段会产生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中一一对应。

覆盖范围

从覆盖范围可以看出，本示例刻意保持"独立 combine kernel"的纯粹性：它不依赖 HCCL 的AllToAllVcollective API，而是直接操作底层 RDMA window 实现变长 return，为上层在 EP 域内的定制化通信留出空间。

入口契约

Kernel Launcher ABI

kernel_launchers.h 声明的 host 侧 launcher 签名如下，moe_combine_kernel.cpp 末尾给出了其实现（内部以MoeCombineKernel<<<launchBlockCount, nullptr, stream>>>的形式把参数直通传给__global__kernel）：

void 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);

运行时输入

这些 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 的kMoeCombineWindowHeadGuardBytes，host 清零时也会覆盖 head guard + live payload 的完整区间，见 main.cpp 的ClearDeviceState）。

A5 localWindowBase [4096B head guard] peerWindow live payload: ptrD countReadySignal[ep] combineDoneSignal[ep]

peerWindow的 live payload 布局由 layout.h 的ComputePeerWindowLayout计算：先放ptrD（expandedRows * K * 2字节，fp16），再放countReadySignal[ep]与combineDoneSignal[ep]（各ep * 4字节，int32），每个字段按 64 字节对齐。

MoeCombineShape

该结构体定义于 common.h，是 host/device 两侧共享的 ABI：

routeMeta布局

routeMeta是显式低层 combine 路由账本。它是本地 GM，不属于 HCCL window。其字节布局由ComputeCombineRouteMetaLayout计算（每个字段 64 字节对齐），字段含义如下：

其中expertNumPadded是expertNum按16（kMoeCombineMetadataPad）向上对齐的结果。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（变长 return）

kernel 遍历所有本地 expert segment：

segment = src_rank * expertPerRank + localExpert globalExpert = myRank * expertPerRank + localExpert rows = routeMeta.peerTokenPerExpert[src_rank, globalExpert]

对每个非空 segment：

1. srcStart由dispatchOffset[localExpert] + prevSumBeforeRank[src_rank, localExpert]计算；
2. dstStart由cumsumPerExpert[src_rank, globalExpert - 1]计算；globalExpert == 0时为0；
3. 如果src_rank == myRank，行被本地复制到本 rank 的peerWindow.ptrD（CopyLocalRowsToPeerWindow，逐行逐 tile 用TLOAD/TSTORE双缓冲搬运）；
4. 否则，PTOTPUT把行 chunk 写入 source rank 的远端 peer window（PutRemoteRowsToOwner）。

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 rank：

TNOTIFY(remotePeer.combineDoneSignal[myRank], AtomicAdd) TWAIT(localPeer.combineDoneSignal[peer] >= 1)

对应源码是NotifyCombineOwners（对所有src槽位做TNOTIFY(sig, 1, NotifyOp::AtomicAdd)，前置pipe_barrier(PIPE_ALL)保证 TPUT 全部落窗）与WaitCombinePhase（以blockId为起点按blockNum步进遍历所有 peer，TWAIT(sig, value, WaitCmp::GE)）。Host 会在每轮迭代前清零包含 head guard 和 live payload 的 HCCL window 区域，因此 kernel 固定等待每个 peer 的一次 notify；前置MPIBarrier只用于保证所有 rank 完成清理后再 launch kernel，不计入combine_e2e。

阶段 3: RestoreOutputRows（加权还原）

每个 AIV block 负责一段连续 token（TokenShardBegin/TokenShardEnd均匀切分）。对每个 token 和每个列 tile：

1. 使用TEXPANDS把输出 tile 清零；
2. 对每个有效 route，加载ptrD[expandedRowIdx]；
3. 使用TAXPY(outTile, ptrTile, prob)累加；
4. 将 fp16 tile 写回outputC。

还原阶段的内部循环在AccumulateRestoreTile中实现：使用pto::Event<pto::Op::TAXPY, pto::Op::TLOAD>把上一次TAXPY与下一次TLOAD串联成 event chain，保证 topK loop 顺序的同时，避免每个 route 后都插入完整 pipe barrier。

优化说明：从源码看设计取舍

该 kernel 是 AIV-only combine kernel。对于K=7168这类 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 block（ReturnExpertRowsToOwners中的分配循环）。
• PTOTPUTping/pong 路径：远端 return 使用TPUT(remoteDst, localSrc, ping, pong)，ping/pong 两块 UB tile 地址分别为0x0与0x1000，通过 UB 双缓冲让 MTE2 load 和 MTE3 store 形成流水。
• Restore route cache：当topK <= 16（kRouteCacheMax）时，每个 token 的 route row 和 prob 会缓存到标量数组（RestoreRouteCache{rows[16], probs[16], count}），减少内层 restore loop 对 route metadata 的重复读取。
• DCCI 批量 acquire：每个 token 在消费返回的ptrD行前先刷新对应 GM range（DcciGmRangeNoFence按 64B cacheline 逐个dcci），然后对本轮 cached routes 做一次dsb(DSB_DDR)，避免逐行刷新的同步开销。
• A5 restore event chain：A5 restore loop 使用 PTO event 串联TAXPY -> next TLOAD，减少 topK loop 内不必要的PIPE_ALLbarrier。
• Soft AIV sync：同一个 kernel 内用pto::SYNCALL<SyncAllMode::Soft>分隔 return、wait、restore 阶段（SoftSyncAiv），workspace 中的localSync区大小按aivBlocks * (8 + expertNumPadded)且不小于 64 个 int32 计算。

Tiling 与默认参数

常用调试命令--aiv-blocks 24，并使用PES=2, M=64, K=7168, topK=8, expertPerPe=2时，各布局大小为：

这些数字与 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 -> RestoreOutputRows

Return phase: routeMeta(peerToken/cumsum/offset) + expertOutput -> local or remote peerWindow.ptrD -> TNOTIFY peer combineDoneSignal[myRank] Restore phase: routeMeta.expandedRowIdx + probs + peerWindow.ptrD -> outputC

Host 侧的完整编排在 main.cpp 中体现：InitRankInfo（MPI rank/device 绑定）→PrepareHostData（确定性数据生成 + CPU golden）→InitHccl（root 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关闭）。

仅编译（手动 CMake）

cmake -S kernels/manual/a5/moe_combine \ -B /tmp/moe_combine_a5_build \ -DRUN_MODE=npu \ -DSOC_VERSION=Ascend950PR_958b cmake --build /tmp/moe_combine_a5_build --target moe_combine -j8

构建细节可参考 CMakeLists.txt：kernel 以 SHARED 库（moe_combine_kernel）形式用 Bisheng CCE 编译（--cce-aicore-arch=dav-c310-vec、--cce-fatobj-link、定义PTO_NPU_ARCH_A5），host 可执行文件链接runtime ascendcl hcomm等库。

A5 快速验证（小 shape）

cd kernels/manual/a5/moe_combine bash run.sh -pes 2 -M 8 -K 64 -topK 2 -expertPerPe 1 --aiv-blocks 24

默认 shape

cd kernels/manual/a5/moe_combine bash run.sh -pes 2 -M 64 -K 7168 -topK 8 -expertPerPe 2 --aiv-blocks 24

主要命令行参数

此外 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 路由账本（GenerateDeterministicInputs：token 输入按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/），便于逐字段排查。

预期成功输出：

verify=PASS

脚本输出的 profile 形式如下（A5 runtime 数字需要在 A5 机器上采集，当前开发环境可编译 A5 工程）：

[PROFILE] CombineTile M=64 K=7168 ranks=2 topK=8 expertPerPe=2 warmup=3 measured=5 samples=5 prepare_fixture: avg=... us max=... us combine_e2e: avg=... us max=... us verify=PASS

关键指标含义：

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