mini-infer系统实战-22-Packed Communication：难点不是 bytes 公式，而是实现和 benchmark 口径一起做对

很多人第一次看 MoE 推理里的 Expert Parallelism，会把注意力全部放在一个问题上：

padded all_to_all_single 的通信量是 ideal 的 2 倍，那把它改成 packed / non-padded 不就完了吗？

如果只是写一页设计文档，这句话没有问题。
但真正把它做成一轮可以落盘、可以复盘、可以拿去面试讲清楚的工程实现时，你会发现这件事远比“把通信 bytes 减半”复杂。

在 mini-infer 的 Phase 19，我做的是一版 synthetic MoE 上的 packed / non-padded expert dispatch。最终正式 benchmark 结果是：

  • dense：22552.86 tok/s
  • ep_padded：43046.93 tok/s
  • ep_packed：52400.26 tok/s
  • ep_packed / dense = 2.323x
  • ep_packed / ep_padded = 1.217x
  • ep_packed_bytes_per_layer = ep_ideal_bytes_per_layer = 262144
  • max_abs_diff_packed = 0.000000

这些结果来自该阶段的正式 benchmark，不是估算值。

这篇文章想讲的不是“如何解释 packed communication 的概念”，而是更实际的问题：

为什么 Phase 19 真正难的不是 bytes 公式，而是要同时把实现路径、控制面开销和 benchmark 口径一起收住。

背景：Phase 18 解决了权重所有权，没有解决通信本身

在 Phase 17，我已经把 synthetic MoELayer、EPMoELayer、EPEngine 和 2 卡 EP 主链路跑通了。
在 Phase 18，我又把 true expert sharding 收住了，让每个 rank 只持有本地 expert shard，正式 benchmark 得到：

  • shard_ratio = 0.5002
  • EP / dense = 1.916x
  • max_abs_diff = 0.000000

但当时还有一个很明确的技术债：

  • ep_ideal_bytes_per_layer = 262144
  • ep_padded_bytes_per_layer = 524288

也就是说，权重所有权已经收住了，通信却还是 prototype。

这正是我把 Phase 19 定成 Non-Padded Expert Dispatch / EP 通信闭环 的原因。目标非常明确：

  1. 保留 true expert sharding
  2. 增加 ep_packed
  3. 在同一套 benchmark 里对比 dense / ep_padded / ep_packed
  4. 把 exact bytes、数值等价和吞吐一起收住

问题定义：想要 exact bytes，首先要搞清楚你到底在测什么

当前通信量的三套口径是：

  • tp_bytes_per_layer = 131072
  • ep_ideal_bytes_per_layer = 262144
  • ep_padded_bytes_per_layer = 524288
  • ep_packed_bytes_per_layer = 262144

这里最关键的不是公式本身，而是你要非常明确哪些东西被统计进 bytes，哪些没有。

在我的 benchmark 里，Phase 19 统计的是：

  • 只算 hidden-state payload
  • 不算 expert-id metadata
  • 不算 valid mask metadata
  • packed path 的 split-size control plane 单独写 note，不偷偷算进 payload bytes

这件事在 benchmark_moe.py 里是显式表达的，而不是“默认大家都懂”。

为什么这很重要？
因为如果不把 payload bytes、metadata bytes、control-plane 开销拆开，你很容易得到一种看起来很漂亮、实际上无法比较的结论：

  • 通信 bytes 看起来已经等于 ideal
  • 但吞吐可能完全没有变好
  • 甚至 benchmark 还可能因为口径问题把 packed 路径测快了

这也是 Phase 19 里我真正踩到的第一个坑。

方案设计：保留 padded baseline，再并行增加 packed path

我没有直接把旧通信路径删掉，而是让 EPMoELayer 同时保留两条 distributed 路径：

  • comm_mode="padded"
  • comm_mode="packed"

对应代码在 moe_layer.py。

