我面对的是一个很具体、也很容易被“理论带宽”带偏的问题:Ascend DeepEP 的 MoE dispatch 吞吐只有硬件理论值的一半,128P 跨机场景尤其明显,而且系统里还有两层路由。麻烦在于,我既不熟悉 MoE dispatch 的接口与数据协议,也没写过典型 URMA/UDMA 算子,更不知道应该在哪里计时、打印和判断瓶颈。
这篇文章不从零散优化技巧出发,而是沿一条 token 的真实旅程,依次读懂 Python 接口、Host 侧 workspace/notify、Device 侧 AIV/UDMA/QP、目标 rank 的接收布局与反向 combine;再把 ascend_deepep、海思 hierarchy 和 MoonEP 的 peer 调度放进同一张拓扑图。最终目标不是立即猜出一个“神奇参数”,而是建立一条可证伪的优化路线:先区分数据面、控制面和同步面,再判断 128P 下真正撞到的是字节带宽、WQE 提交率、P 维路由扫描、拓扑扇出还是最慢 rank 的尾延迟。
问题不只是“带宽没跑满” 看到吞吐只有硬件理论值的一半,第一反应通常是增加 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 2 3 x[token] 一行 hidden topk_idx[token, 0:K] 目标专家 topk_weights[token, 0:K] combine 权重
一个全局专家 expert_id 对应:
1 2 dst_rank = expert_id / experts_per_rank dst_local_expert = 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,紫色是完成和缓冲区复用协议。
{ width=100% }
图 1:根据会话内容与固定 revision 源码整理的自绘示意图。重点不是组件数量,而是 payload、路由 offset 和完成信号不能互相替代。
一个正确的单边通信顺序通常是:
1 2 3 4 5 计算目标地址和接收 offset → 把 payload Put 到目标 SHMEM window → 确认 payload 已经对端可见 → 再发布 ready/count/completion → 接收方消费并复用缓冲区
如果 flag 先于 payload 可见,接收方就可能读到旧数据;如果缓冲区在 quiet 前被复用,后续 WQE 可能读到被覆盖的源数据。因此,通信和计算能否重叠,首先是依赖图与内存可见性问题,其次才是调度技巧。
Rank :EP 通信域中的一个参与者,通常对应一张 NPU 卡上的进程或设备执行上下文。
AIV :Ascend AI Core 的 Vector 侧执行单元。这里不只做数值计算,也负责扫描路由、搬运 GM/UB、构造和提交通信任务。
GM / UB :GM 是设备全局内存,容量大、延迟高;UB 是核内高速缓冲,容量小,常用来做 staging、向量处理和 WQE 暂存。
SHMEM :为每个 rank 建立对称内存窗口,使相同 offset 可以映射到不同 peer 的远端地址,适合 one-sided Put/Get。
URMA/UDMA :URMA 更接近远端内存访问语义和资源体系;UDMA 是这里执行远端搬运的具体 DMA 路径。不要把 API 语义、Runtime 资源和硬件引擎压成同一个概念。
QP(Queue Pair) :远端通信队列。发送侧把工作提交到 QP,底层引擎异步执行。
WQE(Work Queue Element) :一次通信工作的描述符,包含源/目标地址、长度和操作类型等信息。小包很多时,瓶颈可能是每秒能处理多少 WQE,而不是每秒能搬多少字节。
CQE / completion :某批通信工作的完成记录或完成信号。
**quiet()**:等待此前提交到指定范围的非阻塞通信完成,并建立后续复用或发布 flag 所需的顺序边界。它不是普通函数调用开销,而可能直接暴露网络和队列尾延迟。
Barrier :所有参与者到齐后再推进。它很容易保证正确性,也很容易把一个慢 rank 的延迟放大到全局。
从 Python 接口进入真实调用链 本节以 ascend_deepep@18e2dd4b 为代码锚点。Python 入口位于 ascend_deepep/buffer.py::Buffer.dispatch,它先区分是否提供 handle。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 if handle is not None : ( rank_prefix_matrix, channel_prefix_matrix, _recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank_from_handle, _send_head, ) = _unpack_v1_handle(handle) return self ._native.intranode_dispatch( x_tensor, x_scales, None , None , None , is_token_in_rank_from_handle, None , num_recv_tokens, rank_prefix_matrix, channel_prefix_matrix, ..., )
没有 handle 时,需要传入当前路由的统计和矩阵:
1 2 3 4 5 6 if num_tokens_per_rank is None : raise ValueError("num_tokens_per_rank is required" ) if is_token_in_rank is None : raise ValueError("is_token_in_rank is required" ) if num_tokens_per_expert is None : raise ValueError("num_tokens_per_expert is required" )
随后执行 native dispatch,并把这次生成的路由计划打包成 v1_handle:
1 2 3 4 5 6 7 8 v1_handle = ( rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank, send_head, )
这段接口已经透露了一个重要优化边界:hidden 可以每轮变化,但如果路由关系不变,prefix、offset 和反向索引没有必要每轮重建。 cached dispatch 就是在复用这部分控制面工作。
Fresh dispatch 不是正式 API 名称,只是为了区分两条路径使用的简称:
1 2 3 4 5 6 7 8 9 handle is None → 当前路由第一次出现或发生变化 → 重新计算 prefix、offset、src_idx、send_head → 返回新的 handle handle is not None → 路由关系与上次一致 → 复用 handle,只搬运新的 x → cached dispatch
因而“Fresh”更准确的中文是重新规划路由的 dispatch ,不代表第一次调用整个程序,也不代表输入 hidden 必须是新 Tensor。
handle 不是一个不透明的通信句柄,而是一组让 cached dispatch 和 combine 找回数据位置的 Tensor:
项
直观作用
后续消费者
rank_prefix_matrix
各 source rank 向各 destination rank 发送后的累计边界,决定一个源 rank 的数据落在目标接收区哪一段
cached dispatch、combine
channel_prefix_matrix
当前源 rank 内,各 channel 面向每个目标 rank 的累计发送边界
cached dispatch
recv_channel_prefix_matrix
目标 rank 收集到的“各 source rank × channel”前缀,用于还原本地接收段和规划反向发送
combine
recv_src_idx
每一行 compact 接收 token 对应源 rank 上的原始 token 下标
combine 把 expert 结果送回原 token
is_token_in_rank
shape 为 [S,P] 的布尔路由矩阵,表示某 token 是否需要发到某 rank
cached dispatch
send_head
shape 为 [S,P] 的发送序号/存在性表;-1 表示该 token 没有发给这个 contributor
combine 判断需要归约哪些 rank,并找到对应槽位
可以把它压缩成一句话:prefix 决定“段在哪里”,recv_src_idx 决定“这一行原来是谁”,send_head 决定“combine 应该等谁”。
Host 侧先造一张通信地图 进入 csrc/ops/dispatch.cpp 后,Host 侧主要完成四件事:校验输入、规划 symmetric workspace、启动 notify、创建输出 view 并启动真正的 dispatch kernel。
Workspace 不是一个 Tensor,而是一块地址空间 代码首先解析 AIV 数量:
1 2 const int64_t resolved_num_aiv = ResolveNumAiv (num_aiv);const int64_t num_channels = resolved_num_aiv;
随后按极端情况预留内部接收容量:
1 2 const int64_t recv_buffer_tokens = CheckedMulInt64 ( num_tokens, world_size, "dispatch receive capacity" );
原因是:每个 source rank 最多有 num_tokens 行可以发给当前 rank;共有 world_size 个 source rank,所以当前实现给内部接收区保留 num_tokens × world_size 行。
MakeDispatchWorkspaceLayout() 再用一个不断前移的 cursor,把一整块 symmetric buffer 切成多个互不重叠的区域:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 layout.recv_channel_prefix_offset = cursor; cursor = AppendWorkspaceRegion (cursor, ...); layout.x_buffer_offset = cursor; cursor = AppendWorkspaceRegion (cursor, recv_buffer_tokens * hidden_bytes, ...); layout.src_idx_buffer_offset = cursor; cursor = AppendWorkspaceRegion (cursor, recv_buffer_tokens * sizeof (int32_t ), ...); layout.topk_idx_buffer_offset = cursor; ... layout.topk_weights_buffer_offset = cursor; ... layout.x_scales_buffer_offset = cursor;
这些 offset 随后写入 tiling,Device kernel 只要拿到 symmetric buffer 基地址,就能定位 payload、反向索引、Top-K 元数据、scale 和 completion signal。
它不是 NCHW、ND、TND 之类的 Tensor 维度布局,而是一块大 workspace 的内存地图 :
1 2 3 4 5 6 7 8 9 symmetric_buffer base ├── notify workspace ├── recv_channel_prefix [P, num_channels] ├── x payload [S×P, H] ├── src_idx [S×P] ├── topk_idx [S×P, K],可选 ├── topk_weights [S×P, K],可选 ├── x_scales [S×P, scale_cols],可选 └── UDMA completion signals [P, 2, 512B]
offset 就是每个区域相对 symmetric_buffer base 的字节偏移。这样做可以复用一次大的对称内存分配,并保证远端 rank 使用相同 offset 找到对应区域。
Notify 先交换数量和前缀 Host 在 payload dispatch 前启动 notify kernel:
1 2 3 4 5 6 7 8 9 10 11 LaunchNotifyDispatchKernel ( *num_tokens_per_rank, *num_tokens_per_expert, is_token_in_rank, rank_prefix_matrix, channel_prefix_matrix, notify_num_recv_tokens, notify_num_recv_tokens_per_expert, tiling, symmetric_buffer, shmem.FftsAddr ());
notify 的核心任务不是提醒一句“开始接收”,而是把不规则路由转成通信双方都能使用的边界:当前 rank 实际要收多少 token、每个本地 expert 收多少、每个 source rank 和 channel 应写到接收区的什么位置。真正的 payload kernel 依靠这些 prefix,才能让多个 AIV 和多个 rank 并行写入而不互相覆盖。
可以这样理解,但要再精确一步:它不是发送一张 token 明细表,而是在生产计数和 offset 。
1 2 3 4 5 Router 结果 → 每个 token 要去哪些 rank → 每个 source→destination 有多少行 → 每个 channel 负责其中哪一段 → prefix sum 得到接收起点
有了起点后,数据面可以直接执行:
1 2 3 4 remote_address = remote_shmem_base + x_buffer_offset + recv_position * row_stride
所以 notify 更像 dispatch 的控制面规划与发布阶段 。即使调用方提供了固定 num_worst_tokens,notify 仍然要运行;省掉的只是 Host 回读实际接收行数,而不是 prefix 计算本身。
输出长度有动态和固定两种模式 notify 之后,Host 决定公开输出 Tensor 的第一维:
1 2 3 4 5 6 7 8 if (cached_mode) { num_recv_tokens = cached_num_recv_tokens; } else if (num_worst_tokens > 0 ) { num_recv_tokens = num_worst_tokens; } else { auto recv_tokens_cpu = notify_num_recv_tokens.cpu (); num_recv_tokens = recv_tokens_cpu.item <int32_t >(); }
然后用 SHMEM 内部地址创建非 owning view:
1 2 3 4 5 auto recv_x = MakeShmemTensor ( symmetric_buffer, dispatch_workspace.x_buffer_offset, {num_recv_tokens, hidden}, x.scalar_type ());
最后才启动 LaunchV1DispatchKernel()。因此 Host 路径可以概括成:
1 2 3 4 5 6 7 validate → resolve AIV/channel → reserve worst internal workspace → launch notify → choose public output rows → create SHMEM Tensor views → launch dispatch payload kernel
直觉层。 把输出 Tensor 想成提前摆好的货架:
1 2 3 4 5 6 7 8 9 num_worst_tokens = 0 → 等 notify 统计实际收到 237 行 → Host 回读 237 → 输出 shape = [237, H] num_worst_tokens = 256 → Host 不回读实际计数 → 直接输出 shape = [256, H] → 前 237 行有效,后 19 行 padding
源码层。 Device 端执行:
1 2 3 const int32_t actual_recv_tokens = ...;const int32_t recv_tokens_to_copy = DispatchMin (actual_recv_tokens, num_output_tokens);
所以调用方必须保证:
[ actual_recv_tokens \le num_worst_tokens ]
上界太大会增加 padding;上界小于真实值时,当前实现会只公开前 num_worst_tokens 行,存在静默截断风险。尾部 x/scales/weights 填 0,topk_idx 塄为 -1。
还要区分两种容量:内部 SHMEM 仍按 S × P 预留,num_worst_tokens 只决定公开输出长度和是否需要 Device→Host 回读。它当前不会缩小内部 workspace 。
最理想的情况是 layout 阶段已经知道精确接收量,直接把精确值传进来:既没有 Host 同步,也没有 padding。只知道上界时,可以采用安全 bucket;完全不知道时先传 0 保证正确性。
Device 侧如何把 channel 变成远端 Put 普通 dispatch.cpp 的核心所有权模型是:一个 AIV 负责一个 channel,并维护面向所有 destination rank 的游标。
1 2 3 4 5 6 7 8 9 10 for (int32_t dst_rank = 0 ; dst_rank < world_size; ++dst_rank) { const int32_t rank_base = rank > 0 ? rank_prefix_gm.GetValue ((rank - 1 ) * world_size + dst_rank) : 0 ; const int32_t channel_prefix_begin = core_idx > 0 ? channel_prefix_gm.GetValue (dst_rank * num_channels + core_idx - 1 ) : 0 ; recv_position_ub.SetValue ( dst_rank, rank_base + channel_prefix_begin); }
随后它遍历自己负责的 token 范围和全部目标 rank:
1 2 3 4 5 6 7 8 9 for (int32_t token = token_begin; token < token_end; ++token) { for (int32_t dst_rank = 0 ; dst_rank < world_size; ++dst_rank) { if (is_token_in_rank_gm.GetValue ( token * world_size + dst_rank) == 0 ) { continue ; } ... } }
这段代码简单直接,但它也暴露了 128P 的一个结构性成本:即使 Top-K 只产生少数目标 rank,每个 token 仍可能扫描完整的 world_size 维度。
UDMA 版本在 peer 调度上又增加了两个策略:按 source rank 旋转目标顺序,以及用两个 QP 承载 channel。
1 2 3 4 5 6 7 8 9 for (int32_t logical_destination = destination_worker_index; logical_destination < world_size; logical_destination += destination_worker_count) { int32_t dst_rank = rank + logical_destination; if (dst_rank >= world_size) { dst_rank -= world_size; } ... }
如果所有 rank 都按 0,1,2,... 的顺序访问 peer,那么某一时刻可能出现:
1 2 3 4 5 rank 0 → peer 0 rank 1 → peer 0 rank 2 → peer 0 ... rank 127 → peer 0
旋转后则变成:
1 2 3 rank 0 → 0,1,2,... rank 1 → 1,2,3,...,0 rank 2 → 2,3,4,...,1
它的目的只是错开起始 peer,降低同一时刻的集中冲击 。它没有减少 peer 数,也没有理解 server、交换框或 rail;每个源 rank 最终仍会面对全部目标 rank。因此它是一种低成本去同步化调度,不等于拓扑感知算法。
两个 QP 采用 3:1,而不是均分 当前代码固定两个 QP,并按 channel 编号分配:
1 2 3 4 5 6 7 8 __aicore__ inline int32_t DispatchUdmaQpForChannel ( int32_t channel, int32_t qp_count) { if (qp_count == 2 ) { return (channel & 3 ) == 3 ? 1 : 0 ; } return 0 ; }
也就是:
1 2 channel: 0 1 2 3 | 4 5 6 7 QP: 0 0 0 1 | 0 0 0 1
所有 peer 的 payload 提交完成后,每个 QP 分别追加一个 Strong Order completion,再执行 qp_quiet(dst_rank, qp_slot);最后才是:
1 2 aclshmem_quiet ();aclshmemx_barrier_all_vec ();
并行机会 和负载均分 是两个问题。目标是让两个 QP 尽量同时结束:
[ T = \max\left(\frac{W_0}{R_0}, \frac{W_1}{R_1}\right) ]
如果两个 QP 的有效速率相同,理想模型确实倾向 1:1;如果某次实验中 (R_0 \approx 3R_1),3:1 反而更接近同时完成。
但源码只明确说这是 MoonEP 的 two-QP performance experiment,没有记录底层原因。因此只能提出待验证假设:两个 QP 可能存在服务能力差异、共享资源竞争、CQE/queue depth 差异,或者 3:1 只对当时的 shape 和拓扑有效。不能把“QP0 就是三倍快”写成事实。
更重要的是,按 channel 计数 3:1 不等于按字节或 WQE 3:1。不同 channel 的 token 数可能高度不均。优化时应该记录每个 QP 的 WQE 数、payload 字节、提交周期和 quiet 周期,并至少比较 1QP、2QP 1:1、2QP 3:1、2QP 1:3 与 4QP round-robin。
completion signal 的地址计算是:
1 2 signal_slot = rank * kDispatchUdmaQpCount + qp_slot; signal_addr = base + signal_slot * kDispatchUdmaCompletionSignalBytes;
常量是:
1 2 kDispatchUdmaQpCount = 2 ; kDispatchUdmaCompletionSignalBytes = 512 ;
所以每个 symmetric workspace 需要为每个 peer、每个 QP 保留一个独立的 512B completion slot:
[ bytes = world_size \times 2 \times 512 ]
128P 时:
[ 128 \times 2 \times 512 = 131072\ bytes = 128\ KiB ]
512B 是实现选定的 signal slot 间距和对齐单位,不是 uint64_t completion 值本身有 512B;真正的值仍是 8 字节,较大步长用于隔离和满足实现的缓存行/顺序要求。
128P 为什么会放大问题 小规模下,“每个源 rank 直接面向每个目标 rank”简单且可能足够快;扩展到 128P 后,数据量未必线性增长,控制面却很容易增长:
1 2 3 4 5 路由表示: [S,P] token 路由扫描: O(S×P) peer 状态: O(P) rank prefix: O(P²) completion slot:O(P×QP)
与此同时,每个 peer 的消息变小,固定 WQE/quiet 成本占比上升;不同源 rank 若在相近时刻访问相同拓扑区域,还会形成热点。这里要解决的已经不是单核 DataCopy 的问题,而是全系统同时有哪些通信关系处于活跃状态 。
下面把三套代码的调度单位放到同一张图里。
{ width=100% }
图 2:根据会话内容与固定 revision 源码整理的自绘比较。三种实现的硬件与能力边界不同,图中只比较 peer 调度思想,不宣称可以逐行互相替换。
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 2 3 4 5 6 7 fullmesh: token 直接通过 RDMA 发往 Top-K 专家所在 rank hierarchy: token 经过跨机、机内两次发送 仅不同 server 同号卡之间使用 RDMA server 内使用 HCCS
Layered kernel 固定:
1 2 constexpr static uint32_t SERVER_RANK_SIZE = 16 ;serverNum_ = worldSize_ / SERVER_RANK_SIZE;
跨机目标中继 rank 的计算是:
1 2 3 uint32_t dstRankId = rankId_ % SERVER_RANK_SIZE + dstServerId * SERVER_RANK_SIZE;
也就是说,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。跨机阶段先发送到:
rank 37 是 server 2 的同号中继;随后 server 2 内部再把 hidden 送到 rank 42。若同一 token 在 server 2 上命中多个专家,跨机阶段还可以按 server 聚合 hidden,避免为同一 server 重复发送多份大 payload。
代价也很明确:多一次 relay、机内 copy 与控制协议。只有当减少的跨机 peer、重复 hidden 和 RDMA 下发成本,大于新增机内搬运成本时,hierarchy 才会获益。
固定 revision 的 README 声明:fullmesh 支持到 128P 及以上,而 hierarchy 当前只支持 16、32、64P。它仍然非常适合学习“先按 server 聚合,再机内展开”的设计,但不能据此声称这份实现已直接覆盖当前 128P 场景。要迁移思想,必须重新确认 server 数、每 server rank 数、HCCS/RDMA 路径和 tiling 约束。
MoonEP:16-rank group 与拓扑波次 ascend-moonep@b67c6840 为 128P 写出了明确的生产拓扑假设:
1 2 3 4 5 6 constexpr uint32_t kDispatchFrameCount = 2 ;constexpr uint32_t kDispatchRanksPerFrame = 64 ;constexpr uint32_t kDispatchRanksPerDirectRow = 8 ;constexpr uint32_t kDispatchGroupRows = 4 ;constexpr uint32_t kDispatchGroupColumns = 4 ;constexpr uint32_t kMaxDispatchGroupWidth = 16 ;
它把 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 没有按这套物理位置编号,照搬公式可能主动制造热点。
不是。首先,“16 rank”指每个 source rank 当前面对的 16 个目的 peer ,不是全系统只有 16 个 rank 活跃;128 个 source rank 都在同时执行自己的映射。
其次,kernel 固定启动 32 AIV,并明确分工:
1 2 3 4 const int64_t token_aiv_count = GetDispatchGroupWidth (world_size); const bool token_worker = aiv_index < token_aiv_count;const bool weight_worker = !token_worker;
128P 时:
1 2 3 4 5 6 7 8 9 10 11 AIV 0–15 → 16 个 token workers → 每个 AIV 负责当前 group 的一个 peer → 搬运大块 hidden AIV 16–31 → 16 个 weight workers → 并行处理 Top-K weight 元数据 汇合后 → 32 个 AIV 一起执行 DispatchEpilogue
如果 with_weights == 0,后 16 个 AIV 在 token 阶段可能没有有效工作;如果 weight 很小,它们也可能提前到达 barrier。于是 16:16 只是特定 workload 的静态划分,不保证所有 shape 都负载均衡。
三种策略可以用同一组问题比较,而不是只看函数名:
实现
同时面向的通信单位
跨机路径
控制协议
主要边界
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 2 is_token_in_rank[token, rank] = 当前 token 是否有专家位于该 rank
它的优点是判断简单,缺点是在 128P 下表示和扫描都按 (P) 扩展。假设 Top-K 为 8,而这些专家最多落在 8 个 rank 上,那么每个 token 实际只需要保存少量 destination,却要读取 128 个布尔值。
更紧凑的 plan 可以写成:
1 2 3 4 5 6 7 8 9 10 11 token 0 → [dst rank 3, dst rank 17] token 1 → [dst rank 8] token 2 → [dst rank 3, dst rank 9, dst rank 17] destination_offsets: rank 3 → [begin, end) rank 8 → [begin, end) ... 每个 destination 段内再保存: token_id、local_expert_id、topk_slot
hierarchy 模式则先压成 destination server,再在 server 内展开 rank/expert:
1 2 3 4 token → unique destination server list → server segment offset → destination rank / local expert list
这样既能减少 [S,P] 扫描,也能识别“同一 token 在同一 server 命中多个专家”,为跨机 hidden 去重创造条件。需要注意,plan 生成本身也有成本;如果把 planning 时间排除在 benchmark 外,就可能得到一个对端到端系统无效的优化结论。
原来的表示像一张 128 列考勤表,即使一个 token 只去两个 rank,也要逐列看完:
1 token 7: 0 0 0 1 0 0 ... 0 1 ... 0 共 P 列
compact plan 只保存非零项:
1 token 7: [rank 3, rank 65]
再按 destination 聚合,就得到真正适合通信提交的连续段:
1 2 rank 3 : token [7, 12, 19, ...] rank 65 : token [7, 21, ...]
优化点不是把 bool 换成更小的 dtype,而是把工作量从“查看全部可能目标”改成“只处理真实目标”。理想情况下复杂度更接近 S × unique_dst_per_token,而不是 S × P。
先把时间和字节测对 在改调度前,需要建立一个同时覆盖 Host、Device、网络和多 rank 尾部的测量面。最少应该回答:
1 2 3 4 5 本 rank 实际发送/接收了多少 hidden 字节? 生成并提交了多少 WQE? 每个 QP 分到多少 WQE 和字节? route scan、WQE submit、qp_quiet、final barrier、epilogue 各花多久? 最慢 rank 是谁,它的 route histogram 是否异常?
有效 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 2 3 4 5 event_start → layout/notify/dispatch enqueue event_end → 最后统一 synchronize → elapsed(event_start, event_end)
不要每个小阶段后立刻 Host synchronize,否则测到的是被人为串行化后的路径。正确性调试阶段可以同步定位,性能测量阶段则应连续提交多轮、丢弃 warmup,并明确是否冲刷 L2。
Device 侧计时 MoonEP 已经展示了一个适合借鉴的模式:用编译期 Profile 开关读取 cycle,把每个 AIV 的阶段周期和字节数写入独立 GM 槽位;Profile=false 时编译成零开销路径。
建议每个 AIV 至少记录:
1 2 3 4 5 6 7 8 9 route_scan_cycles wqe_build_cycles wqe_submit_cycles wqe_count payload_bytes qp0_quiet_cycles qp1_quiet_cycles barrier_cycles epilogue_cycles
记录位置必须按 [rank, aiv, metric] 或其他独占方式切开,不能为了计时再引入热点原子加。最后由 Host 一次性读取并聚合。
printf 适合确认控制流,不适合放在 token、peer 或 WQE 热循环中。调试时可以使用三道阀门:
1 2 3 4 5 if constexpr (Debug) { if (rank == debug_rank && aiv_index == debug_aiv) { } }
推荐顺序是:
先把计数、cycle 和关键状态写到每 AIV 独占 GM buffer。
kernel 结束后由 Host 统一读取并格式化。
真正需要确认死锁位置时,再限制到一个 rank、一个 AIV、一个 peer 打印。
高频 Device printf 会改变指令调度、同步和队列时序,甚至让原本不复现的竞态消失,因此不能把带打印版本的吞吐当作性能证据。
把优化变成可证伪的实验 优化顺序不应是“哪个技巧听起来高级就先做哪个”,而应是每一步只验证一个瓶颈假设。
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 2 3 4 5 group width → 活跃 peer 数 → 每轮 WQE 数和字节 → 每轮 quiet/barrier → max-rank latency
如果硬件拓扑确认与 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 2 3 4 5 1 QP 2 QP 1:1 2 QP 3:1 2 QP 按累计字节加权 4 QP round-robin
若增加 QP 只增加了 WQE/CQE 管理而没有降低 quiet 周期,瓶颈可能在共享链路或共享引擎,不在单 QP 队列深度。
P3:压缩路由计划 从 Top-K expert 直接生成 unique destination rank/server list,并按 destination compact;尽量融合:
1 2 3 4 expert → rank/server 映射 + count + prefix + token reorder plan
目标是避免多次读取 [S,P],也避免所有 rank 重复读取完整 source-rank 计数。计划生成要纳入端到端计时,除非业务确实能长期复用 cached handle。
P4:再做细粒度通算掩盖 把工作拆成真正无依赖或仅有局部依赖的阶段:
1 2 3 4 5 6 7 8 当前 peer group 的 UDMA 在途 ↔ 下一 group 的 route/offset/WQE 准备 hidden token AIV ↔ weight metadata AIV ping buffer 的 GM→UB ↔ pong buffer 的 UB→remote GM
全局 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。
你的理解方向基本正确:根据计算逻辑建立依赖图,把计算、通信、访存、原子和同步拆开,再根据 shape 精细调度,使多个硬件资源尽量同时工作。
但要补上两道门:
先证明独立性。 单边通信中的 payload、flag、credit 和 buffer reuse 有严格顺序,不能因为代码段看起来分开就并行。
先证明当前上限。 如果瓶颈是 WQE 提交率,增加 payload AIV 没用;如果瓶颈是最慢 expert,优化平均 GM 带宽也没用。
常见“几板斧”可以归纳为:
减数据 :量化、去重、同 server 命中多个专家时 hidden 只跨机一次、避免重复元数据。
减任务 :连续 token 合并成大 WQE,批量下发,降低 CQE、doorbell 和 quiet 次数。
懂拓扑 :按 server/frame/rail 分组,限制每轮 fan-out,用双射或错峰 schedule 避免热点。
做流水 :双缓冲,当前通信与下一批 planning/pack 重叠,hidden 与 weight 分工。
减同步 :用局部 completion/credit 替代不必要的全局 barrier,但不破坏内存可见性。
压访存 :连续布局、对齐、UB tiling、减少 GM 往返、避免伪共享和热点原子。
做均衡 :按实际字节/WQE 给 AIV 与 QP 分工,关注 expert/rank skew 和 max-rank tail。
按 shape 特化 :为常见 P/H/K/S 选择 group width、AIV 比例、QP 数和 bucket,避免一套参数覆盖所有场景。
真正可复用的方法不是背下这些技巧,而是把每一项写成:
1 假设 → 改动 → 指标 → 预期现象 → 失败时说明什么
一份可以直接执行的实验表
优先级
假设
最小改动
必看指标
支持假设的现象
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 :Python Buffer.dispatch、Host csrc/ops/dispatch.cpp、Device kernels/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 中的 128P group schedule、QP 策略、AIV 分工和 profiling。
AIV MoE All-to-All :dispatch/combine 的 one-sided Put、ready/count 与反向归约背景。
DMA Communication Terminology :UDMA、IPC/MTE、SHMEM 与完成语义的概念边界。