CANN pto-isa A2/A3 MoE Combine Kernel：基于 PTO TPUT 的变长 Return 与 Gate 加权还原实现-平芜编程栈

CANN pto-isa A2/A3 MoE Combine Kernel：基于 PTO TPUT 的变长 Return 与 Gate 加权还原实现

【免费下载链接】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 仓库中 moe_combine 算子文档为主体，结合 kernel 源码、布局计算与运行脚本逐层剖析这一 AIV-only 通信算子的完整契约：读懂它之后，你能掌握 PTO 通信指令（TPUT/TNOTIFY/TWAIT）如何在 HCCL peer window 上实现变长 all-to-all-like 的 MoE combine return，并用显式路由账本完成 token 级加权还原。

一、算子定位：dispatch-compute-combine 流水的 return 半段

本算子运行在 Ascend A2/A3 系列芯片（已在 Atlas 910B1 验证），是 MoE 推理流水 dispatch-compute-combine 三段中的最后一段 combine。本地 expert 完成计算后，combine kernel 承担两个职责：

Return：把expertOutput中的 expert 输出行，按上游整理好的路由账本routeMeta返还给 token 所在 rank；
Restore：用 gate 权重probs把topK路 expert 输出加权求和，还原出每个 token 的最终输出outputC。

文档给出的数据流概览：

expertOutput[local expert rows, K] -> 通过 HCCL peerWindow.ptrD 做变长 return -> 通过 TNOTIFY/TWAIT 做跨 rank 完成同步 -> 加权还原: outputC[token, :] = sum(topK probs * returned rows)

需要注意的是，当前 kernel 是一个独立 combine kernel，入口消费的是显式低层路由账本routeMeta。expert_ids、assist_info_for_combine、ep_send_counts等路由信息由上游或 host 侧按本算子的账本布局整理后显式写入routeMeta传入，kernel 本身不做路由决策。

目录结构与文件分工

kernels/manual/a2a3/moe_combine/ ├── CMakeLists.txt # Bisheng CCE + host 构建配置 ├── run.sh # 一键构建和运行脚本，发现 MPI，估算 HCCL_BUFFSIZE ├── 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 / golden.cpp # CPU golden 数据结构和路由构造、输出校验 ├── hccl_context.h # HCCL window 初始化和 peer-window 地址交换 ├── comm_mpi.h # MPI 动态加载封装 ├── README.md # 英文 README └── README_zh.md # 中文 README（本文主体）

从源码结构看，common.h 中的MoeCombineShape、WorkspaceLayout、CombineRouteMetaLayout、PeerWindowLayout四组 struct 是 host 与 device 共享的 ABI：host 侧 layout.h 用它们计算字节布局，device 侧 moe_combine_kernel.cpp 中的MakeWorkspaceLayout/MakeCombineRouteMetaLayout/MakePeerWindowLayout用同一套规则在 kernel 内部重算偏移，两侧字段名严格对齐，保证任何一边改动都会立刻暴露 ABI 不一致。

二、算子语义：从 routeMeta 到加权还原

2.1 输入与计算流程

对每个 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]。

2.2 加权还原公式

对本 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, :]；row < 0表示该 route 无效（例如未命中该 rank 的 expert），直接跳过。这一语义在 kernel 源码的LoadRestoreRoute中以*ptrDRow < 0分支精确实现，host 侧 CPU golden 在 golden.cpp 中用同一公式独立计算，供逐元素比对。

2.3 覆盖范围

包含	不包含
EP 域内基于 HCCL window 的 combine return	Dispatch pack/gather kernel
使用`TPUT`实现变长 all-to-all-like return	HCCL collective`AllToAllV`API
使用`probs`做加权还原	Expert FFN/GMM 计算
显式低层`routeMeta`契约	量化、TP ReduceScatterV、shared/copy/const expert
A2/A3 HCCL peer-window 路径	上层公共 ABI 适配层

三、入口契约：Kernel Launcher ABI 与缓冲区布局

3.1 Launcher ABI

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