这背后的设计取舍很直接：

  1. padded 是稳定 baseline
    它的通信量不是最优，但实现简单，行为稳定，适合作为 benchmark 的保底对照。

  2. packed 是真正想验证的新路径
    它追求 exact split-size 的 hidden-state dispatch / return path，但天然会引入更复杂的 split-size 控制面。

  3. 两条路径必须在同一个 compare benchmark 里对照
    否则你没法判断 packed 到底是“真的更好”，还是只是换了 workload、换了权重、换了输入。

因此 Phase 19 的 benchmark 入口最终被扩成了三路 compare：

  • dense
  • ep_padded
  • ep_packed

并且强制使用：

  • 同一组 synthetic MoE 权重
  • 同一组 hidden states
  • 同一套官方 workload

实现细节：packed 路径到底改了什么

1. EPMoELayer 新增 comm_mode

在 moe_layer.py 里，EPMoELayer 现在根据 comm_mode 走两条不同的 distributed forward：

  • _forward_distributed_padded()
  • _forward_distributed_packed()

padded 的思路是固定 chunk：

  • 每个 rank 都按 num_tokens * top_k 预留固定长度
  • 通过 valid mask 过滤空位
  • 回传时也是同样的 padded 逻辑

packed 的思路则是 exact split-size：

  • source rank 先做 router 和 dispatch
  • 拿到 send_counts
  • 用 all_to_all_single(..., input_split_sizes=..., output_split_sizes=...)
    做真实 split-size 通信
  • 回传路径也按真实 split-size 走

这一点在 moe_layer.py 的 prepare_packed_source_context() 和 _forward_distributed_packed() 可以直接看到。

2. PackedSourceContext 的作用不是“缓存加速”，而是职责拆分

我新增了 PackedSourceContext，其中包含：

  • route
  • layout
  • dispatched_x
  • leading_shape

它的作用不是为了偷偷把更多工作挪出计时窗口，而是为了把 source rank 的一次性路由和 layout 结果明确表达出来，避免 packed path 的数据流散在多个函数里，不利于 review 和 benchmark 对齐。

这一步非常像真实工程里常见的重构动作：
不是为了“代码更优雅”，而是为了让实现职责和 benchmark 口径能对应起来。

3. EPEngine 不只是透传 comm_mode

在 ep_engine.py 里，EPEngine 做了两件重要的事：

  1. worker 能按 comm_mode 选择通信路径
  2. benchmark_forward() 的 steady-state 计时必须和 benchmark 脚本里的吞吐口径一致

很多 EP 原型的问题不是 forward 数学错，而是 benchmark 口径和 worker 生命周期根本没对齐。
Phase 19 的一个重点，就是不让 packed 变成“功能对、benchmark 假快”的路径。

两个真实坑：都是 benchmark 口径问题，不是 MoE 数学问题

这轮最有价值的经验，不是 exact split-size 的 API 用法，而是我连续踩到了两次 benchmark 口径坑。

坑 1：把 .cpu().tolist() 放进了 EPMoELayer 热路径

第一版 packed 实现里，我把 GPU 上的 send_counts 在 EPMoELayer 的 per-layer distributed forward 中拉回 CPU：

1
send_counts_cpu = send_counts.cpu().tolist()

这在功能上没有错，但 review 立即把它拦下了。原因很明确：

  • 这是 per-layer 路径
  • 它会引入 host sync
  • 层数一多，这种东西会直接污染 steady-state 吞吐

这个问题的关键不在于“能不能 work”，而在于这种实现根本不能拿去做性能结论。

坑 2：把 source-rank router / dispatch 预计算挪到了计时窗口外

第二版修复时，我为了把 host sync 从 EPMoELayer 热路径里移出去，把 packed path 的 source-rank 预处理放到了 worker 循环外。

结果呢？

  • ep_packed 看起来突然很快
  • 甚至远快于 ep_padded

这次不是实现错，而是 benchmark 口径错了。因为我等于在测：

  • 不含 source-rank router
  • 不含 dispatch layout
  • 不含 token packing

的 packed 路径吞吐。

