Ascend DeepEP MoE Dispatch Optimization
问题不只是“带宽没跑满”
看到吞吐只有硬件理论值的一半,第一反应通常是增加 AIV、增加 QP、做双缓冲或者把通信和计算异步起来。这些方向并非错误,但它们默认了一个尚未证明的前提:当前瓶颈就是负责搬 payload 的硬件引擎。
MoE dispatch 的端到端时间更接近:
[
T_{dispatch} = T_{layout} + T_{notify} + T_{route\ scan} + T_{WQE\ submit}
- T_{network} + T_{quiet/barrier} + T_{epilogue}
]
这几项不一定完全串行,实际延迟取决于依赖关系和重叠程度;多 rank 场景最终看到的又往往不是平均值,而是最慢 rank:
[
T_{step} \approx \max_{r \in ranks} T_{dispatch}^{(r)}
]
于是,“理论带宽的一半”至少可能来自五类不同问题:
- 数据面受限:hidden payload 真正压满了链路、UDMA 引擎或 GM 读写带宽。
- 控制面受限:单包太小,WQE 生成、提交、CQE/quiet 的固定成本占比过高。
- 路由规划受限:Top-K 很小,却仍按
[num_tokens, world_size]扫描和保存路由。 - 拓扑受限:128 个 peer 同时展开,跨机链路、交换框或某条 rail 出现热点。
- 尾延迟受限:平均 rank 不慢,但少数 expert/rank 收到更多 token,所有人最终等它完成。
这也是后续读代码时最重要的视角:不要只找“拷贝循环”,要同时找 payload 在哪里走、offset 谁算、完成由谁确认、下一阶段在等谁。
一条 token 如何走完 MoE
设 EP 域共有 (P) 个 rank,每个 rank 持有 (E_l) 个本地专家;输入 x 的 shape 为 [S, H],Router 为每个 token 选择 (K) 个专家:
1 | x[token] 一行 hidden |
一个全局专家 expert_id 对应:
1 | dst_rank = expert_id / experts_per_rank |
如果一个 token 的 Top-2 专家分别位于 rank 0 和 rank 3,dispatch 就要把同一行 hidden 复制到两个目标 rank;目标 rank 按本地专家重排成连续输入,完成 expert FFN 后,再把两份结果发回 token owner,最后按 Top-K 权重归约。dispatch 是一趟不等长、由路由结果驱动的 All-to-AllV;combine 是携带反向索引的逆过程。
下面这张图把容易混在一起的三条线分开:蓝色是 payload,绿色是路由与 offset,紫色是完成和缓冲区复用协议。
一个正确的单边通信顺序通常是:
1 | 计算目标地址和接收 offset |
如果 flag 先于 payload 可见,接收方就可能读到旧数据;如果缓冲区在 quiet 前被复用,后续 WQE 可能读到被覆盖的源数据。因此,通信和计算能否重叠,首先是依赖图与内存可见性问题,其次才是调度技巧。
从 Python 接口进入真实调用链
本节以 ascend_deepep@18e2dd4b 为代码锚点。Python 入口位于 ascend_deepep/buffer.py::Buffer.dispatch,它先区分是否提供 handle。
1 | if handle is not None: |
没有 handle 时,需要传入当前路由的统计和矩阵:
1 | if num_tokens_per_rank is None: |
随后执行 native dispatch,并把这次生成的路由计划打包成 v1_handle:
1 | v1_handle = ( |
这段接口已经透露了一个重要优化边界:hidden 可以每轮变化,但如果路由关系不变,prefix、offset 和反向索引没有必要每轮重建。cached dispatch 就是在复用这部分控制面工作。
Host 侧先造一张通信地图
进入 csrc/ops/dispatch.cpp 后,Host 侧主要完成四件事:校验输入、规划 symmetric workspace、启动 notify、创建输出 view 并启动真正的 dispatch kernel。
Workspace 不是一个 Tensor,而是一块地址空间
代码首先解析 AIV 数量:
1 | const int64_t resolved_num_aiv = ResolveNumAiv(num_aiv); |
随后按极端情况预留内部接收容量:
1 | const int64_t recv_buffer_tokens = CheckedMulInt64( |
原因是:每个 source rank 最多有 num_tokens 行可以发给当前 rank;共有 world_size 个 source rank,所以当前实现给内部接收区保留 num_tokens × world_size 行。
MakeDispatchWorkspaceLayout() 再用一个不断前移的 cursor,把一整块 symmetric buffer 切成多个互不重叠的区域:
1 | layout.recv_channel_prefix_offset = cursor; |
这些 offset 随后写入 tiling,Device kernel 只要拿到 symmetric buffer 基地址,就能定位 payload、反向索引、Top-K 元数据、scale 和 completion signal。
Notify 先交换数量和前缀
Host 在 payload dispatch 前启动 notify kernel:
1 | LaunchNotifyDispatchKernel( |
notify 的核心任务不是提醒一句“开始接收”,而是把不规则路由转成通信双方都能使用的边界:当前 rank 实际要收多少 token、每个本地 expert 收多少、每个 source rank 和 channel 应写到接收区的什么位置。真正的 payload kernel 依靠这些 prefix,才能让多个 AIV 和多个 rank 并行写入而不互相覆盖。
输出长度有动态和固定两种模式
notify 之后,Host 决定公开输出 Tensor 的第一维:
1 | if (cached_mode) { |
然后用 SHMEM 内部地址创建非 owning view:
1 | auto recv_x = MakeShmemTensor( |
最后才启动 LaunchV1DispatchKernel()。因此 Host 路径可以概括成:
1 | validate |
Device 侧如何把 channel 变成远端 Put
普通 dispatch.cpp 的核心所有权模型是:一个 AIV 负责一个 channel,并维护面向所有 destination rank 的游标。
1 | for (int32_t dst_rank = 0; dst_rank < world_size; ++dst_rank) { |
随后它遍历自己负责的 token 范围和全部目标 rank:
1 | for (int32_t token = token_begin; token < token_end; ++token) { |
这段代码简单直接,但它也暴露了 128P 的一个结构性成本:即使 Top-K 只产生少数目标 rank,每个 token 仍可能扫描完整的 world_size 维度。
UDMA 版本在 peer 调度上又增加了两个策略:按 source rank 旋转目标顺序,以及用两个 QP 承载 channel。
1 | for (int32_t logical_destination = destination_worker_index; |
两个 QP 采用 3:1,而不是均分
当前代码固定两个 QP,并按 channel 编号分配:
1 | __aicore__ inline int32_t DispatchUdmaQpForChannel( |
也就是:
1 | channel: 0 1 2 3 | 4 5 6 7 |
所有 peer 的 payload 提交完成后,每个 QP 分别追加一个 Strong Order completion,再执行 qp_quiet(dst_rank, qp_slot);最后才是:
1 | aclshmem_quiet(); |
128P 为什么会放大问题
小规模下,“每个源 rank 直接面向每个目标 rank”简单且可能足够快;扩展到 128P 后,数据量未必线性增长,控制面却很容易增长:
1 | 路由表示: [S,P] |
与此同时,每个 peer 的消息变小,固定 WQE/quiet 成本占比上升;不同源 rank 若在相近时刻访问相同拓扑区域,还会形成热点。这里要解决的已经不是单核 DataCopy 的问题,而是全系统同时有哪些通信关系处于活跃状态。
下面把三套代码的调度单位放到同一张图里。
ascend_deepep:扁平直达
当前 UDMA kernel 通过 rank 旋转错开起始目标,但没有显式的 server/frame/rail 分组。每个 source rank 仍会依次处理全部 destination rank;每个 peer 的两个 QP 完成后分别 quiet,最后进入全局 barrier。
这个方案的优点是地址关系直接、没有额外 relay;风险则是 128P 下 peer 扇出、路由扫描、completion/quiet 和尾延迟同时放大。若性能从 64P 到 128P 明显掉档,而 hidden size 增大后有效带宽反而改善,应优先怀疑固定控制成本与拓扑调度,而不是立即优化向量计算。
海思 hierarchy:跨机只面向 server
ops-transformer@8dc2fa04 的公开说明把两种算法写得很直接:
1 | fullmesh: |
Layered kernel 固定:
1 | constexpr static uint32_t SERVER_RANK_SIZE = 16; |
跨机目标中继 rank 的计算是:
1 | uint32_t dstRankId = |
也就是说,source rank 的 local lane 为 5 时,跨机只和目标 server 的 lane 5 通信;数据到达后,再通过 Win2Ipc()、SendTokenToIpc()、Ipc2Out() 等机内路径转发到真正持有 expert 的 rank。
举一个抽象例子:每 server 16 rank,目标 expert 在 server 2 的 rank 42,而 source rank 的 local lane 是 5。跨机阶段先发送到:
1 | 2 × 16 + 5 = rank 37 |
rank 37 是 server 2 的同号中继;随后 server 2 内部再把 hidden 送到 rank 42。若同一 token 在 server 2 上命中多个专家,跨机阶段还可以按 server 聚合 hidden,避免为同一 server 重复发送多份大 payload。
代价也很明确:多一次 relay、机内 copy 与控制协议。只有当减少的跨机 peer、重复 hidden 和 RDMA 下发成本,大于新增机内搬运成本时,hierarchy 才会获益。
MoonEP:64P group 与 128P 拓扑波次
如果更想先理解设计动机,可打开 MoonEP Group Design:分步交互解释,从直接发送、接收竞争、错峰与放行许可逐步回到 64P;下面保留源码层面的分组细节。
先澄清实现归属。MoonshotAI 的原版 MoonEP 解决的是动态冗余 expert、固定每 rank 工作量与 zero-copy 通信;下面的 64P/128P group 调度来自固定 revision 的 Ascend 移植内核,不是 MoonEP 算法天然规定的拓扑。[^moonep-64p-group]
64P 不是连续 16 rank 分组
64P 走的是 Legacy 映射,而不是后面 128P 专用的 two-frame 映射。kernel 先固定 group_width = 16,再把 64 rank 均分为两个 32-rank half。每个逻辑 group 从两个 half 各取 8 个同号 lane:
[
G_q={8q+l\mid 0\le l<8}\cup
{32+8q+l\mid 0\le l<8},\qquad q\in{0,1,2,3}.
]
因此实际成员是:
| 逻辑 group | half 0 | half 1 |
|---|---|---|
G0 |
rank 0–7 | rank 32–39 |
G1 |
rank 8–15 | rank 40–47 |
G2 |
rank 16–23 | rank 48–55 |
G3 |
rank 24–31 | rank 56–63 |
这里最容易误读的是 group_rank_index。它不是“当前 rank 所属 group 的编号”,而是一个 0–15 的组内 lane:前 8 个 lane 指向 half 0,后 8 个 lane 指向 half 1。每个 rank 启动的前 16 个 AIV 分别占一个 lane,所以一个 source rank 在当前轮会同时面向目标 group 的 16 个 rank;后 16 个 AIV 则与 token 搬运并行处理 weight。
下面保留决定 64P 映射的源码主干;省略的是 128P 专用分支、参数校验和无关样板。group_rank_index 来自当前 AIV 编号,而 rank_group 才是 source rank 所属的逻辑 group:
1 | const uint32_t group_rank_index = |
其中 rank 是当前 source rank,group 是循环中的逻辑轮次 g,不是一个固定 group ID;返回值才是这个 AIV 在当前轮负责的 dst_rank。[^moonep-64p-group]
目标 group 按四轮循环平移。若 source rank 属于 Gq,其 group 编号为 (q_{src}),逻辑轮次为 (g),则目标编号为:
[
q_{dst}=(q_{src}+g)\bmod 4.
]
| 轮次 | G0 的目标 |
G1 的目标 |
G2 的目标 |
G3 的目标 |
|---|---|---|---|---|
g=0 |
G0 |
G1 |
G2 |
G3 |
g=1 |
G1 |
G2 |
G3 |
G0 |
g=2 |
G2 |
G3 |
G0 |
G1 |
g=3 |
G3 |
G0 |
G1 |
G2 |
b67c6840。我第一次枚举源码公式时,把“每轮双射”误写成了 rank-to-rank 双射,程序马上给出了反例。准确的不变量其实分两层:
- group 层是双射。 每轮的
q → (q+g) mod 4是四个 source group 到四个 destination group 的置换,所以同一轮不会有两个 source group 集中到同一个 destination group。 - group 内是全连接。 一旦
Gq → Gp配对,Gq的 16 个 source rank 都各自用 16 个 token AIV 覆盖Gp的 16 个 destination rank,形成完整的 (K_{16,16})。四轮结束后,每个 source rank 恰好覆盖全部 64 个 destination rank。
换句话说,group 调度没有减少 All-to-All 的总语义;它把一次面对 64 个 peer 的控制与下发过程,整理成每轮 16 个 peer、共四轮的波次。
源码中的 completion/credit 链又给波次加了一层流控。每个 token AIV 向当前 dst_rank 追加 completion signal,同时等待由 GetGroupIncomingPeer() 反算出的 source peer;收到后,接收方才向下一轮的 source peer 发布 credit。64P 及以下还会在 completion 后和 group 尾部各执行一次 16-AIV dissemination barrier,四步传播即可让 16 个 token AIV 锁步。于是某个 AIV 不能在其他 lane 尚未完成时独自冲进下一轮;128P 则关闭这层严格 group 同步,让各 AIV 沿 completion/credit 链异步推进。
128P 改用 two-frame 映射
ascend-moonep@b67c6840 为 128P 写出了明确的生产拓扑假设:
1 | constexpr uint32_t kDispatchFrameCount = 2; |
它把 128 rank 看成两个 64P frame;每个 frame 是 8 × 8,再切成四个 4 × 4 = 16 rank 的 group。全局共有 8 个 group,因此每个 rank 用 8 个逻辑轮次覆盖所有目标 group。GetGroupDstRank() 保证每轮映射是双射,使每个目标 group 在一轮中只接收一个源 group,避免多组同时集中冲击同一目标。
这里最特殊的不是数字 16,而是 rank ID 被当成物理坐标使用:frame、row、column、group side 都由 rank 算术得到。如果真实 rank table 没有按这套物理位置编号,照搬公式可能主动制造热点。
三种策略可以用同一组问题比较,而不是只看函数名:
| 实现 | 同时面向的通信单位 | 跨机路径 | 控制协议 | 主要边界 |
|---|---|---|---|---|
ascend_deepep |
最终覆盖全部 peer;按 source rank 旋转顺序 | source → destination rank | 两个 QP、per-peer quiet、全局 barrier | 无显式 topology group |
| 海思 hierarchy | 先 destination server,再 server 内 rank | source lane → 目标 server 同 lane 中继 | RDMA + HCCS/IPC + flag | 当前公开支持范围不是 128P |
| MoonEP | 每轮一个 16-rank group,8 轮覆盖 128P | 由 2×64、8×8 坐标映射决定 | 1–4 QP、completion/credit、group schedule | 拓扑与 rank 编号假设很强 |
P 维路由为什么值得单独优化
当前 is_token_in_rank 是 [S,P] 布尔矩阵:
1 | is_token_in_rank[token, rank] = |
它的优点是判断简单,缺点是在 128P 下表示和扫描都按 (P) 扩展。假设 Top-K 为 8,而这些专家最多落在 8 个 rank 上,那么每个 token 实际只需要保存少量 destination,却要读取 128 个布尔值。
更紧凑的 plan 可以写成:
1 | token 0 → [dst rank 3, dst rank 17] |
hierarchy 模式则先压成 destination server,再在 server 内展开 rank/expert:
1 | token |
这样既能减少 [S,P] 扫描,也能识别“同一 token 在同一 server 命中多个专家”,为跨机 hidden 去重创造条件。需要注意,plan 生成本身也有成本;如果把 planning 时间排除在 benchmark 外,就可能得到一个对端到端系统无效的优化结论。
先把时间和字节测对
在改调度前,需要建立一个同时覆盖 Host、Device、网络和多 rank 尾部的测量面。最少应该回答:
1 | 本 rank 实际发送/接收了多少 hidden 字节? |
有效 payload 带宽应按真实路由计算,而不是直接拿 S × P × H 当分子:
[
BW_{effective}^{(r)} =
\frac{actual\ received\ hidden\ bytes^{(r)}}
{T_{dispatch}^{(r)}}
]
端到端报告同时给出平均延迟和 max-rank latency。只看 rank 0 或所有 rank 平均值,可能完全看不到负载不均和拓扑热点。
Host 侧计时
Host 侧适合用 NPU event 包住稳定的阶段边界:
1 | event_start |
不要每个小阶段后立刻 Host synchronize,否则测到的是被人为串行化后的路径。正确性调试阶段可以同步定位,性能测量阶段则应连续提交多轮、丢弃 warmup,并明确是否冲刷 L2。
Device 侧计时
MoonEP 已经展示了一个适合借鉴的模式:用编译期 Profile 开关读取 cycle,把每个 AIV 的阶段周期和字节数写入独立 GM 槽位;Profile=false 时编译成零开销路径。
建议每个 AIV 至少记录:
1 | route_scan_cycles |
记录位置必须按 [rank, aiv, metric] 或其他独占方式切开,不能为了计时再引入热点原子加。最后由 Host 一次性读取并聚合。
把优化变成可证伪的实验
优化顺序不应是“哪个技巧听起来高级就先做哪个”,而应是每一步只验证一个瓶颈假设。
P0:建立可信基线
先固定:
S/H/K/P、dtype、expert 数、capacity/alignment。- rank 到 server、frame、交换路径和 local lane 的映射。
- 实际 route histogram:每个 token 的 unique destination rank/server 数,每个 rank/expert 的接收量。
- 计时是否包含 layout/notify/plan,是否使用 cached handle,是否使用
num_worst_tokens。 - 带宽分子是 hidden payload、全部元数据,还是链路物理字节。
同时做小规模分解:8P、16P、32P、64P、128P;均匀路由和倾斜路由;小 H 和大 H。只有 128P 掉档时,优先检查拓扑和 P 维控制面;所有规模都只有一半时,再检查数据路径、统计口径和底层引擎。
P1:先处理拓扑扇出
低风险起点是保留 rank 旋转,并把“全 128 peer”改成可调波次,例如每轮 8/16/32 peer,记录:
1 | group width |
如果硬件拓扑确认与 MoonEP 的 2×64、8×8、4×4 一致,可以研究它的双射 group schedule;如果真实系统更接近每 server 16 rank,则可研究 hierarchy 的“同 lane 跨机 + server 内展开”。必须先建立 rank→物理位置表,再选择公式。
P2:减少 WQE 和完成等待
对相同 peer 的连续 token 做聚合,使一次 WQE 搬更多字节;把 quiet 放在一批 payload 之后,而不是每个小块之后,同时保证 completion flag 晚于 payload 可见。
两个 QP 的分流不要只看 channel 数。把策略做成可配置,对比:
1 | 1 QP |
若增加 QP 只增加了 WQE/CQE 管理而没有降低 quiet 周期,瓶颈可能在共享链路或共享引擎,不在单 QP 队列深度。
P3:压缩路由计划
从 Top-K expert 直接生成 unique destination rank/server list,并按 destination compact;尽量融合:
1 | expert → rank/server 映射 |
目标是避免多次读取 [S,P],也避免所有 rank 重复读取完整 source-rank 计数。计划生成要纳入端到端计时,除非业务确实能长期复用 cached handle。
P4:再做细粒度通算掩盖
把工作拆成真正无依赖或仅有局部依赖的阶段:
1 | 当前 peer group 的 UDMA 在途 |
全局 barrier 若只用于保护某个 peer buffer,可以尝试改成 per-peer/per-group completion 与 credit;但必须保留“payload 可见后才能发 ready”“consumer 完成后才能复用 slot”这两条 happens-before。
P5:最后压内存细节
在确认 GM/UB 搬运成为热点后,再处理:
- 让同一 peer/expert 的 token 在 GM 中连续,扩大 DataCopy 和 WQE 粒度。
- 保证 UB 起点与传输粒度对齐,避免邻近 AIV 写同一缓存行。
- 使用 ping-pong staging,把 MTE2 读、Vector 处理、MTE3/UDMA 提交重叠。
- 复用 symmetric workspace 与 cached plan,避免热路径分配和初始化。
- 针对常见
H/K/P做静态 specialization,同时保留通用 fallback。
一份可以直接执行的实验表
| 优先级 | 假设 | 最小改动 | 必看指标 | 支持假设的现象 |
|---|---|---|---|---|
| P0 | 带宽口径或尾延迟判断有误 | 记录实际 payload 和所有 rank 延迟 | bytes、avg/p95/max rank latency | max 明显高于平均,或实际字节远小于理论分子 |
| P1 | 128P peer 扇出/拓扑是主瓶颈 | group width 扫描 8/16/32/128 | 每轮 quiet、总 latency、链路分布 | 中等 group 明显降低 max latency |
| P2 | WQE/quiet 固定成本过高 | 合并同 peer 连续 token | WQE/byte、submit/quiet cycles | WQE 数下降且有效带宽上升 |
| P2 | 3:1 QP 不适合当前 shape | 扫描 1QP、1:1、3:1、1:3、4QP | per-QP bytes/WQE/quiet | 某策略稳定降低最后完成 QP 的时间 |
| P3 | [S,P] route scan 在 128P 过重 |
compact destination list 原型 | route cycles、GM bytes、plan time | P 扩展时 route 时间不再近似线性增长 |
| P4 | 全局同步吞掉重叠 | per-group completion/credit 原型 | barrier cycles、并发在途 WQE | barrier 降低且 correctness 不变 |
| P4 | AIV 分工失衡 | 扫描 token/weight AIV 比例 | 每组 AIV 完成周期 | 两组结束时间接近,总时间下降 |
| P5 | GM/UB 路径受限 | 重排连续化与 ping-pong | MTE cycles、GM throughput | 大 H/大 batch 下收益更明显 |
每次只改一个主要变量。若同时改拓扑、QP 数、路由格式和 AIV 分工,即使性能上涨,也无法知道收益来自哪里,更无法判断换一个 shape 后为什么回退。
回到“只有理论一半”
现在可以把最初的问题改写得更精确:
在 128P 两层拓扑上,当前 dispatch 是受 hidden payload 带宽限制,还是受全 peer 扇出、WQE/quiet、
[S,P]规划和最慢 rank 同步限制?
源码已经能确认几件事:ascend_deepep 当前使用扁平 peer 扫描、source-rank 旋转、固定两个 QP 和 3:1 channel 分配;Host 侧始终按 S×P 预留内部接收区;num_worst_tokens 能避免 Host 回读但不会缩小这块 workspace;海思 hierarchy 展示了“按 server 跨机聚合、机内再展开”的协议;MoonEP 展示了“16-rank group、拓扑双射波次、token/weight AIV 分工和局部 completion/credit”的另一条路线。
源码还不能证明的是:3:1 为何在当前硬件上最优、MoonEP 的 2×64 拓扑是否等于实际部署、hierarchy 的 relay 收益能否覆盖当前 128P,以及理论带宽的一半究竟丢在哪个阶段。这些问题不能继续靠静态阅读裁决,只能靠分阶段 cycle、per-QP WQE/bytes、route histogram 和 max-rank latency 去收束。
因此第一版优化不必追求“大改架构”。先把计时面补齐,再做两个最小实验:QP 策略扫描和peer group width 扫描。如果前者无效、后者在 128P 明显有效,主矛盾就从“单 QP 不够并行”转向“拓扑与控制面扇出”;届时再投入 compact plan 或 hierarchy,路径会清楚得多。
源码锚点
本文只描述以下固定 revision 和会话中已经走读的实现,不把实验性策略升级为通用硬件结论:
ascend_deepep@18e2dd4b:PythonBuffer.dispatch、Hostcsrc/ops/dispatch.cpp、Devicekernels/dispatch.cpp与kernels/dispatch_udma.cpp。ops-transformer@8dc2fa04:README 中的fullmesh/hierarchy约束与moe_distribute_dispatch_v2_layered.h。ascend-moonep@b67c6840:moonep_dispatch_udma.cpp中的 64P/128P group schedule、QP 策略、AIV 分工和 profiling。
[^moonep-64p-group]: MoonEP 官方仓库说明动态冗余 expert、固定 S×K 工作量与 zero-copy 契约;Ascend 分组实现以 moonep_dispatch_udma.cpp@b67c6840 为准:GetDispatchGroupWidth()、GetGroupRankIndex()、GetGroupDstRank() 和 GetGroupIncomingPeer() 定义 64P legacy 映射,DispatchScatterToken() 与 CompleteTokenScatterPeerWithSignal() 定义四轮 peer 遍历及 completion/credit/strict-group-sync 生命周期。本文另用独立 Python 枚举了全部 64×4×16 映射,验证每个 source rank 四轮覆盖全部 destination、每轮 group 映射为置换、paired groups 展开为完整 (K_{16,16}),以及 incoming 公式为 sender 的精确逆映射;这属于对源码公式的可复现推导,不是硬件性能测量。
Ascend DeepEP MoE Dispatch Optimization