该函数在 moe_combine_kernel.cpp 末尾以标准MoeCombineKernel<<<launchBlockCount, nullptr, stream>>>(" 三参数 launch 形式提交，launchBlockCount` 决定 kernel 使用的逻辑 AIV block 数。

3.2 运行时输入

参数	方向	存储	含义
`shape`	输入	值传递	静态 shape 和 AIV block 数，如`ep`,`m`,`k`,`topK`,`expertPerRank`,`aivBlocks`
`myRank`	输入	值传递	EP 域内 rank id
`expertOutput`	输入	`aclrtMalloc`GM	本地 expert 输出行，形状`[maxOutputSize, K]`，fp16
`probs`	输入	`aclrtMalloc`GM	gate 权重，形状`[M, topK]`，fp32
`outputC`	输出	`aclrtMalloc`GM	还原后的 token 输出，形状`[M, K]`，fp16
`routeMeta`	输入	`aclrtMalloc`GM	显式 combine 路由账本
`peerWindow`	输入/输出	HCCL RDMA window	远端可见的`ptrD`return buffer 和 signal
`hcclCtx`	输入	`aclrtMalloc`GM	设备侧所有 rank 的 HCCL window 地址
`workspace`	临时	`aclrtMalloc`GM	本地 AIV soft sync 区
`stream`	输入	ACL stream	kernel launch stream
`launchBlockCount`	输入	值传递	kernel 使用的 AIV block 数

3.3`MoeCombineShape`字段

字段	含义
`ep`	EP rank 数
`m`	每 rank token 数
`k`	hidden size
`topK`	每 token 的 expert 路由数
`expertPerRank`	每 rank 本地 expert 数
`expertNum`	全局 expert 数，通常为`ep * expertPerRank`
`maxOutputSize`	每 rank expert 输出最大行容量
`aivBlocks`	逻辑 AIV block 数；A3 默认`24`，可传参覆盖

aivBlocks为 0 时的降级行为在源码中有明确定义：layout.h 的 EffectiveAivBlocks 将其视为 1，而 run.sh 会把命令行传入的 0 替换为默认值 24，两条路径的默认值不同，分别服务于"host 布局计算"和"实际运行"两个场景。

3.4`peerWindow`内容

localWindowBase是 HCCL window 的起始地址。A2/A3 上传给 kernel 的peerWindow指向localWindowBase处的 live payload；本 layout 不额外保留 A5 的 4096B head guard。

A2/A3 localWindowBase peerWindow live payload: ptrD countReadySignal[ep] combineDoneSignal[ep]

字段	位置	内容
`ptrD`	HCCL window live payload	return 目标行，被远端`TPUT`写入
`countReadySignal[ep]`	HCCL window live payload	per-rank ready 计数区
`combineDoneSignal[ep]`	HCCL window live payload	per-rank 完成计数器；远端 rank 完成写入本 rank`ptrD`后`TNOTIFY`对应槽位

3.5`routeMeta`布局

routeMeta是显式低层 combine 路由账本，属于本地 GM，不属于 HCCL window：

字段	形状	含义
`peerTokenPerExpert`	`[ep, expertNumPadded]`int32	每个 source rank 到每个 global expert 的行数
`expandedRowIdx`	`[M * topK]`int32	token route 到`peerWindow.ptrD`的行映射；`-1`表示无效 route
`cumsumPerExpert`	`[ep, expertNumPadded]`int32	每个 source rank 内按 global expert 的 inclusive prefix：`cumsum[src,e] = sum(peerTokenPerExpert[src,0..e])`
`dispatchOffset`	`[expertPerRank]`int32	每个本地 expert 在`expertOutput`中的基地址行
`prevSumBeforeRank`	`[ep, expertPerRank]`int32	某 source rank 在本地 expert 行段中的前缀偏移

expertNumPadded是expertNum向上对齐到 metadata pad 粒度 16（kMoeCombineMetadataPad）后的值，用于让peerTokenPerExpert/cumsumPerExpert两行的行宽固定，方便设备侧以整行 stride 寻址。

3.6 布局计算与 HCCL_BUFFSIZE 估算

三份布局（workspace / routeMeta / peerWindow）的 host 实现集中在 layout.h。关键规则：

64 字节字段对齐：每个字段追加前先把游标对齐到 64B（AppendField），与 device 侧AppendFieldDevice完全一致；
workspace：只含localSync软同步区，槽位数是aivBlocks * (8 + expertNumPadded)，不足 64 槽时取 64；
peerWindow：ptrD占M * topK * K * 2字节（fp16），后接两个[ep]int32 信号数组；
溢出保护：所有乘法经CheckedMul检查，避免 shape 参数异常导致 uint64 溢出。

run.sh会先用同样的 bash 逻辑预估三类布局字节数，再按EstimateHcclBuffSizeMb的规则（layout.h L122-L128）估算 HCCL 缓冲：peerWindow 总字节 + 64 MiB 安全余量，向上对齐到 MiB 后导出HCCL_BUFFSIZE；用户也可用--hccl-buffsize-mb手动覆盖。

四、Kernel 细节：三阶段实现

Kernel 入口 MoeCombineKernel 的骨架非常清晰：

ReturnExpertRowsToOwners(...); // 阶段 1: return WaitCombinePhase(...); // 阶段 2: 跨 rank 等待 SoftSyncAiv(...); // 软同步：进入 restore 前全体 block 会合 RestoreOutputRows(...); // 阶段 3: 加权还原 SoftSyncAiv(...); // 软同步：收尾

其中SoftSyncAiv包装的是 PTO 软同步指令pto::SYNCALL<Soft>（源码 L251-L255），它只在同一个 kernel 的多个 AIV block 之间建立栅栏，代价远低于 kernel 级同步——这也是文档中"Soft AIV sync"优化项的落地方式。

4.1 阶段 1：ReturnExpertRowsToOwners

kernel 遍历所有本地 expert segment（共ep * expertPerRank个）：

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

对每个非空 segment（LoadReturnSegment）：

srcStart = dispatchOffset[localExpert] + prevSumBeforeRank[src_rank, localExpert]，即本 source rank 的行在本 expert 行段内的起点；
dstStart = cumsumPerExpert[src_rank, globalExpert - 1]（globalExpert == 0时为 0），即该行段在 owner rank 的ptrD中的目标行；
按固定 8 行 chunk 切分（kMoeCombineRowChunk），以chunkBase % blockNum轮转分配给各 AIV block，实现 L471-L483 的 load balance；
src_rank == myRank时走本地路径CopyLocalRowsToPeerWindow：用TLOAD -> TSTORE把行经 UB tile 复制到本 rank 的peerWindow.ptrD（L371-L405）；否则走远端路径PutRemoteRowsToOwner。

远端路径的核心是一次 PTO 通信指令调用（L407-L427）：

pto::comm::TPUT(remoteDst, localSrc, ping, pong);

remoteDst不是本 rank 的地址，而是先用 RemotePtr 把本地指针相对本 rank window base 的偏移，平移到hcclCtx->windowsIn[peerRank]上得到远端 window 内同布局的ptrD地址——这正是 HCCL window "同偏移寻址"契约的体现。ping/pong是分配在 UB 两个地址（0x0/0x1000）上的双缓冲 tile，让 MTE2 load 与 MTE3 远端 store 形成流水。该指令的语义可参考仓库通信 ISA 文档 TPUT。

所有 return 完成后，每个 AIV block 对分配给它的 source rank 集合执行通知（NotifyCombineOwners）：

pto::comm::TNOTIFY(sig, 1, pto::comm::NotifyOp::AtomicAdd); // sig = 远端 combineDoneSignal[myRank]

语义可参考 TNOTIFY。

4.2 阶段 2：WaitCombinePhase

每个 rank 等待所有 peer 的写入完成（源码 L263-L273）：

TNOTIFY(remotePeer.combineDoneSignal[myRank], AtomicAdd) // 阶段 1 尾部发出 TWAIT(localPeer.combineDoneSignal[peer] >= 1) // 本阶段等待

等待谓词是WaitCmp::GE、阈值kMoeCombineSignalValue = 1。Host 会在每轮迭代前清零combineDoneSignal（main.cpp 的 ClearDeviceState 中aclrtMemset整块 peerWindow），因此 kernel 固定等待每个 peer 恰好一次 notify。TWAIT的语义参见 TWAIT 文档。

4.3 阶段 3：RestoreOutputRows

每个 AIV block 负责一段连续 token（TokenShardBegin/End做余数感知的均匀切分）。对每个 token 的每个 1024 列 tile：

TEXPANDS(outTile, 0.0)把输出 tile 清零；
对每个有效 route，TLOAD加载ptrD[expandedRowIdx]的一行片段；
TAXPY(outTile, ptrTile, prob)完成y = A*x + y形式的累加；
全部 route 累加完毕后TSTORE把 fp16 tile 写回outputC。

这一步有两条值得注意的实现细节：

Restore route cache + DCCI 批量 acquire（PrepareRestoreRouteReads）：当topK <= 16时，token 的全部 route 行号与 prob 先缓存进标量数组；每条 route 的ptrD行用dcci逐 cache line 刷新后只做一次dsb(DSB_DDR)，把内层 restore loop 中逐条 route 的 GM acquire 开销摊薄为每 token 一次。
TLOAD -> TAXPY event chain（AccumulateRestoreTile）：利用 PTO 的 event 依赖机制，pto::Event<Op::TAXPY, Op::TLOAD>串联上一轮 AXPY 与下一轮 LOAD，pto::Event<Op::TLOAD, Op::TAXPY>串联本轮 LOAD 与 AXPY，形成 load/axpy 交替流水，避免显式 flag 等待打断指令发射。

五、优化手段汇总

该 kernel 是 AIV-only combine kernel。对于K=7168这类 hidden size，一行 fp16 数据是 14 KiB，整体主要受 GM/HCCL window 搬运带宽影响。优化目标是让数据搬运尽量流式化，同时降低控制面元数据开销：

优化	说明	源码佐证
显式 routeMeta	路由元数据作为独立 GM buffer 传入；`peerWindow`只保留远端可见 return 数据和信号，`workspace`只保留本地 AIV soft sync 区	common.h ABI
chunk 化 return 分片	return 阶段遍历`src_rank x local_expert`segment，按`chunkBase % blockNum`把 8 行 chunk 分给 AIV block	L470-L483
PTO TPUT ping/pong 路径	`TPUT(remoteDst, localSrc, ping, pong)`通过 UB 双缓冲让 MTE2 load 与 MTE3 store 流水	L422-L426
Restore route cache	`topK <= 16`时 route 行号与 prob 缓存到标量数组，减少内层 loop 对 route metadata 的重复读取	L488-L531
DCCI 批量 acquire	每 token 消费`ptrD`行前按 64B cache line 刷新对应 GM range，本轮 cached routes 只做一次`dsb(DSB_DDR)`	L202-L228
TLOAD 到 TAXPY event chain	restore 阶段通过 PTO event 依赖完成加载与计算衔接	L533-L558
Soft AIV sync	同一 kernel 内用`pto::SYNCALL<Soft>`分隔 return、wait、restore 阶段	L251-L255

六、Tiling 与默认参数

参数	默认值	说明
`PES`/`ep`	`2`	EP rank 数
`M`	`64`	每 rank token 数
`K`	`7168`	hidden size
`topK`	`8`	每 token expert 路由数
`expertPerPe`	`2`	每 rank 本地 expert 数
`expertNum`	`4`	`PES * expertPerPe`
`maxOutputSize`	`PES * M * topK`	默认容量；默认 shape 下为`1024`
`aivBlocks`	`24`	A3 默认逻辑 AIV block 数
内部 Vector tile 列宽	`1024`	示例实现固定值（`kMoeCombineTileCols`）
内部 return chunk	`8 rows`	固定的 return 阶段行 chunk（`kMoeCombineRowChunk`）
内部 metadata pad	`16`	expert metadata 对齐粒度（`kMoeCombineMetadataPad`）

使用PES=2, M=64, K=7168, topK=8, expertPerPe=2, aivBlocks=24时，各布局大小为：

Layout	字节数
`workspace`	`2304`
`routeMeta`	`2432`
`peerWindow`	`7340160`

这三个数字可以由 layout.h 的公式直接推导验证：例如peerWindow = 64*8*7168*2 + 2*2*4 + 对齐 = 7340032 + 信号区 ≈ 7340160（对齐到 64B）。run.sh启动时会打印workspace_bytes / route_meta_bytes / peer_window_bytes，与上表一致。

七、Host 侧整体架构

Host: ParseArgs -> ComputeWorkspaceLayout / ComputeCombineRouteMetaLayout / ComputePeerWindowLayout -> PrepareHostData and CPU golden -> Init HCCL peer window -> 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

从 main.cpp 的编排看，几个环节值得注意：

HCCL window 初始化：rank 0 通过 hccl_context.h 中的InitHcclRootInfo获取 root 信息，经 MPI broadcast 给所有 rank，再由InitHcclWindowContext完成 window 分配与 peer-window 地址交换，最终所有 rank 的 window base 填入设备侧HcclDeviceContext.windowsIn/windowsOut（结构定义见 common.h L58-L66），供 kernel 内RemotePtr寻址；
MPI 封装：rank 映射默认从mpirun环境获取（--rank-from-mpi 1），MPI 符号通过 comm_mpi.h 动态加载，run.sh会先检查mpirun是否存在；
每轮迭代清理：ClearDeviceState用aclrtMemset清零 workspace、routeMeta 与整个 peerWindow（含combineDoneSignal），保证 TWAIT 的"等待一次 notify"语义每轮都成立；
Profile 聚合：rank 0 用MpiGatherBytes收集各 rank 的prepare_fixture / combine_e2e / total_e2e计时，按迭代取跨 rank 最大值后打印，避免个别慢 rank 掩盖统计。

八、构建与运行

8.1 环境

source /usr/local/Ascend/cann-8.5.0/set_env.sh

执行run.sh前需要先在 shell 中加载 CANN 环境。如果 shell 中没有mpirun，请先配置 MPI 环境（脚本会直接报错退出）。run.sh还会把Ascend910B1simulator 库目录加入LD_LIBRARY_PATH，并按--keep-hccl-shm决定是否清理/dev/shm/sem.hccl*与共享内存段。

8.2 仅编译

cd kernels/manual/a2a3/moe_combine bash run.sh --skip-build 0 --clean-build 1

脚本内部流程是cmake -DRUN_MODE=npu -DSOC_VERSION=Ascend910B1 .. && make -j16，产物为build/moe_combine可执行文件，再以mpirun -n ${PES}拉起全部 rank。

8.3 快速验证

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

8.4 默认 shape

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

8.5 主要命令行参数

参数	默认值	含义
`-pes`	`2`	rank 数
`-M`	`64`	每 rank token 数
`-K`	`7168`	hidden size
`-topK`	`8`	每 token route 数
`-expertPerPe`	`2`	每 rank expert 数
`--max-output-size`	`PES * M * topK`	expert output 行容量
`--aiv-blocks`	`0 -> 24`	逻辑 AIV block 数，用于匹配不同硬件资源规划
`--device-base`	`0`	rank 到 device 映射使用的起始 device id
`--ndevices`	`PES`	示例 launcher 使用的可见 device 数

脚本还会强制两条约束（run.sh L133-L154）：所有 shape 字段必须为正数，且maxOutputSize >= PES * M * topK（不支持容量不足导致的丢行）。此外run-mode第一版只支持npu，--case预设不被接受，必须显式传 shape 参数。

8.6 实测性能与 profile 指标

工程可直接在 A2/A3 机器（Atlas 910B1）上运行，脚本输出如下 profile：

[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

关键指标含义：

指标	含义
`combine_e2e`	combine kernel launch 到 stream sync；不包含 clear、fixture、verify，也不包含 kernel launch 窗口之外的 MPI barrier
`verify=PASS`	device`outputC`与 CPU golden 一致

计时口径上，warmup 迭代被排除，每个 measured 样本取跨 rank 最大值（对应 main.cpp 的 PrintProfileSummary）。

九、验证方法

Host 会构造确定性的 CPU golden 路由账本（含无效 route 覆盖），将其写入routeMeta，拷贝expertOutput，启动 kernel，再将outputC与 CPU golden 输出逐元素对比（相对/绝对容差默认均为1e-2）。默认开启验证，预期成功输出：

verify=PASS

十、小结与延伸阅读

moe_combine 展示了 PTO 通信指令集在 A2/A3 上的一个完整落地范式：

数据面：以 HCCL window 同偏移寻址 +TPUT变长远端写入，绕开 collective 库，直接实现 all-to-all-like 的 return 数据流；
控制面：routeMeta显式账本（peerToken / cumsum / offset 三类元数据）把 dispatch 阶段的路由决策与 combine 执行解耦；
同步面：TNOTIFY(AtomicAdd) + TWAIT(GE 1)完成跨 rank 屏障，pto::SYNCALL<Soft>完成 rank 内 AIV block 之间的阶段分隔；
性能面：UB ping/pong 双缓冲、event chain、DCCI 批量 acquire 等技巧把控制面开销压到最低，让 14 KiB/行的大数据搬运尽量流式化。

CANN pto-isa A2/A3 MoE Combine Kernel：基于 PTO TPUT 的变长 Return 与 Gate 加权还原实现