这件事在 review 里再次被拦下。最终我只能把 packed path 的这些工作重新放回每次 timed iteration 里，同时又保持：

  • host sync 不回到 EPMoELayer per-layer hot path
  • split-size control-plane 仍然在 worker 侧处理

这也是为什么我最后在 benchmark note 里明确写出：

ep_packed_note=... includes source-rank router/dispatch + split-size control plane

这是一个非常典型、也非常面试友好的工程经验：

很多时候你不是“把代码写对了”就完了，你还得把 benchmark 的边界画对。否则你测到的根本不是你以为的那个系统。

实验结果：Phase 19 真正完成了什么

正式 workload 是：

  • batch_size=4
  • seq_len=16
  • hidden_size=512
  • intermediate_size=1024
  • num_experts=8
  • top_k=2
  • dtype=float16
  • warmup=2
  • runs=5
  • src_rank=1

同配置跑了 2 次，平均结果是：

指标 结果
Dense throughput 22552.86 tok/s
EP padded throughput 43046.93 tok/s
EP packed throughput 52400.26 tok/s
EP packed / dense 2.323x
EP packed / EP padded 1.217x
max_abs_diff_padded 0.000000
max_abs_diff_packed 0.000000
shard_ratio 0.5002
ep_ideal_bytes_per_layer 262144
ep_packed_bytes_per_layer 262144

两次波动分别是：

  • dense：0.69%
  • ep_padded：1.61%
  • ep_packed：3.29%

都远低于 20% 的异常阈值，所以这组结果是稳定的。

从结论上看，Phase 19 完成了三件事：

  1. packed 的 payload bytes 已经和 ideal 对齐
  2. packed 的吞吐没有因为 control-plane 开销塌掉
  3. packed 的数值和 dense oracle 仍完全一致

这轮真正的收获：工程上要同时收三条线

如果只从“功能实现”角度看，Phase 19 好像是在做一件很直白的事：

  • 把 padded dispatch 改成 packed dispatch

但从工程上讲，这轮真正要同时收住的是三条线：

  1. 实现线
    exact split-size dispatch / return path 要真的能跑

  2. 口径线
    bytes、metadata、control-plane 和 steady-state timing 的边界要讲清楚

  3. review 线
    per-layer host sync 和“计时窗口外偷算工作”这两类问题都必须挡住

这也是为什么我现在会认为，Phase 19 比单纯把 all_to_all_single 改成 split-size 更有价值。
因为这轮真正练到的是：

  • 怎么把一个通信优化做成可解释的 benchmark
  • 怎么识别“实现没错但测法错了”的情况
  • 怎么让 review 真正成为性能结论的防线

还有什么没做

Phase 19 完成了 communication 闭环，但它并不是最终形态。

当前明确没做的东西有：

  • ep_size > 2
  • grouped GEMM
  • dispatch / communication overlap
  • 端到端 serving 集成
  • 消除 host-side split-size Python list

所以现在最准确的表述不是“MoE EP 已经产品级”，而是：

mini-infer 已经完成一版 synthetic MoE 上的 non-padded expert dispatch 闭环，能把 exact bytes、真实吞吐和数值等价同时说清楚。

这对学习项目和面试表达都已经很强，但和真正的生产推理系统之间，仍然隔着 runtime 与 serving 这一层。

总结

Phase 19 给我的最大结论是：

Packed communication 的难点从来不只是 bytes 公式，而是要把实现、控制面和 benchmark 口径一起做对。

如果你只是把 ep_padded_bytes_per_layer 变成 ep_packed_bytes_per_layer，你只能证明你会写一段 split-size all_to_all_single。
如果你还能把：

  • per-layer host sync
  • source-rank 预计算是否计入吞吐
  • payload bytes 和 metadata/control-plane 的边界
  • dense / ep_padded / ep_packed 的公平 compare

都一起讲清楚，那你讲的就不再是“一个能跑的 demo”，而是一段真正像 AI infra 推理工程的工作。

对想冲字节 AML、Doubao、推理系统核心组的人来说，这种差别非常关键。

系列导航