【轻量化复刻大模型 / KIMI K3】MoonEP 源码深度解读

【轻量化复刻大模型 / KIMI K3】MoonEP 源码深度解读

目录

0x00 概要

0.1 主要内容

定位：MoonEP 是 Moonshot AI（Kimi 团队）开源的 Expert Parallelism 通信库，通过动态冗余专家实现完美负载均衡。本文基于 MoonEP-master/ 官方源码，逐模块拆解其设计与实现。

原始来源：MoonEP 没有正式的 arXiv 论文，以 GitHub 开源仓库形式发布：

  • 标题：MoonEP: A Perfectly Balanced Expert Parallelism Library via Dynamic Redundant Experts
  • 作者：Yutian Chen, Cong Li, Yucheng Wang, Ming Wei
  • 机构：Moonshot AI（Kimi）
  • 年份：2026
  • 仓库：https://github.com/MoonshotAI/MoonEP

致谢与灵感来源：README 第 196-201 行列出灵感来源：

  • DeepEP（DeepSeek）：EP 通信库先驱
  • Echo（arXiv:2603.07685）：冗余专家调度理论
  • UltraEP（Dots-Infra）：EP 优化实现
  • AcclEP（阿里巴巴）：EP 通信库

符号约定：S = 每 rank 输入 token 数，K = 路由 top-k，E = EP 组内总路由专家数，R = EP rank 数，B = 每 rank 权重预取槽数，NvS = 每 rank 派发 token 槽位数（S*K real tokens + padding），H = 隐藏维度，H' = 专家 FFN 中间维度。

0.2 注意说明

  • 项目https://github.com/RossiLH/Model-Viz 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
  • 涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
  • 全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
  • 模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
  • 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
  • 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本（plot_*.py），可生成静态图。
  • 所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
  • 为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。

0x01 演进脉络与定位

1.1 前身：标准 EP 的三重痛点

标准专家并行（如 Megatron-LM 的 EP）的做法是：每个 rank 持有 E/R 个专家，路由后直接把 token 发到对应 rank。这带来三个致命问题：

痛点 具体表现 根因
负载不均 Hot rank 的 token 数远超平均，拖慢整次同步 Top-K 路由的随机性
显存碎片 每步每层的激活 shape 不同，无法预分配固定缓冲 路由结果动态变化
Host 同步停滞 每层都要 host-device 同步才知道该分配多大显存 形状未知

1.2 同期竞品

方案 核心机制 优势 劣势
DeepEP v2 固定 EP + 动态 token 分配 简单，无冗余开销 负载不均时性能退化，显存碎片
Echo 冗余专家理论证明 理论基础扎实 无工程实现
UltraEP EP 通信优化 低延迟 仍受负载不均影响
MoonEP 动态冗余专家 + GPU planning 完美均衡，零碎片 需要 NVLink 共享内存

1.3 定位一句话

不同于 DeepEP v2 的"固定分配 + 动态适应"做法，MoonEP 通过在线规划动态冗余专家实现了完美负载均衡——每个 rank 始终恰好计算 S×K 个 token，彻底消除 hot rank、显存碎片和 host 同步。

1.4 MoonEP vs DeepEP

我们首先比对下这两个方案。

  • DeepEP：固定专家计算位置，优化数据搬运路径
  • MoonEP：先调整专家计算位置，再按照新调度方案驱动数据流转
  • MoonEP：本质是内置调度器的完整 MoE 执行运行时
  • DeepEP：本质是面向对称显存、多 GPU 构建的 MoE 专用并行通信底座
  • MoonEP：会先做一层调度判断：倘若专家归属关系（原始专家归属）造成部分 rank 负载过热，不必强制把全部 token 送往原归属 rank。它会将本轮部分专家的计算任务临时迁移至空闲 rank，再基于新的执行位置编排后续通信。
  • DeepEP：EP 直接接收路由器输出的专家归属关系。专家驻留在哪一张 GPU rank，对应的 token 就必须发送至该 rank；库的核心目标是以最小化 SM 占用、HBM 拷贝、网络往返开销，将 all‑to‑all 通信压榨至硬件带宽极限。
  • MoonEP plan：改写输出目标 rank，再为 GEMM 计算分配最终执行槽位。
  • DeepEP layout：针对已经确定的目标 rank，分配通信槽位。
  • MoonEP：允许专家参数的计算位置临时漂移。前向传播需要预取 gate/up/down 权重；反向传播中，冗余槽位产出的 fp32 梯度不能直接交付优化器，必须通过reduce_grad归并回原始主参数。对接上层框架时，需要同步修改权重 VMM、共享预取池、临时梯度 buffer、优化器执行时序等多处逻辑。
  • DeepEP：遵循「token 去往参数所在位置」范式，原生适配标准 autograd 链路。dispatch 的反向为 combine，combine 的反向复用缓存 dispatch；专家参数始终仅在归属 home rank 完成计算，参数位置不会发生漂移。

0x02 原理讲解

2.1 直觉类比

想象一个快递站有 8 个分拣员（rank），每个分拣员专门负责一类包裹（专家）。传统做法是包裹来了就分给对应的分拣员，但如果某类包裹特别多，那个分拣员就会忙死，其他人闲着。

MoonEP 的做法是：先让调度员（planning kernel）看一眼当前所有包裹的分布，然后临时借调几个分拣员去帮忙处理那类特别多的包裹。借调的分拣员不是随便找的，而是从包裹最少的分拣员那里借——这样最后每个分拣员手里的包裹数完全一样。借调的分拣员需要先预习（prefetch）那类包裹的处理方法（权重），处理完后再把结果归还（reduce_grad）给原来的分拣员。

2.2 准专业水平视角

备注：此处符号参见后文 “符号与 shape 约定"。

输入输出 shape：

输入：
  hidden_sh:        [S, H]      bf16  — 本 rank 的 S 个 token
  route_weights_sk: [S, K]      fp32  — 每 token 的 K 个路由权重
  topk_experts_sk:  [S, K]      int32 — 每 token 的 K 个路由专家 id
  tokens_per_expert: [E]        int32 — 每个专家在本 rank 收到的 token 数

输出（dispatch）：
  hidden_nvsh:       [NvS, H]   bf16  — 派发到 VM group 位置的 token
  route_weights_nvs: [NvS]      fp32  — 对应的路由权重
  cu_seqlens:        [E+B]      int32 — 每个 VM group 行的 padded token 末偏移
  plan:              MoonEPCommPlan — 保存用于后续 combine/反向的通信计划

核心公式：

MoonEP 的核心不是数学公式，而是一个构造性算法：

  1. 目标：让每个 rank 恰好持有 S×K 个 token
  2. 容量界：M(I) ≤ E/R，即每 rank 最多冗余 E/R 个专家副本
  3. 构造：反复从"最过载 rank"搬运 token 到"最欠载 rank"，每次精确填满欠载者

与标准做法的对比：

标准 EP:
  rank_i 的 token 数 = Σ_e token_per_expert[e]  (e 属于 rank_i 的专家)
  → 不确定，可能从 0.5*S*K 到 2*S*K 不等

MoonEP:
  rank_i 的 token 数 = S*K  (恒等)
  → 静态形状，固定缓冲，无 host 同步

2.3 专家视角

为什么这个设计在理论上 work？

MoonEP 的理论基础是K3 report 附录 E 的定理：M(I) ≤ E/R。直觉是：即使路由极端不均（所有 token 都路由到同一个专家），那个专家最多需要从其他 rank 借 E/R 个副本——因为每个 rank 只持有 E/R 个本地专家，所以冗余副本数不可能超过 E/R。

KIMI-K3-report-E

关键设计决策：

  1. 为什么用贪婪而非最优？ 贪婪算法 O(R²) 复杂度，最优解是 NP-hard。实践中贪婪已足够好——每轮填死一个 rank，至多 R-1 轮收敛。
  2. 为什么只从一个 remote home group 复制？ 简化 planning kernel 的实现：每个 rank 的冗余专家全部来自同一个远程 rank，减少了跨 rank 通信的复杂度。
  3. 为什么 B = E/R（训练）/ B = 3-4（推理）？ 训练必须精确匹配（否则梯度 reduce 会错），推理可以容忍少量跨 rank 读取。

2.4 负载均衡的数学保证

附录 E 的定理给出冗余容量界：

M(I) ≤ E/R  恒成立

其中 M(I) = max_r redundant_experts(r)
即任何 rank 的冗余专家副本数不超过 E/R。

构造性证明（Model-Viz/src/parallel/moonep.py:110-190 的 CPU 实现）：

  1. 初始化：每 rank 只持有 home 专家的 token
  2. 主循环：找欠载 rank u 和过载 rank o，从 o 搬 need = target - load(u) 个 token 到 u
  3. 终止：每轮填死一个 rank，至多 R-1 轮

2.5 拷贝权重

我们的问题是：在MoonEP真实场景下，拷贝权重不耗费时间嘛？

拷贝权重当然耗时而且是 ms 量级的真实开销，作者也把它算进了 benchmark。——MoonEP 的做法不是"不耗时"，而是三件事：摊薄、隐藏、并且承认它。

README :25 自己有说明：

The comparison counts MoonEP's extra kernels: MoonEP adds planning and weight-prefetch kernels that DeepEP does not need, and they are already stacked in the bars above. Even with the whole critical path included, total dispatch time is on par with DeepEP v2's dispatch alone and pulls ahead under imbalance.

也就是说：开销真实存在，作者没有藏它，benchmark 柱状图里 planning 与 prefetch 两个额外 kernel 是叠进去的。结论是"算上这笔开销后仍然不亏、且在不均衡时净赚"。这跟"不耗时"是两回事。

comm_vs_deepep

MoonEP 让它变得可接受靠三件事：

  • ① 预取 slot 的物理内存跨层共享，显存代价是 B 份权重总计而非每层（README:54）；
  • ② 拷贝跑在高优先级独立 comm stream 上、只用 2 个 warp 做纯 TMA 传输，藏进主流计算的影子里，反向还能跳过重复预取（api.py:864-879, 840-843）；
  • ③ 训练期 B = E/R 由定理保证拷贝量有紧上界、不随路由倾斜恶化，推理期可降到 3-4 并让溢出专家走对称内存直读（README:56-59）。

省下的零拷贝 epilogue、hot-rank 等待、per-layer host 同步和显存碎片，加起来超过这笔支出——这是一笔用有界带宽换无界抖动的交易，不是一笔免费的午餐。

下面拆开讲为什么这笔账划得来。

2.5.1 成本有多大？先把数量级算清楚

按 README :76-77 的默认配置（H20, EP=8）：

S=4096, H=7168, K=8, E=256, R=8  →  B = E/R = 32

单个专家的三个投影（gate/up/down）各 [H, H'] bf16。K2 类模型 H'≈2048：

单专家单投影 = 7168 × 2048 × 2 B ≈ 29.4 MB
单专家三投影 ≈ 88 MB
B=32 个 slot 全填  ≈ 2.8 GB

NVLink（H20 约 900 GB/s 双向、单向 ~450 GB/s）搬 2.8 GB ≈ 6 ms 量级。如果搬运多个专家，这显然不是可以忽略的数字——如果它在关键路径上，这个方案就无意义了。

所以关键全在后面三招。

2.5.2 第一招：物理内存池跨层共享 —— 把成本从"每层"摊成"每步"

这是最容易被忽略的一条。README :54：

Rows [E, E+B): local prefetch slots, filled by buffer.prefetch_weight; ... Their physical memory comes from a process-global pool shared by all layers, so the extra cost is B expert weights per projection in total, not per layer.

显存成本是 B 份权重总计，不是每层 B 份。60 层模型不会变成 60×2.8 GB。

但请注意区分：显存是跨层共享的，拷贝动作仍然每层都要做（每层路由不同、experts_to_copy 不同）。所以这一招解决的是"显存爆炸"，不是"带宽开销"。带宽靠下面两招。

2.5.3 第二招：独立 comm stream 异步 —— 把拷贝藏进计算的影子里

这是真正让开销"消失"的机制。api.py:816-879 prefetch_weight：

if not async_finish:                                  # :855  同步模式（会阻塞）
    self._run_prefetch_weight_on_current_stream(...)
    return None

main_stream = torch.cuda.current_stream()
comm = self._comm_stream                              # :864  独立通信流
self._record_streams((plan.experts_to_copy, *weight_prefetch_args), comm)  # :867
input_ready = main_stream.record_event()              # :868
comm.wait_event(input_ready)                          # :869  只等输入就绪

with torch.cuda.stream(comm):                         # :871  ★在 comm stream 上跑
    self._run_prefetch_weight_on_current_stream(...)
    done = comm.record_event()                        # :877
return done                                           # :879  返回 event，不阻塞

而这个 comm stream 被刻意设成高优先级（api.py:450, 470-471）：

comm_stream_priority: priority of the async comm stream; the default **-1** gives it **higher priority than the main stream**.

调用顺序（README :86-102）是精心排的：

buffer.dispatch(...)        → 返回 plan（含 experts_to_copy）
buffer.prefetch_weight(...) → 立刻在 comm stream 上启动权重搬运
        ┊                      （此时主流还在跑 attention / 上一层 FFN / dispatch 尾部）
        ┊  ← 权重拷贝在这段影子里完成
group GEMM 需要权重时才等 event

关键洞察：权重拷贝与 token dispatch 走不同硬件路径。token dispatch 是 NVLink 上的小消息散射（延迟敏感），权重 prefetch 是大块连续 TMA 传输（带宽敏感）。两者叠加时，大块传输能把 NVLink 上 dispatch 留下的带宽空隙填满，而不是简单相加。

prefetch.py 的 kernel 实现也是为"尽量不占 SM"设计的：只用 2 个 warp（PRODUCER_WARP / CONSUMER_WARP，:216, :239），全程 TMA 异步拷贝（cp_async_bulk），SM 几乎只做地址计算。:272-275 那段早释放注释解释了为什么不滞后 stages-1 组——就是为了把单 SM 带宽榨到 tile_bytes/remote_latency 以上。

api.py:840-843 还有一处省钱：

Kept separate from dispatch so the plan-reuse path (combine bwd) can skip re-prefetching.

combine 的 backward 复用已存的 plan（README :144："planning is skipped and no prefetch is needed"）——反向传播不重复拷权重。所以一次前反向只拷一次，不是两次。

2.5.4 第三招：训练强制 B = E/R —— 用定理保证"拷贝量有上界"

README :56-59 这段是整个设计的收口：

Training: must use B = E/R — the planner duplicates experts from at most one remote home group per rank (≤ E/R experts), so every expert the group GEMM touches is local.

Inference (prefetch only, no gradients): B < E/R is allowed, B = 3–4 is recommended. If a rank ever needs more distinct remote experts than B, the group GEMM reads the overflow weights straight from the home rank through the symmetric mapping — slightly slower, with no impact on correctness.

两个信息量极大的点：

① 训练期 B = E/R 的理由正是那条定理。 "duplicates experts from at most one remote home group per rank" —— 这跟本项目 moonep.py:26-32 docstring 写的证明完全同源：

因为每个 rank 至多被填一次，它收来的全部远端 token 来自同一个其他 rank，那个 rank 只 home E/R 个专家 —— 于是每 rank 冗余副本 ≤ E/R。

所以 B = E/R 不是拍脑袋的经验值，而是理论紧界。拷贝量因此有硬上界，不会随路由倾斜恶化——这正是 README :32 说的 "iteration time stays flat at every imbalance level"。

② 推理期 B 可以远小于界（3-4），代价是"溢出走 symmetric memory 直读"。 这是个漂亮的降级路径：拷不下的专家，group GEMM 直接通过对称内存映射读远端权重（memory semantics，靠 cu_seqlens 寻址），"slightly slower, with no impact on correctness"。推理没有梯度、B 可以砍到 1/8，拷贝成本随之暴跌。

2.5.5 账最终怎么算平的？

MoonEP 花的：planning kernel + weight prefetch（有界、异步、跨层摊薄） MoonEP 省的（README :23-24, :31-32）：

  1. 零拷贝省掉的 epilogue：token 直接写到远端 rank 上的最终专家分组位置，"no permute in, no permute out"，消掉了 "comm-buffer → user-buffer copy that dominates the epilogue"。这笔省下的量本身就很大。
  2. 不再被最热 rank 拖累：DeepEP 的延迟由 hottest rank 决定，maxvio 涨则稳定恶化；MoonEP 因为每 rank 恰好 S×K，comm 时间"almost flat"。
  3. 消掉 per-layer host 同步：静态 shape 让 grouped GEMM 形状 launch 前已知（:9）。这一条省的是流水线停顿，往往比带宽更值钱。
  4. 不再 OOM：DeepEP 的动态激活 shape 碎片化显存，"training OOMs at high imbalance"；MoonEP 全静态，"never OOMs"。

净结果（:25）：总 dispatch 时间与 DeepEP v2 的 dispatch 单独持平，不均衡下反超；combine 在所有档位都明显更快。

0x03 子知识点拆解

3.1 子知识点树状结构

MoonEP 源码级分析
├── 4.0 共享抽象与设计动机
├── 4.1 Buffer API 与 NVL 共享内存管理
│   ├── 4.1.1 VMM 对称内存分配（buffer.py）
│   ├── 4.1.2 Multicast 视图（buffer.py）
│   ├── 4.1.3 Context 创建与缓冲区布局（api.py）
│   └── 4.1.4 Buffer 类 API（api.py）
├── 4.2 Planning Kernel（4-phase 贪婪规划）
│   ├── 4.2.1 Phase A: 排序与计数
│   ├── 4.2.2 Phase B: 贪婪冗余规划
│   ├── 4.2.3 Phase C: 分配目标位置
│   ├── 4.2.4 Phase D: 扫描输出结构
│   └── 4.2.5 MoonEPCommPlan 数据结构
├── 4.3 Dispatch Kernel（warp-specialized 派发）
│   ├── 4.3.1 G2S Producer（TMA 流水线）
│   ├── 4.3.2 S2G Consumer（payload 写入 + weight scatter）
│   ├── 4.3.3 Zero Warp（段填充清零）
│   └── 4.3.4 Dedup Builder Warps（dedup 结构构建）
├── 4.4 Dispatch Epilogue（in-place 重复展开）
├── 4.5 Combine Kernel（3-stage fp32 累加）
├── 4.6 Combine Prologue（in-place 重复累加）
├── 4.7 Prefetch Kernel（2D TMA 权重预取）
├── 4.8 Grad Reduce Kernel（fp32 梯度归约）
└── 4.9 跨 rank 同步基础设施

3.2 官方代码位置标注

子知识点 源码文件 关键行号
Buffer API moonep/api.py Buffer.__init__:439, dispatch:685, combine:881
VMM 分配 moonep/buffer.py create_nvl_dist_tensor:224, create_nvl_single_owner_tensor:345
Planning Kernel moonep/planning.py PlanningKernel.__call__:365, kernel:519
MoonEPCommPlan moonep/planning.py MoonEPCommPlan:31
Dispatch Kernel moonep/dispatch.py DispatchKernel.kernel:268, launch_dispatch:839
Dispatch Epilogue moonep/dispatch_epilogue.py DispatchEpilogueKernel.kernel:168
Combine Kernel moonep/combine.py CombineKernel.kernel:210, launch_combine:565
Combine Prologue moonep/combine_prologue.py CombinePrologueKernel.kernel:197
Prefetch Kernel moonep/prefetch.py PrefetchKernel.kernel:137, launch_prefetch:317
Grad Reduce moonep/grad_reduce.py GradReduceKernel.kernel:150, launch_grad_reduce:437
PTX 工具 moonep/_common.py grid_sync:249, cross_rank_barrier:274, cp_async_bulk_g2s:414
常量 moonep/constants.py RANK_BITS=7, KIDX_BITS=7, DEDUP_BUILDER_WARPS=4
跨 rank 同步 moonep/inter_rank_sync.py launch_inter_rank_sync:119
CPU 教学实现 Model-Viz/src/parallel/moonep.py plan_redundant_experts:110, MoonEP:207

完整执行流程图如下：

用户调用:
  buffer.dispatch(hidden_sh, route_weights_sk, topk_experts_sk, tokens_per_expert)
    │
    ├── inter_rank_sync()        ← 跨 rank 对齐
    ├── planning()               ← 4-phase GPU planning
    │   ├── Phase A: 排序 + 贪婪规划
    │   ├── Phase B: 分配目标位置
    │   ├── Phase C: dst canonicalization
    │   └── Phase D: 扫描输出结构
    ├── dispatch()               ← warp-specialized 派发
    │   ├── Warp 0: G2S TMA
    │   ├── Warp 1: S2G + weight scatter
    │   ├── Warp 2: zero-fill
    │   └── Warp 3+: dedup builder (3-pass)
    ├── dispatch_epilogue()      ← in-place dedup 展开
    └── 返回 (hidden_nvsh, route_weights_nvs, cu_seqlens, plan)

  buffer.prefetch_weight(plan, full_gate_weight, full_up_weight, full_down_weight)
    └── prefetch() × 3          ← 2D TMA 权重预取 (gate/up/down)

  # Expert GEMM 在外部执行，输入/输出都是 NVL buffer view

  buffer.combine(plan, hidden_nvsh, route_weights_nvs)
    ├── inter_rank_sync()
    ├── combine_prologue()      ← in-place fp32 dedup 累加
    ├── combine()               ← 3-stage fp32 ACC + weight gather
    └── 返回 (hidden_sh, route_weights_sk)

  # Backward:
  buffer.dispatch(grad_output_sh, plan=plan)   ← re-dispatch (复用 plan)
  buffer.combine(plan, grad_hidden_nvsh)       ← combine bwd
  buffer.reduce_grad(plan, ...)                ← fp32 梯度归约

3.3 本项目复现

参考代码：MoonEP-master/moonep/ 全部 12 个模块

教学实现：parallel/moonep.py（288 行 CPU 模拟）

MoonEP 的核心优化（VMM、NVLink multicast、cp.async.bulk、PDL）都是 GPU 专用硬件特性，无法在 CPU 上复现。教学实现覆盖了规划算法但不覆盖内核级优化。

原因：硬件限制。

0x04 子知识点深度分析

4.0 共享抽象与设计动机

MoonEP 的所有 kernel 共享一套核心抽象，理解它们是读懂整个代码库的前提。

4.0.1 符号与 shape 约定

符号          含义                              典型值
─────────────────────────────────────────────────────────
S             每 rank 输入 token 数              4096
K             路由 top-k                        8
E             EP 组内总路由专家数                 256
R             EP rank 数                        8
epn           每 rank 本地专家数 = E/R           32
B             权重预取槽数（训练=E/R，推理=3-4）    32
N = S*K       总 token-专家指派数                32768
NvS           派发后每 rank token 槽位数         N + padding
H             隐藏维度                          7168
H'            专家 FFN 中间维度                  2048

所有 rank 通过 NVLink 共享两块关键缓冲区：

hidden_buf [R * NvS_padded, H] bf16 — token 数据：

┌─────────────────────────────────────────────────────┐
│ rank 0 的 NvS_padded 行  │ rank 1 的 ...  │ ...      │
│ [0, NvS) 有效, [NvS, NvS_padded) padding   │         │
└─────────────────────────────────────────────────────┘

meta_buf [R * meta_chunk_padded] int32 — 元数据（多用途复用）：

偏移区域                    用途
[0, NvS)                   weights (fp32, alias 为 int32)
[NvS, NvS + R*E)           tpe gather (tokens_per_expert 跨 rank 收集)
[NvS + R*E, ...)           planning 内部 scratch/staging
[..., +N4)                 topk0 offload scratch
[..., +N4)                 ORDER, c1 order consumed by c2/passB
[..., +N4)                 ORDER0, rank1 rank0-c1 scratch
[..., +3)                  cross-rank phase barrier
[..., +NvS)                src_info scratch (planning 发布, dispatch builder 消费)

这种布局的关键设计：所有跨 rank 数据交换都通过 NVLink 共享内存完成，无需显式通信原语。

4.0.3 MoonEPCommPlan 数据结构

官方代码中，planning.py:31-88 定义了 planning kernel 的输出，贯穿整个通信流程：

@dataclass(frozen=True, slots=True)
class MoonEPCommPlan:
    dst: torch.Tensor              # [N] int32 — 每个 token 的目标位置
                                   # 非负: dest_rank * NvS + local_offset (拷贝 payload)
                                   # 负值: -raw_dst - 1 (duplicate, 只散播 weight)
    experts_to_copy: torch.Tensor  # [R, B] int32 — 每 rank 要预取的远程专家 id
    zero_fill_ranges: torch.Tensor # [E+B, 2] int32 — 每个 VM group 行的 padding 范围
    remote_stats: torch.Tensor     # [2] int32 — 远程专家统计
    dup_groups: torch.Tensor       # [NvS, 3] int32 — (primary_loff, dup_start, dup_n)
    dup_loffs: torch.Tensor        # [NvS] int32 — duplicate slot 偏移
    dup_counts: torch.Tensor       # [2] int32 — (n_groups, n_dup_loffs)
    N, R, E, B, NvS, K: int        # 形状常量

dst 的负编码是 MoonEP 的核心技巧：同一个 token 被路由到同一 rank 的多个专家时，只有第一个 k 路径拷贝 payload（dst >= 0），后续路径标记为 duplicate（dst < 0），只散播 weight。重复行通过 dispatch epilogue 在 NVL shard 上 in-place 展开。

4.1 Buffer API 与 NVL 共享内存管理

4.1.1 要解决的问题

EP 通信的核心挑战是跨 rank 数据交换。传统做法用 all_to_all 通信原语，但：

  • 每次 all_to_all 需要 host-device 同步
  • 数据在发送端和接收端各存一份，浪费显存
  • 无法做到 zero-copy

MoonEP 的解决方案：用 NVLink 共享内存（VMM）让所有 rank 直接读写同一块物理内存。

4.1.2 技术难点

  1. VMM 对称内存分配：需要 CUDA driver 级别的 cuMemCreate + cuMemAddressReserve + cuMemMap，不是 PyTorch 的 torch.cuda.tensor 能搞定的
  2. 跨 rank 内存共享：每个 rank 的 GPU 都能看到同一块物理内存，但需要通过 fabric handle（跨节点）或 POSIX fd（节点内）交换映射信息
  3. Multicast 支持：NVSwitch 的 SHARP 硬件可以让一次写操作同时复制到所有 rank 的本地 chunk

4.1.3 架构与类依赖关系图

buffer.py
├── nvl_dist_alloc()           ← pybind11 → cuMemCreate (CUDA driver)
├── nvl_dist_map()             ← pybind11 → cuMemMap (CUDA driver)
├── nvl_fabric_supported()     ← 检查 NVLink fabric 支持
├── nvl_multicast_create()     ← cuMulticastCreate (NVSwitch SHARP)
├── nvl_multicast_bind_map()   ← 绑定物理内存到 multicast 对象
├── create_nvl_dist_tensor()   → 所有 rank 可读写的 NVL 张量
├── create_nvl_dist_multicast_tensor() → NVL 张量 + multicast 视图
└── create_nvl_single_owner_tensor()   → 单 owner 的 NVL 张量

api.py
├── _create_context()          → 分配 hidden_buf + meta_buf + 本地 scratch
└── Buffer 类                  → 顶层 API，封装 dispatch/combine/prefetch/reduce_grad

4.1.4 数据流

_create_context() 被 Buffer.__init__ 调用:
  1. 计算 NvS = S*K + (token_padding-1)*2*E/R  (逻辑槽位)
  2. NvS_padded = pad_dim0_for_alignment([NvS, H], bf16)  (VMM 对齐)
  3. hidden_buf = create_nvl_dist_tensor([NvS_padded, H], bf16)  ← NVL 共享
  4. meta_buf, meta_mc = create_nvl_dist_multicast_tensor([meta_chunk_padded], int32)  ← NVL + multicast
  5. 分配本地 scratch: alloc, group_tokens, z, local_hist, grid_sync_bar, primary_packed, kmask, kidx_to_loff, builder_bar

4.1.5 关键类/函数职责表

函数/类 文件:行号 输入 输出 职责
nvl_dist_alloc buffer.py (pybind11) shape, dtype, use_fabric keepalive, shareable, owned_handle 分配 VMM 物理内存
nvl_dist_map buffer.py (pybind11) chunk_shape, dtype, shareables, local_rank, world_size, use_fabric full_tensor 创建跨 rank 映射
create_nvl_dist_tensor buffer.py:224 chunk_shape, dtype, local_rank, world_size, group tensor 所有 rank 可读写的 NVL 张量
create_nvl_dist_multicast_tensor buffer.py:251 chunk_shape, dtype, local_rank, world_size, group (full_tensor, mc_view) NVL 张量 + multicast 视图
create_nvl_single_owner_tensor buffer.py:345 shape, dtype, owner_rank, local_rank, group tensor 单 owner 的 NVL 张量
_create_context api.py:213 S, H, K, E, num_ep_ranks, ... ctx dict 预分配所有缓冲区
Buffer.__init__ api.py:439 S, H, K, E, num_ep_ranks, ... Buffer 创建 context + comm stream
Buffer.dispatch api.py:685 hidden_sh, route_weights_sk, topk_experts_sk, tokens_per_expert (hidden_nvsh, route_weights_nvs, cu_seqlens, plan) dispatch 前向
Buffer.combine api.py:881 plan, hidden_nvsh, route_weights_nvs (hidden_sh, route_weights_sk, event) combine 前向

4.1.6 实现机制详解

VMM 分配流程（buffer.py:224-248）：

def create_nvl_dist_tensor(chunk_shape, dtype, local_rank, world_size, group):
    use_fabric = _use_fabric_for_group(group)           # 1. 决定用 fabric 还是 fd
    keepalive, shareable, owned_handle = nvl_dist_alloc(
        shape=chunk_shape, dtype=dtype, use_fabric=use_fabric)  # 2. cuMemCreate
    try:
        return _map_nvl_dist_tensor(                    # 3. cuMemMap
            chunk_shape, dtype, shareable, keepalive,
            local_rank, world_size, group, use_fabric,
        )
    finally:
        nvl_release_mem_handle(owned_handle)             # 4. 释放 handle（映射已建立）

Fabric vs FD 的选择（buffer.py:37-63）：

  • 跨节点（world_size > 1）：必须用 fabric handle（NVLink fabric 支持）
  • 节点内（world_size == 1）：可以用 POSIX fd（更简单）
  • 环境变量 MOONEP_MEM_HANDLE_TYPE 可覆盖：auto/fabric/fd

4.1.7 对比基线

维度 标准 EP (all_to_all) MoonEP (NVL VMM)
数据拷贝次数 2（发送端打包 + 接收端解包） 0（直接读写共享内存）
Host 同步 每次 all_to_all 需要 无需
显存占用 2× buffer（发送 + 接收） 1× buffer（共享）
通信延迟 受 PCIe/NVLink 带宽限制 NVLink 直连，延迟更低

4.1.8 反方辩手

a) 反对最强理由：VMM 对称内存需要特定硬件支持（NVLink + NVSwitch），不是所有 GPU 都能用。

b) 最可能错的假设：假设 NVLink 带宽足够覆盖所有 EP 场景。当 EP 组很大（R=64+）时，跨节点 NVLink fabric 的延迟可能成为瓶颈。

c) 胜出条件：节点内 EP（R ≤ 8），NVLink 带宽充足；失败条件：跨节点 EP，或 NVLink 带宽不足。

d) 更稳健的做法：DeepEP v2 的 all_to_all 方案，虽然性能差一点，但兼容所有硬件。

e) 受伤害的利益相关方：使用非 NVIDIA GPU 的训练方（如 AMD MI300X）。

4.1.9 计算复杂度分析

  • FLOPs：NVL 分配本身无计算，纯内存操作
  • 内存：hidden_buf = R × NvS_padded × H × 2 bytes (bf16)；meta_buf = R × meta_chunk_padded × 4 bytes (int32)
  • 参数数量：无新增可训练参数（通信库，不是模型组件）

4.1.10 常见疑问 Q&A

Q: 为什么 meta_buf 用 int32 而不是 fp32？
A: meta_buf 是多用途复用的——权重区域逻辑上是 fp32，但通过 view(torch.int32) 共享同一块内存。int32 的位操作（atomicAdd、atomicOr）在 planning/dispatch builder 中更方便。

Q: zero_copy=True 为什么不能跨通信调用保存 tensor？
A: 因为返回的是 NVL buffer 的 view，下一次 dispatch/combine 会覆盖同一块内存。autograd 不能把 view 存到 SavedTensor 里——那会导致 backward 时读到被覆盖的数据。

4.2 Planning Kernel（4-phase 贪婪规划）

4.2.1 要解决的问题

Planning kernel 是 MoonEP 的"大脑"——它决定每个 token 应该被发到哪个 rank 的哪个专家槽位。核心挑战：

  1. 路由不均：不同专家收到的 token 数差异巨大
  2. 完美均衡：每个 rank 必须恰好收到 S×K 个 token
  3. 冗余规划：需要决定哪些远程专家要被复制到本地
  4. 在线计算：规划必须在 GPU 上完成，不能回传 host

4.2.2 技术难点

  • 4 个 phase 必须在一个 cooperative grid launch 中完成（不能拆成多个 kernel，因为中间结果在 SMEM/GMEM 之间传递）
  • Phase B 的贪婪算法需要warp 级别的 reduction（找最大/最小值）
  • Phase C 的二分搜索需要 O(log R) 步骤，每步读 global memory
  • 输出结构（dst, cu_seqlens, experts_to_copy, zero_fill_ranges）的大小与 E, R, B 相关，必须在编译时确定

4.2.3 架构与执行流

PlanningKernel.__call__ (planning.py:365)
  └── kernel (planning.py:519)
       ├── Phase A (planning.py:601-702):
       │   ├── Step 1: tpe gather → rank0 chunk
       │   ├── Step 2: cross_rank_barrier
       │   ├── Step 3: tpe_cumsum 计算
       │   ├── Step 4: 贪婪冗余规划 (balance loop)
       │   ├── Step 5: alloc_tensor 计算
       │   └── Step 6: cu_seqlens + zero_fill_ranges + experts_to_copy
       │
       ├── Phase B (planning.py:961-977):
       │   └── C2: 根据 ORDER 将 token 分配到目标 rank/offset
       │
       ├── Phase C (planning.py:979-1077):
       │   └── 跨 rank barrier + dst canonicalization (负编码 duplicate)
       │
       └── Phase D (planning.py:1003-1131):
           └── 扫描输出结构写入 cu_seqlens, zero_fill_ranges, experts_to_copy

4.2.4 Phase A 详解：排序与冗余规划

Phase A 是最复杂的部分，包含 6 个步骤：

Step 1-2: tpe gather + cross-rank barrier（planning.py:602-609）

# 每个 rank 把自己的 tpe (tokens_per_expert) 写入 meta_buf 的 rank 0 区域
copy_v4_remote(meta, TPE_OFF + rank * E, tpe, E, pid, tid, num_threads, num_sms)
# rank 0 还收集所有 rank 的 topk 原始数据
if rank == 0:
    copy_v4_remote(meta, ms + TOPK0_OFF, topk, N, pid, tid, num_threads, num_sms)
# 跨 rank 同步，确保所有 rank 的 tpe 都可见
cross_rank_barrier(meta, ms, BARRIER_OFF, rank, R, bar_p, num_sms, num_threads, tid)

Step 3: tpe_cumsum 计算（planning.py:625-670）

rank 0 计算跨 rank 的 token 数前缀和：

# 对每个 expert，计算 rank 间的累积和
# tpe_cumsum[r, e] = sum(tpe[0..r-1, e]) + tpe[r, e]
# 这告诉每个 rank：expert e 的第 global_rank 个 token 应该发到哪里

Step 4: 贪婪冗余规划（planning.py:671-702）

这是 MoonEP 的核心算法，在 GPU 上实现：

# balance[j] = group_tokens[j] - CAP  (CAP = S*K)
# 循环：
#   surplus = max(balance) → 有盈余的 rank
#   deficit = min(balance) → 有亏空的 rank
#   move_tokens = -deficit → 从 surplus 搬 move_tokens 个 token 到 deficit
#   z[surplus_rank, deficit_rank] = move_tokens
# 终止条件：surplus <= 0 或 deficit >= 0

warp 级别的 argmax/argmin 用 reg_scan_argmax_min_idx（planning.py:258-263）实现，纯寄存器操作，无需 shared memory。

Step 5: alloc_tensor 计算（planning.py:703-798）

对每个 owner rank，计算其他 rank 要从它这里拿多少 token：

# alloc_tensor[rank, expert] = 该 rank 为 expert 计算的 token 数
# alloc_cumsum[expert, rank] = alloc_tensor 的行前缀和（供 Phase C 二分搜索）

Step 6: cu_seqlens + zero_fill_ranges（planning.py:799-954）

对每个 dest_rank，计算 VM group 的 padded token 末偏移：

# cu_seqlens[dest_rank, group_idx] = padded_end_offset
# zero_fill_ranges[dest_rank, group_idx] = (pad_start, pad_count)
# experts_to_copy[dest_rank, slot] = 要预取的远程专家 id

4.2.5 Phase B: C2 分配目标位置

Phase B（planning.py:961-1077）根据 ORDER（Phase A 产出的排序序列）将每个 token 分配到目标 rank 和 offset：

# 对每个 token offv = ORDER[idx]:
#   expert_idx = topk[offv]
#   prev = tpe_cumsum[rank-1, expert_idx]  (本 rank 之前的累积)
#   global_rank = prev + (idx - s_expoff[expert_idx])  (全局排名)
#   二分搜索找到目标 rank: alloc_cumsum[expert_idx, mid] > global_rank
#   lo = expert_off[lo, expert_idx]  (目标 offset)
#   dst[offv] = lo * NvS + bo + (global_rank - pc)
#   src_info[lo * ms + SRC_INFO_OFF + loff] = rank * NvS + offv  (发布源信息)

4.2.6 Phase C: dst canonicalization

Phase C（planning.py:1079-1113）处理 duplicate 检测：

# 对每个 token 的 K 个路由路径，检测是否有重复目标 rank
# seen_lo, seen_hi: 128-bit 位掩码
# 如果 dst[k] 的目标 rank 已经在 seen 中，则标记为 duplicate:
#   dst[base_idx + k] = -(dst_vals[k]) - 1
# 否则标记为 primary:
#   seen_lo |= (1 << d)

这个位掩码技巧让 duplicate 检测完全在寄存器中完成，无需 shared memory。

4.2.7 Phase D: 扫描输出结构

Phase D（planning.py:1003-1131）将 planning 的结果从 meta_buf 拷贝到输出张量：

# cu_seqlens[group_idx] = s_pd_cu_stage[local_group]
# zero_fill[group_idx*2] = s_pd_zfr_stage[local_group*2]
# experts_to_copy[etc_idx] = s_pd_etc_stage[local_idx]
# remote_stats[i] = meta[PB + STATS_SUB + rank*2 + i]

使用 cp.async.bulk G2S + S2G 的流水线，避免直接从 global memory 读取。

4.2.8 数据流图

输入:
  topk_experts_sk: [S, K] int32  →  topk_flat: [N] int32
  tokens_per_expert: [E] int32   →  tpe: [E] int32

Phase A:
  tpe → meta_buf[TPE_OFF + rank*E]  (每个 rank 写自己的)
  cross_rank_barrier
  rank 0: tpe_cumsum[R, E], group_tokens[R], z[R,R]
  rank 0: alloc_tensor[R, E], alloc_cumsum[E, R]
  rank 0: cu_seqlens[R, E+B], zero_fill_ranges[R, E+B, 2]
  rank 0: experts_to_copy[R, B], remote_stats[R, 2]
  rank 0: multimem_st → meta_mc (multicast 到所有 rank)

Phase B:
  ORDER[N] → dst[N] int32, src_info[NvS] int32

Phase C:
  dst[N] canonicalized (负编码 duplicate)

Phase D:
  meta_buf → cu_seqlens[E+B], zero_fill_ranges[E+B, 2], experts_to_copy[R, B]

输出:
  dst[N]              → plan.dst
  cu_seqlens[E+B]     → 返回值
  experts_to_copy[R,B] → plan.experts_to_copy
  zero_fill_ranges[E+B,2] → plan.zero_fill_ranges
  dup_groups/dup_loffs/dup_counts → 由 dispatch builder 填充

4.2.9 对比基线

维度 标准 EP MoonEP Planning
规划位置 CPU GPU (cooperative grid)
规划延迟 O(E log E) host O(E + R²) GPU warp-level
中间结果 需要 host-device 同步 全部在 GPU 上完成
负载均衡 无保证 完美均衡 (S*K 恒等)

4.2.10 反方辩手

a) 反对最强理由：planning kernel 的 4-phase 设计过于复杂，cooperative grid launch 的调试极其困难。
b) 最可能错的假设：假设 routing 分布不会极端到让 B = E/R 不够用。如果某个 rank 需要 > E/R 个不同的远程专家，多余的会直接从 home rank 通过 NVLink 读取（慢但正确）。
c) 胜出条件：E/R ≥ K（每 rank 至少有 K 个本地专家），routing 有一定随机性；失败条件：E/R < K 且 routing 极端集中。
d) 更稳健的做法：静态预分配 + all_to_all，牺牲一点性能但实现简单。
e) 受伤害的利益相关方：推理服务方（latency 敏感，planning 开销占比更高）。

完成度：已实现并运行通过（moonep/planning.py）

4.2.11 计算复杂度分析

  • FLOPs：Phase A 的 tpe_cumsum: O(E×R)；贪婪规划: O(R²)；Phase B 的 C2: O(N×log R)；总: O(E×R + R² + N×log R)
  • 内存：SMEM: O(E + R² + NvS)；GMEM: meta_buf 的 planning 区域
  • 参数敏感性：
    • num_sms：越大规划越快，但受 SM 数量限制
    • BLOCK_SIZE_P2 = 2048：warp 级别操作的粒度，过大降低并行度
    • token_padding = 128：padding 越大 NvS 越大，浪费显存但对齐更好

4.3 Dispatch Kernel（warp-specialized 派发）

4.3.1 要解决的问题

Dispatch kernel 的任务是：根据 planning 产出的 dst 数组，把 [S, H] 的 token 数据直接写到 NVL 共享内存中对应 rank 的对应槽位——不需要 all_to_all，不需要显式 permute。

核心挑战：

  1. 零拷贝：token 直接写到目标 rank 的最终位置，无需中间缓冲
  2. Per-row TMA：每行用 cp.async.bulk 做一次 128-bit 对齐的 TMA 拷贝
  3. Dedup 结构构建：在 dispatch 的同时，builder warps 需要扫描 src_info 构建 dup_groups/dup_loffs/dup_counts
  4. 段填充清零：padding 行需要被清零（expert GEMM 会读到这些行）

4.3.2 Warp 分工

dispatch.py:50-82 定义了 warp 布局：

Warp 0 (PRODUCER):  G2S producer — 从 hidden_sh 读一行到 stage_smem
Warp 1 (CONSUMER):  S2G consumer — 从 stage_smem 写到目标 rank 的 NVL 位置
Warp 2 (ZERO):      zero warp — 清零 padding 行
Warp 3+ (DEDUP_BUILDER): dedup builder warps — 构建 dedup 结构（仅 fresh planning）

总线程数：96 + 32 * DEDUP_BUILDER_WARPS（dispatch.py:77）

4.3.3 G2S Producer（TMA 流水线）

dispatch.py:360-383：

# Warp 0 驱动 TMA G2S 流水线
# 每个 token 一行：cp_async_bulk_g2s(smem_row, gmem_row, H_BYTES, mbar)
# PipelineTmaAsync 管理多阶段流水线（stages = 2..16，取决于 H 和 smem budget）
for li in cutlass.range(n_tok, unroll=1):
    load_pipe.producer_acquire(load_state)
    if cute.arch.lane_idx() == 0:
        cp_async_bulk_g2s(s_row_int, g_row_int, H_BYTES, mbar_int)
    load_state.advance()

4.3.4 S2G Consumer（payload 写入 + weight scatter）

dispatch.py:387-448：

# Warp 1 从 stage_smem 读取，根据 dst 决定写到哪里
for li in cutlass.range(n_tok, unroll=1):
    load_pipe.consumer_wait(use_state)
    if cute.arch.lane_idx() == 0:
        for k in range(K):
            dst_val = dst[sK + k]
            store_token = dst_val >= 0          # 非负 = primary, 拷贝 payload
            raw_dst = dst_val if store_token else -dst_val - 1
            drank = raw_dst // NvS
            loff = raw_dst % NvS
            if store_token:
                drow = drank * NvS_padded + loff
                cp_async_bulk_s2g(s_row, gmem_dst + drow * H, H_BYTES)
            # weight scatter: 所有路径都写 weight（包括 duplicate）
            meta_tensor[drank * meta_stride + weights_off + loff] = wb
    use_state.advance()

关键点：weight scatter 对所有 K 路径都执行（包括 duplicate），因为 route weights 是 per-topk 的值，不能 dedup。

4.3.5 Zero Warp（段填充清零）

dispatch.py:467-497：

# Warp 2 并行清零 padding 行
# zero_fill_ranges[e] = (pad_start, n_pad_rows)
while e < zero_groups:
    pad_start, n_pad = zero_fill_ranges[e*2], zero_fill_ranges[e*2+1]
    if n_pad > 0:
        for j in range(n_pad):
            loff = pad_start + j
            local_row = rank * NvS_padded + loff
            cp_async_bulk_s2g(zero_smem, gmem_dst + local_row * H, H_BYTES)
            if with_weights:
                meta_tensor[rank * meta_stride + weights_off + loff] = 0
    e += num_sms  # grid-stride

Zero warp 使用预初始化的 zero_smem（所有线程协作清零一次，后续直接复用），避免每次清零都写一遍。

4.3.6 Dedup Builder Warps（3-pass 构建 dedup 结构）

dispatch.py:502-681，这是 dispatch kernel 中最复杂的部分：

Pass 1: Primary election（dispatch.py:557-576）

# 每个 builder warp 扫描自己负责的 src_info chunk
# 对每个 (source_rank, token) key：
#   packed = (kidx << NvS_BITS) | loff
#   atom_min_relaxed_gpu_s32(primary_packed[key], packed)  // 选最小 loff 作为 primary
#   atom_or_relaxed_gpu_b32(kmask[key], 1 << kidx)        // 记录哪些 kidx 存在
#   kidx_to_loff[key * K + kidx] = loff                    // 记录 kidx → loff 映射

Pass 2a: Classification（dispatch.py:593-611）

# 每个 lane 分类自己 chunk 的 slot：
#   primary_loff = packed & NvS_MASK
#   dup_count = popc(mask) - 1
#   if loff == primary_loff and dup_count > 0:
#       lane_grp_n += 1     // 本 lane 负责的 group 数
#       lane_dup_n += dup_count  // 本 lane 负责的 dup 数

Pass 2b: Output emission（dispatch.py:648-681）

# warp-inclusive-scan 得到 lane 间的偏移
# 每个 lane 用 2 个 atomicAdd 预留输出空间
# 然后写入 dup_groups[my_grp*3] = (loff, my_dup, dup_count)
#         dup_loffs[my_dup + pos] = cur_dup_loff

4.3.7 执行时序

Warp 0 (producer):  acquire → G2S → advance → acquire → G2S → ...
Warp 1 (consumer):  wait → 读 dst → S2G(多个 k) → commit → wait → ...
Warp 2 (zero):      并行清零 padding 行
Warp 3+ (builder):  Pass 1 → cross_warp_sync → Pass 2a → cross_warp_sync → Pass 2b
Exit:               cross_rank_barrier → pdl_trigger_dependents

4.3.8 数据流

输入:
  hidden_sh: [S, H] bf16           → stage_smem → NVL hidden_buf
  route_weights_sk: [S, K] fp32    → meta_buf weights 区域
  dst: [N] int32                    → 决定每行写到哪个 rank 的哪个 offset
  zero_fill_ranges: [E+B, 2] int32 → zero warp 读取
  src_info: [NvS] int32            → builder warps 读取

输出:
  hidden_buf: [R*NvS_padded, H] bf16  → 写入目标 rank 的 payload 行
  meta_buf weights: [NvS] fp32         → 写入 route weights
  dup_groups: [NvS, 3] int32           → builder 构建
  dup_loffs: [NvS] int32               → builder 构建
  dup_counts: [2] int32                 → builder 构建

4.3.9 对比基线

维度 标准 EP (all_to_all + permute) MoonEP Dispatch
数据拷贝 2次（打包 + 解包） 1次（直接写到最终位置）
Permute 显式 permute/unpermute 融合在 dispatch/combine 中
缓冲区 2× S×K×H 1× NvS×H
Dedup 无（每个 token 独立） 构建 dedup 结构 + epilogue 展开

4.3.10 反方辩手

a) 反对最强理由：dedup builder 的 3-pass 设计增加了 dispatch kernel 的延迟，尤其是 Pass 1 的 atom_min_relaxed_gpu_s32 可能成为热点。
b) 最可能错的假设：假设 K ≤ 32。如果 K > 32，kmask 的 u32 位掩码不够用。
c) 胜出条件：K ≤ 32，routing 有一定重复（dedup 收益大）；失败条件：K > 32 或 routing 完全无重复（builder 白跑）。
d) 更稳健的做法：不做 dedup，每个 token 的 K 路径都独立拷贝 payload（牺牲带宽但逻辑简单）。
e) 受伤害的利益相关方：K 很大的模型（如 K=64 的 routing）。

4.3.11 参数敏感性

  • DEDUP_BUILDER_WARPS = 4（constants.py:17）：builder warp 数量。太少则 Pass 1 延迟高；太多则 warp 级别同步开销大。4 是实测甜点。
  • stages（pipeline depth）：取决于 H 和 smem budget。H=7168 时通常 2-4 stages。
  • num_sms：dispatch 的并行度，越大吞吐越高。

4.4 Dispatch Epilogue（in-place 重复展开）

4.4.1 要解决的问题

Dispatch 只拷贝了 primary 行到 NVL shard。对于 duplicate 的 token（被路由到同一 rank 的多个专家），它们的 payload 行还没有被写入。Dispatch Epilogue 负责从 primary 行展开到所有 duplicate slot。

4.4.2 实现机制

dispatch_epilogue.py:44-61：

class DispatchEpilogueKernel:
    """Round-robin duplicate-group expansion on the local NVL shard.

    两个 warp 镜像 dispatch 的数据路径：
    - Warp 0 (producer): 把 B 个 primary 行 G2S 到 smem pipeline stage
    - Warp 1 (consumer): 把每行 S2G 到所有 duplicate slot
    """

核心逻辑（dispatch_epilogue.py:208-305）：

# 从 dup_counts[0] 读取有效 group 数
n_groups = dup_counts_tensor[0]

# Warp 0 (producer): 批量 G2S
while gi < n_groups:
    n = min(B, n_groups - gi)
    # 从 hidden_buf 读取 B 个 primary 行到 stage_smem
    for idx in range(n):
        loff = dup_groups[(gi+idx) * 3]  # primary_loff
        cp_async_bulk_g2s(stage_smem[idx], hidden[loff], H_BYTES, mbar)

# Warp 1 (consumer): 批量 S2G 到 duplicate slot
while gi < n_groups:
    for idx in range(n):
        dup_start = dup_groups[(gi+idx)*3 + 1]
        dup_n = dup_groups[(gi+idx)*3 + 2]
        for k in range(dup_n):
            slot = dup_loffs[dup_start + k]
            cp_async_bulk_s2g(hidden[slot], stage_smem[idx], H_BYTES)

关键设计：primary 行只从 gmem 读一次，然后 fan-out 到所有 duplicate slot。这比每个 duplicate 都从 gmem 读 primary 要快得多。

4.4.3 对比 Combine Prologue

维度 Dispatch Epilogue Combine Prologue
方向 primary → duplicate (fan-out) duplicate → primary (reduce)
数据类型 bf16 (拷贝) fp32 (累加)
无跨 rank barrier 是（纯本地操作） 是（纯本地操作）
被谁消费 Expert GEMM Combine kernel

4.4.4 反方辩手

a) 反对最强理由：每个 duplicate group 都要从 gmem 读一次 primary 行，如果 duplicate 很多（K 很大），读带宽可能成为瓶颈。
b) 最可能错的假设：假设 duplicate group 的大小（dup_n）通常很小（≤ K-1）。如果 routing 高度集中，某些 group 的 dup_n 可能很大。
c) 胜出条件：K ≤ 16，routing 有一定分散性；失败条件：K 很大且 routing 极端集中。
d) 更稳健的做法：在 dispatch 中直接写所有 duplicate 行（不做 dedup），牺牲带宽但延迟更低。

完成度：已实现并运行通过（moonep/dispatch_epilogue.py）

4.5 Combine Kernel（3-stage fp32 累加）

4.5.1 要解决的问题

Combine kernel 执行 dispatch 的逆操作：从 NVL 共享内存中收集 expert output，按 token 聚合（K 次累加），写回 [S, H] 的 token-major 格式。

核心挑战：

  1. FP32 累加：expert output 是 bf16，但累加必须在 fp32 中进行（避免精度损失）
  2. Dedup 优化：duplicate 的贡献已经由 combine prologue 预累加到 primary slot，combine 只需拉取 primary
  3. Weight gather：route weights 需要同步 gather 回 token-major 格式

4.5.2 Warp 分工

combine.py:40-61：

Warp 0 (PRODUCER):  G2S producer — 从 NVL hidden_buf 读一行到 stage_smem
Warp 1-4 (ACC):     fp32 ACC consumers — 128 线程累加 fp32 寄存器
Warp 5 (S2G):       S2G consumer — 把累加结果写回 gmem output
Warp 6 (WEIGHT):    weight-gather epilogue — 并行 gather route weights（仅 with_weights）

4.5.3 关键实现细节

FP32 ACC in Registers（combine.py:363-435）：

# 每个 ACC 线程拥有 H/ACC_THREADS 个 fp32 寄存器
H_PER_THR = H // ACC_THREADS  # = H // 128
acc_reg = cute.make_rmem_tensor(H_PER_THR, Float32)

for s_local in range(n_tok):
    # Prologue: 零初始化累加器
    for j in range(H_PER_THR):
        acc_reg[j] = Float32(0.0)

    for k in range(K):
        dst_val = dst[sK + k]
        accumulate_token = dst_val >= 0  # 只累加 primary (dst >= 0)
        if accumulate_token:
            load_pipe.consumer_wait(use_state)
            for j in range(H_PER_THR):
                idx = j * ACC_THREADS + acc_tid
                acc_reg[j] += Float32(stage_smem[idx, use_state.index])
            load_pipe.consumer_release(use_state)

    # Epilogue: fp32 → bf16 写回
    for j in range(H_PER_THR):
        out_smem[idx, stg_o] = BFloat16(acc_reg[j])
    acc_bar.arrive_and_wait()  # 4 个 ACC warp 同步
    out_pipe.producer_commit(out_state)

Weight Gather（combine.py:485-511）：

# Warp 6 独立于数据路径，直接从 meta_buf 读取 route weights
# 对所有 K 路径（包括 duplicate）都执行 gather
for sk_chunk in range(n_chunks):
    dst_val = dst[sk_global]
    raw_dst = dst_val if dst_val >= 0 else -dst_val - 1
    drank = raw_dst // NvS
    loff = raw_dst % NvS
    wb = meta_tensor[drank * meta_stride + weights_off + loff]
    output_sk[sk_global] = wb

4.5.4 与 DeepEP v2 的性能对比

README 的性能数据（README.md:13-32）：

  • MoonEP 的 combine 时间在所有 imbalance 级别下都显著低于 DeepEP v2
  • MoonEP 的总 dispatch 时间（含 planning + weight prefetch）与 DeepEP v2 的 dispatch 相当，但在 imbalance 时领先
  • MoonEP 的 combine 明显更快，因为 dedup 减少了需要拉取的行数

e2e_vs_deepep

comm_vs_deepep

4.5.5 反方辩手

a) 反对最强理由：6-7 个 warp 的设计增加了 occupancy 压力，可能在某些 GPU 上导致 register spilling。
b) 最可能错的假设：假设 H % 128 == 0。如果 H 不是 128 的倍数，ACC 线程数需要调整。
c) 胜出条件：H ≥ 128 且 K ≥ 2（dedup 收益明显）；失败条件：H 很小（如 H=64）且 K=1（无 dedup）。
d) 更稳健的做法：不做 dedup 优化，直接拉取所有 K 行。

完成度：已实现并运行通过（moonep/combine.py）

4.6 Prefetch + GradReduce（权重预取 + 梯度归约）

4.6.1 Prefetch Kernel

要解决的问题：远程专家的权重在其他 rank 的 GPU 上，dispatch 后需要在本地 GEMM 之前把权重拉过来。

prefetch.py:29-76：

class PrefetchKernel:
    """Persistent 2D TMA remote-expert prefetch."""
    M_BLOCK = 128  # tile 高度
    N_BLOCK = 128  # tile 宽度
    # 用 2D TMA 拷贝 H×H' 的权重 tile

执行流（prefetch.py:137-277）：

# 预扫描：构建 exp_tab/slot_tab（active slot 的 compact 列表）
n_active = 0
for b in range(B):
    e = experts[b]
    if e >= 0:
        exp_tab[n_active] = e
        slot_tab[n_active] = b
        n_active += 1

# Warp 0 (producer): 2D TMA G2S
for tile in range(bidx, total_tiles, num_sms):
    i = tile // TILES_PER_EXPERT
    mt = (tile % TILES_PER_EXPERT) // NTILES
    nt = (tile % TILES_PER_EXPERT) % NTILES
    expert = exp_tab[i]
    cute.copy(tma_g2s, tSgS[(None, (expert*MTILES+mt, nt))], tSsS[(None, stage)])

# Warp 1 (consumer): 2D TMA S2G
for tile in range(bidx, total_tiles, num_sms):
    b = slot_tab[i]
    cute.copy(tma_s2g, tDsD[(None, stage)], tDgD[(None, (b*MTILES+mt, nt))])

4.6.2 Grad Reduce Kernel

要解决的问题：backward 时，remote expert 的权重梯度在 reduce buffer 中，需要归约回 home rank。

grad_reduce.py:46-54：

class GradReduceKernel:
    M_BLOCK = 128
    N_BLOCK = 128
    STAGES = 3           # load pipeline depth
    ACC_THREADS = 128    # 4 个 fp32 ACC warp
    NUM_THREADS = 160    # 1 load + 4 acc

两阶段执行（grad_reduce.py:297-389）：

# Phase 1: warp-specialized accumulate（无 clear）
for work_idx in range(bidx, total_work, num_sms):
    local_expert = alist[work_idx // TILES_PER_EXPERT]
    # Warp 0: streaming remote reduce tiles (G2S TMA)
    # Warps 1-4: seed local grad + accumulate remote + store back

# Phase 2: cross_rank_barrier then local clear
cross_rank_barrier(...)
# 每个 rank 清零自己的 consumed slots（16B vector stores）

关键设计：accumulation never writes the reduce buffers。所有累加在 fp32 寄存器中完成，最后一次性写回 home rank 的 grad buffer。清零只清自己的 consumed slots，不做 remote write clear（因为 NVLink 双向带宽共享，remote write 会拉长 phase 1）。

4.6.3 反方辩手

a) Prefetch: 2D TMA 的 tile 大小固定 128×128，如果 H 或 H' 不是 128 的倍数会浪费带宽。
b) GradReduce: prescan 的 match_any_b32 + atomicAdd 在 warp 0 上串行执行，当 R*B 很大时可能成为瓶颈。
c) 更稳健的做法：不预取权重，直接从 home rank 的 NVLink mapping 读取（内存语义，慢但正确）。README 提到推理时可以这样做。

完成度：已实现并运行通过（moonep/prefetch.py + moonep/grad_reduce.py）

4.7 跨 rank 同步基础设施

4.7.1 grid_sync（_common.py:249-271）

def grid_sync(bar_ptr, nsm, tid):
    """Self-resetting cooperative grid barrier."""
    # SM0 发送 TAG - (nsm-1)，其他 SM 发送 +1
    # 总增量 = TAG = 0x80000000，低 31 位始终为 0
    # top bit 翻转表示所有 SM 到达

4.7.2 cross_rank_barrier（_common.py:274-347）

def cross_rank_barrier(meta_buf, meta_stride, barrier_off, rank, num_ranks, bar_ptr, nsm, num_threads, tid):
    """Self-resetting cross-rank barrier."""
    # Entry fence: fence.proxy.alias (multicast → unicast coherence)
    grid_sync(bar_ptr, nsm, tid)
    # Block 0: 发送 +/-1 信号到 peer 的 phase slot
    # 自旋等待自己的 phase slot 达到目标值
    # Exit fence: fence.proxy.async.global (generic → TMA coherence)
    grid_sync(bar_ptr, nsm, tid)

3 个 int32 slot 的自复位设计（_common.py:286-303）：

  • +0/+1：两个 phase 信号（交替使用）
  • +2：本地 phase/sign 状态

4.7.3 cross_warp_sync（_common.py:349-390）

def cross_warp_sync(bar_ptr, nparticipants, leader):
    """All-participating-warp barrier."""
    # leader warp 发送 TAG - (nparticipants-1)
    # 其他 warp 发送 +1
    # 总增量 = TAG，自复位

用于 dispatch builder 的 4 个 warp 之间的同步。

4.7.4 PTX 工具函数

函数 行号 用途
cp_async_bulk_g2s _common.py:414 TMA G2S 拷贝
cp_async_bulk_s2g _common.py:438 TMA S2G 拷贝
atom_add_release_gpu _common.py:75 GPU scope atomicAdd (release)
atom_add_relaxed_gpu_s32 _common.py:92 GPU scope atomicAdd (relaxed)
atom_min_relaxed_gpu_s32 _common.py:109 GPU scope atomicMin (relaxed)
atom_or_relaxed_gpu_b32 _common.py:125 GPU scope atomicOr (relaxed)
popc_b32 _common.py:159 32-bit popcount
ctz_b32 _common.py:176 Count trailing zeros
ld_acquire_gpu_s32 _common.py:199 GPU scope load (acquire)
ld_acquire_sys_s32 _common.py:232 System scope load (acquire)
red_add_release_sys _common.py:216 System scope reduce add
pdl_trigger_dependents _common.py:394 Programmatic dependent launch trigger
pdl_wait_predecessor _common.py:404 PDL wait

0x05 本项目的实现 & 参考 & 讨论

5.1 具体代码

r"""MoonEP：完美均衡的专家并行调度（论文 §5.2.1 + 附录 E）。

设计动机（why）
-------------
专家并行的老问题是**路由负载不均**：Top-K 路由下不同专家在不同
micro-batch、不同层的负载都是波动的，某几个"热"专家会把 token 堆到少数
rank 上，带来三重代价：
  1. 计算瓶颈：hot rank 变慢，拖慢整次同步；
  2. 显存碎片：路由后的激活 shape 每步每层都不同，无法预分配固定缓冲；
  3. host 同步停滞：形状未知，host 每层都要和设备同步一次才能知道该分配多大
     显存、怎么发 grouped GEMM，直接打穿流水线。

MoonEP 的关键洞察（附录 E）
----------------------------
只要允许 **复制**（redundant expert），就可以让每个 TF 拿到的 token 数
**严格相等**（每个 rank 恰好 ``S*K`` 个），从而重获静态 shape 的三大好处：
  - 通信缓冲固定为 ``S*K``（而不是最坏 ``S*K*R``）；
  - grouped GEMM 的形状在 launch 前已知，无需 host 同步；
  - 负载零波动，无 hot rank。

问题变成：复制多少份才够？**附录 E 的定理**给出了答案：
    M(I) <= E/R  恒成立，且界基本是紧的。
即每 rank 预留给"非本地主家专家"的槽位（redundant slots）只需 ``E/R`` 个，
就**永远**存在一个可行规划——训练永不因规划失败中断。

构造性证明（plan_redundant_experts 就是它的直接翻译）：
  反复挑选"欠载 rank u"和"过载 rank o"，把 ``target - load(u)`` 个 token
  从 o 搬到 u。因为 ``load(o) > target >= deficit(u)``，o 永远拿得出足够
  token，u 一次就被填到 target、之后不再被碰。每轮填死一个 rank，至多
  ``R-1`` 轮结束；并且因为每个 rank 至多被填一次，它收来的全部远端 token
  来自**同一个**其他 rank，那个 rank 只 home ``E/R`` 个专家 —— 于是每 rank
  冗余副本 ≤ ``E/R``。

本实现在 CPU 上把"规划层"完整落地（plan_redundant_experts → RedundancyPlan
→ build_permutation / dispatch / combine），并用数值实验验证"恒定 S*K 缓冲
vs 最坏 S*K*R"以及 M(T)<=E/R 的界。真实 GPU 上的对称内存 / zero-copy
dispatch / prefetch 等内核级优化不属于本教学实现的范围（见 4.10 升华）。
"""

@dataclass
class RedundancyPlan:
    """一次规划的结果：每个 (rank, expert) 持有的 token 数。"""

    n_experts: int
    n_ranks: int
    experts_per_rank: int
    #: ``holdings[rank][expert] = 该 rank 要为 expert 计算的 token 数``。
    holdings: list[dict[int, int]] | None = None
    #: 每个 rank 计算的**非本地**（redundant）专家集合。
    redundant: list[set[int]] | None = None
    #: 规划器执行的"搬运"次数。
    moves: int = 0

    def __post_init__(self) -> None:
        if self.holdings is None:
            self.holdings = [{} for _ in range(self.n_ranks)]
        if self.redundant is None:
            self.redundant = [set() for _ in range(self.n_ranks)]

    def home_rank(self, expert: int) -> int:
        """专家 ``expert`` 的主家（home）rank：按连续切片分配。"""
        return expert // self.experts_per_rank

    def tokens_per_rank(self) -> list[int]:
        """每个 rank 实际计算的 token 数（应全部等于 target）。"""
        return [sum(h.values()) for h in self.holdings]

    def max_redundant(self) -> int:
        """所有 rank 中冗余副本数目的最大值（界要 <= E/R）。"""
        return max((len(r) for r in self.redundant), default=0)

    @property
    def capacity(self) -> int:
        """附录 E 的界：E/R 个冗余槽位/rank（向上取整）。"""
        return -(-self.n_experts // self.n_ranks)

    def is_balanced(self) -> bool:
        """所有 rank 的负载严格相同。"""
        counts = self.tokens_per_rank()
        return len(set(counts)) <= 1

    def respects_bound(self) -> bool:
        """冗余副本数不超过容量界。"""
        return self.max_redundant() <= self.capacity

    def summary(self) -> dict[str, float]:
        return {
            "balanced": float(self.is_balanced()),
            "max_redundant": float(self.max_redundant()),
            "capacity": float(self.capacity),
            "moves": float(self.moves),
        }


def plan_redundant_experts(
    loads: list[int] | Tensor,
    n_ranks: int,
    *,
    experts_per_rank: int | None = None,
) -> RedundancyPlan:
    """规划一个完美均衡的分配（附录 A 构造的直接翻译）。

    Args:
        loads: ``loads[e]`` = 路由到专家 ``e`` 的 token-专家指派个数。
        n_ranks: 专家并行度 R。
        experts_per_rank: 每 rank 本地专家数（缺省 = ceil(E/R)）。

    Returns:
        均衡的 RedundancyPlan。要求总指派数能被 R 整除
        （Top-K × S × R 的算术让 target 恰为 S*K）。

    算法（一步步验证界）：
      1. 初始：每个 rank 只持有它 home 专家的 token（``home = min(e // epr, R-1)``）。
      2. 主循环：
         - 找第一个欠载 rank ``u``（load(u) < target）；
         - 找当前最过载的 rank ``o``；
         - 从 o 的 holdings 中（从大专家开始，控制冗余副本数）取出
           ``need = target - load(u)`` 个 token 搬到 u；
         - o 一定是够拿的（load(o) > target >= need），所以 u 恰好落 target；
      3. 终止：某个 rank 一旦被填满再也不被触碰，至多 R-1 轮收敛。
    """
    if isinstance(loads, Tensor):
        loads = [int(x) for x in loads.tolist()]
    n_experts = len(loads)
    experts_per_rank = experts_per_rank or -(-n_experts // n_ranks)

    total = sum(loads)
    if total % n_ranks != 0:
        raise ValueError(
            f"总 token-专家指派数 {total} 不能被 {n_ranks} 个 rank 整除；"
            "MoonEP 的静态 shape 需要一个能整除的总量"
        )
    target = total // n_ranks

    plan = RedundancyPlan(n_experts, n_ranks, experts_per_rank)
    # 初始摆放：每家放自家专家的 token。
    for e, load in enumerate(loads):
        if load:
            home = min(e // experts_per_rank, n_ranks - 1)
            plan.holdings[home][e] = plan.holdings[home].get(e, 0) + load

    current = [sum(h.values()) for h in plan.holdings]

    while True:
        # 找到第一个欠载 rank——它的负载将被精确补齐。
        under = next((r for r in range(n_ranks) if current[r] < target), None)
        if under is None:
            break  # 都已均衡（含整除舍入内）
        # 找最过载的 rank 作为"供应者"。
        over = max(range(n_ranks), key=lambda r: current[r])
        if current[over] <= target:
            break  # 整数舍入内已均衡
        need = target - current[under]

        # 从 over 手里按"大专家优先"搬 need 个 token 给 under。
        # 大专家优先 => 尽可能少搬不同的专家 => 冗余副本数少。
        for expert in sorted(plan.holdings[over], key=lambda e: -plan.holdings[over][e]):
            if need == 0:
                break
            take = min(plan.holdings[over][expert], need)
            plan.holdings[over][expert] -= take
            if plan.holdings[over][expert] == 0:
                del plan.holdings[over][expert]
            plan.holdings[under][expert] = plan.holdings[under].get(expert, 0) + take
            # 非本地专家级别才计入冗余。
            if expert // experts_per_rank != under:
                plan.redundant[under].add(expert)
            need -= take

        current = [sum(h.values()) for h in plan.holdings]
        plan.moves += 1
        if plan.moves > n_ranks * n_ranks:  # 保险丝：理论上不会触发
            raise RuntimeError("MoonEP 规划器未收敛")

    return plan

def redundancy_lower_bound(loads: list[int], n_ranks: int, experts_per_rank: int) -> int:
    """冗余副本数的一个下界（廉价健全性检查）。

    只要有任一 rank 的 home 负载低于 target，它就必须至少算 1 个非本地专家，
    所以 max_redundant >= 1 一旦路由不均衡就成立。论文用整数线性规划求精确
    最优作参考，这里用下界防止规划器做匪夷所思的事。
    """
    target = sum(loads) // n_ranks
    home = [0] * n_ranks
    for e, load in enumerate(loads):
        home[min(e // experts_per_rank, n_ranks - 1)] += load
    return 1 if any(h < target for h in home) else 0

class MoonEP:
    """静态 shape 下的专家派发与汇总。

    plan -> build_permutation -> dispatch/combine 三段式。由于每个 rank 都拿到
    恰好 ``S*K`` 个 token，通信缓冲形状已知，permutation 一步到位（fuse 的
    permute/unpermute），grouped GEMM 的形状也能在 launch 前确定。
    """

    def __init__(self, n_experts: int, n_ranks: int):
        self.n_experts = n_experts
        self.n_ranks = n_ranks
        self.experts_per_rank = -(-n_experts // n_ranks)
        self.last_plan: RedundancyPlan | None = None

    # -- 规划 ----------------------------------------------------------------

    def plan(self, expert_ids: Tensor) -> RedundancyPlan:
        """从一维路由专家 id 张量规划。

        Args:
            expert_ids: ``(N*K,)`` 每 token 的 K 个路由专家 id（扁平）。

        Returns:
            均衡计划。单进程实现不跨 rank 收集，loads 直接用 bincount。
        """
        loads = torch.bincount(expert_ids.reshape(-1), minlength=self.n_experts)
        plan = plan_redundant_experts(loads, self.n_ranks, experts_per_rank=self.experts_per_rank)
        self.last_plan = plan
        return plan

    # -- 派发与合成 ----------------------------------------------------------

    def build_permutation(
        self, expert_ids: Tensor, plan: RedundancyPlan
    ) -> tuple[Tensor, Tensor, Tensor]:
        """返回 ``(order, destination_rank, expert_of_slot)``。

        - ``order``：把"扁平指派序列"重排为"按 rank / 按专家分组"的顺序，
          派发只需做一次 gather；
        - ``expert_of_slot``：目标槽位该算哪个专家，供 grouped GEMM 不看数据
          直接 launch。
        """
        flat = expert_ids.reshape(-1)

        # 每个 (rank, expert) 的剩余容量，直接来自 plan。
        remaining = [dict(h) for h in plan.holdings]
        dest = torch.empty_like(flat)
        cursor = {}
        for idx in range(flat.numel()):
            e = int(flat[idx])
            r = cursor.get(e)
            if r is None or remaining[r].get(e, 0) == 0:
                r = next(
                    (rr for rr in range(self.n_ranks) if remaining[rr].get(e, 0) > 0),
                    None,
                )
                if r is None:
                    raise RuntimeError(f"专家 {e} 已无容量")
                cursor[e] = r
            remaining[r][e] -= 1
            dest[idx] = r

        # 按 (rank, expert) 排序 => 每个 rank 的切面连续、段内按专家连续，
        # 恰好构成一个 grouped GEMM。
        key = dest.to(torch.int64) * self.n_experts + flat.to(torch.int64)
        # 老 torch（1.11）的 argsort 不支持 stable=，加微小扰动保证原序内的稳定。
        idx = torch.arange(key.numel(), device=key.device, dtype=torch.int64)
        order = torch.argsort(key * key.numel() + idx)
        return order, dest, flat[order]

    def dispatch(self, tokens: Tensor, order: Tensor) -> Tensor:
        """fused permute：一次 gather 进专家分组的槽位。"""
        return tokens.index_select(0, order)

    def combine(self, expert_out: Tensor, order: Tensor) -> Tensor:
        """fused unpermute：dispatch 的精确逆操作。"""
        out = torch.empty_like(expert_out)
        return out.index_copy_(0, order, expert_out)

    def static_shape(self, tokens_per_rank: int) -> tuple[int, ...]:
        """每个 rank 使用的计算 shape（launch 前已知）。"""
        return (tokens_per_rank,)

5.2 MoonEP-master 中的"决定拷贝哪些"与"实际拷贝"

我们主要参考 MoonEP-master 的 Python 代码。

这两件事被严格分离在两个文件里——这是 MoonEP 的核心设计：规划（planning）与执行（prefetch）解耦。

5.2.1 决定拷贝哪些专家：moonep/planning.py 的 PlanningKernel

生产实现是 CUDA kernel（planning.py:353 class PlanningKernel，入口 :1294 launch_planning），可读性差。但仓库贴心地提供了语义等价的 Python 参考实现，决策逻辑一目了然：

tests/planning_reference.py:149-166 —— 这就是"决定拷贝哪些"的地方：

for d in range(R):                                    # :149  遍历每个 rank
    local_start = d * epn
    local_end = local_start + epn

    # ① 筛出"该 rank 要算、但不是自己 home 的"专家 = 远端专家
    remote_experts = [
        e for e in range(E)
        if alloc[e, d].item() > 0 and not (local_start <= e < local_end)
    ]                                                 # :153-156

    # ② 按 token 数降序排（tie-break 用 expert id 保证确定性）
    remote_experts.sort(key=lambda e: (alloc[e, d].item(), e), reverse=True)   # :160

    # ③ 只取前 B 个，写进 experts_to_copy
    for b, e in enumerate(remote_experts[:B]):        # :162
        experts_to_copy[d, b] = e                     # :163  ★决策落盘
        remote_stats_all[e // epn, 1] += 1

决策规则三句话：

  1. 候选：alloc[e,d] > 0 且 e 不属于 rank d 的 home 区间 → 远端专家（:153-156）
  2. 排序：按该 rank 要为它算的 token 数降序（:160），注释 :158-159 说明动机："Pick the B remote expert segments with the most tokens for VM prefetch"
  3. 截断：只拷 前 B 个（:162，B = prefetch slot 数），其余远端专家走原路径

关键产物是 experts_to_copy，形状 [R, B]，初值 -1（:139），-1 表示该 slot 空置。

紧接着 :166-193 用这个决策重排物理 token 顺序：VM group id 0..E-1 是全局专家组，E..E+B-1 是被选中远端专家的 prefetch slot（注释 :169-171）；:177 if g not in experts_to_prefetch 确保被 prefetch 的专家不在原位重复占空间。

为什么按 token 数排序？因为拷权重的成本固定（H×H' 一份），而收益 ∝ 该 slot 要处理的 token 数。拷"最热"的远端专家，单位带宽换来的计算本地化收益最大。

5.2.2 实际拷贝专家权重：moonep/prefetch.py

文件头 :4 一句话点题：

"Copies selected remote expert weights from remote_expert[E, H, H'] into ..."

入口：prefetch.py:317 launch_prefetch(...)，签名第二个参数就是 experts_to_copy: torch.Tensor（:320），文档 :328 说明 "Entries are ... " ，并有断言 :340-341（必须 contiguous int32）、:351-352（长度必须等于 B）。注意这里是 [B]（单 rank 视角），planning 出的是 [R,B]，调用时切自己那一行。

实际搬运代码：PrefetchKernel.kernel（:138），分三段：

① 压缩稀疏计划（:207-214）

n_active = Int32(0)
for b in cutlass.range_constexpr(B):
    e = experts[b]                    # experts = experts_to_copy
    if e >= Int32(0):                 # :210  跳过 -1 空槽
        exp_tab[n_active] = e         # :211  源专家 id
        slot_tab[n_active] = Int32(b) # :212  目标 slot 号
        n_active += Int32(1)
total_tiles = n_active * TILES_PER_EXPERT   # :214

注释 :202-203 说明动机：压缩后 tile 循环只访问活跃 slot，不必在 B * TILES_PER_EXPERT 个索引里空转。

② 生产者 warp：远端 → shared memory（:216-237）

expert = exp_tab[i]                   # :226  查出源专家
src_mt = expert * MTILES + mt         # :229  ★源地址由 expert id 算出
cute.copy(tma_g2s, tSgS[(None,(src_mt,nt))], tSsS[...], tma_bar_ptr=...)  # :230-235

③ 消费者 warp：shared memory → 本地 prefetch buffer（:239-277）

b = slot_tab[i]                       # :253  查出目标 slot
dst_mt = b * MTILES + mt              # :256  ★目标地址由 slot 号算出
cute.copy(tma_s2g, tDsD[...], tDgD[(None,(dst_mt,nt))])                   # :257-261

这就是真正的拷贝：TMA 双跳 G2S→S2G，按 (M_BLOCK, N_BLOCK) 分块流水（stages 级 pipeline，:194-200），BF16（:167）。:272-275 的 issued >= 2 早释放策略还专门注释了为什么不滞后 stages-1 组（避免 producer 只剩 ~1 个 in-flight load，把带宽打成 tile_bytes / remote_latency）。

5.2.3 一张表总结 MoonEP-master

关注点 位置 说明
决定拷哪些（生产） moonep/planning.py:353 PlanningKernel，入口 :1294 launch_planning CUDA 实现
决定拷哪些（可读参考） tests/planning_reference.py:149-166 ★候选 :153-156 / 排序 :160 / 截断+落盘 :162-163
决策产物 experts_to_copy [R, B]，初值 -1（:139） -1 = 空槽
实际拷贝（入口） moonep/prefetch.py:317 launch_prefetch 消费 experts_to_copy
实际拷贝（kernel） prefetch.py:138 kernel 压缩 :207-214、G2S :230、S2G :257
源/目标寻址 :229 src_mt = expert*MTILES+mt / :256 dst_mt = b*MTILES+mt expert id → slot 号的地址映射

5.3 本项目 src/parallel/moonep.py 中的对应位置

5.3.1 本项目只实现了规划层，没有实现权重拷贝

这一点文件 docstring 已明确声明（moonep.py:34-37）：本实现在 CPU 上把"规划层"完整落地（plan_redundant_experts → RedundancyPlan → build_permutation / dispatch / combine），并用数值实验验证"恒定 S*K 缓冲 vs 最坏 S*K*R"以及 M(T)<=E/R 的界。真实 GPU 上的对称内存 / zero-copy dispatch / prefetch 等内核级优化不属于本教学实现的范围。

所以问题二要分两半回答。

5.3.2 决定拷贝哪些专家：plan_redundant_experts:170-183

对应 MoonEP-master 的 planning 层，但算法不同（见下节）：

# 从 over 手里按"大专家优先"搬 need 个 token 给 under。
# 大专家优先 => 尽可能少搬不同的专家 => 冗余副本数少。          :170-171
for expert in sorted(plan.holdings[over], key=lambda e: -plan.holdings[over][e]):  # :172 ★选择
    if need == 0:
        break
    take = min(plan.holdings[over][expert], need)      # :175
    plan.holdings[over][expert] -= take                # :176  源减
    if plan.holdings[over][expert] == 0:
        del plan.holdings[over][expert]                # :177-178
    plan.holdings[under][expert] = plan.holdings[under].get(expert, 0) + take  # :179 目标加
    # 非本地专家级别才计入冗余。
    if expert // experts_per_rank != under:            # :181  ★判定是否为冗余副本
        plan.redundant[under].add(expert)              # :182  ★决策落盘
    need -= take                                       # :183

决策落盘处：plan.redundant[under].add(expert)（:182）——这个集合就是本项目版本的 experts_to_copy。数据结构定义在 :64-65：

#: 每个 rank 计算的**非本地**（redundant）专家集合。
redundant: list[set[int]] | None = None

选择规则：:172 的 sorted(..., key=lambda e: -holdings[over][e]) —— 大专家优先，注释 :171 给了理由："大专家优先 => 尽可能少搬不同的专家 => 冗余副本数少"。

谁给谁搬由主循环决定（:159-168）：

under = next((r for r in range(n_ranks) if current[r] < target), None)  # :161 欠载 rank
over  = max(range(n_ranks), key=lambda r: current[r])                   # :165 最过载 rank
need  = target - current[under]                                         # :168 精确补齐量

5.3.3 "实际拷贝"的对应物：只有 token 搬运，没有权重拷贝

本项目里唯一能称为"拷贝"的是 token 数的账面转移（:176-179 那四行加减法）以及数据的 gather/scatter：

def dispatch(self, tokens, order):                     # :277
    return tokens.index_select(0, order)               # :279  fused permute

def combine(self, expert_out, order):                  # :281
    out = torch.empty_like(expert_out)
    return out.index_copy_(0, order, expert_out)       # :284  fused unpermute

但请注意：这是拷贝 token（激活），不是拷贝专家权重。MoonEP-master 的 prefetch.py 拷的是 remote_expert[E, H, H'] 这个权重张量。本项目没有任何对应代码——因为它是单进程 CPU 教学实现，所有专家权重本来就在同一块内存里，不存在"远端权重"这个概念。

build_permutation（:239-275）里的 remaining/cursor 逻辑（:252-267）是在为每个路由指派挑一个有余量的 rank，然后 :271-274 按 (rank, expert) 排序让每个 rank 的切面连续、段内按专家连续——"恰好构成一个 grouped GEMM"（:269-270）。这对应 MoonEP-master 中 planning 的第二部分（planning_reference.py:166-199 那段重排物理 token 顺序、算 cu_seqlens 的逻辑），而不是 prefetch。

5.3.4 本项目对应位置总表

关注点 位置 说明
决定拷贝哪些（主循环选 under/over） moonep.py:159-168 :161 欠载 rank、:165 最过载 rank、:168 精确补齐量
决定拷贝哪些（选具体专家） moonep.py:172 大专家优先排序，理由见注释 :170-171
决策落盘 moonep.py:181-182 plan.redundant[under].add(expert)，判定条件 expert // epr != under
决策产物数据结构 moonep.py:64-65 redundant: list[set[int]] 相当于 master 的 experts_to_copy
token 账面转移 moonep.py:176-179 源减 / 目标加，四行加减法
激活数据搬运 moonep.py:277-284 dispatch/combine index_select / index_copy_
专家权重拷贝 不存在 见 docstring :34-37 显式声明超出范围

5.3.5 两侧的实质差异（不只是"简化"，而是解决不同问题）

这是最容易误读的地方。表面看两边都在算"哪些专家要在非 home rank 上算"，但优化目标完全不同：

维度 MoonEP-master 本项目 moonep.py
算法目标 给定已有的 alloc[e,d]，挑 B 个最热的远端专家做权重预取 从零构造 holdings，使每 rank token 数严格相等
alloc 的来源 输入（launch_planning 的上游已算好容量分配） 输出（plan_redundant_experts 自己构造）
筛选依据 token 数降序取前 B（:160,:162）——B 是硬件 slot 数 token 数降序尽量少搬（:172）——为压低冗余副本数
约束 B = prefetch buffer 槽位数（硬件预算） E/R = 附录 E 定理的理论界（:88-90 capacity）
未入选的远端专家 走原路径（token 跨 rank 传输，不预取权重） 不存在"未入选"——规划保证总能在界内完成
拷什么 权重 remote_expert[E,H,H']（prefetch.py:4） token/激活（dispatch/combine）
验证方式 tests/test_prefetch.py、test_planning.py is_balanced() :92-95、respects_bound() :97-99

MoonEP-master 的 planning 回答"哪些远端专家的权重值得预取"（带宽换计算本地化，B 是硬件预算）；本项目的 plan_redundant_experts 回答"如何复制专家使负载严格均衡"（复制换静态 shape，E/R 是理论界）。前者是执行层优化，后者是论文定理的构造性证明。

本项目 docstring :21-32 把这个定理和它的证明都写清楚了：

M(I) <= E/R 恒成立，且界基本是紧的 …… 因为每个 rank 至多被填一次，它收来的全部远端 token 来自同一个其他 rank，那个 rank 只 home E/R 个专家 —— 于是每 rank 冗余副本 ≤ E/R。

这个"每 rank 至多被填一次"的性质正是 :161 用 next(...) 找第一个欠载 rank、且 :168 一次性精确补到 target 所保证的——填满后 current[u] == target，此后 :161 再也不会选中它。:187-188 的 moves > R*R 保险丝注释也写了"理论上不会触发"。

5.3.6 如果要在本项目补上 prefetch 层

按 master 的结构，缺的是这样一段：

plan.redundant[rank]  ──►  按 holdings[rank][e] 降序取前 B 个  ──►  experts_to_copy[rank]
                                                                        │
                                        expert_weights[e] ──copy──► prefetch_buffer[b]

即：在现有 RedundancyPlan 上加一个 select_prefetch(rank, B) 方法（对应 planning_reference.py:160-163），再加一个按 slot 搬权重的函数（对应 prefetch.py:230/:257 的两跳，CPU 版就是 buffer[b].copy_(expert_w[e])）。这样 before/after 对比就能展示"预取热专家权重 vs 跨 rank 传 token"的带宽差异。

5.4 MoonEP 如何解决静态 shape 问题？

核心思路：用"专家复制"换取"每 rank 恒定负载"。

具体机制（见 src/models/llm/kimi_k3/parallel/moonep.py:1-37）：

  1. 允许冗余专家副本：每个 rank 除了持有自己的 E/R 个本地专家，还预留最多 ⌈E/R⌉ 个冗余槽位（附录 E 定理保证这个界永远够用）；
  2. 构造性规划（plan_redundant_experts，moonep.py:110-190）：反复找"欠载 rank u" + "最过载 rank o"，从 o 搬 target - load(u) 个 token 给 u，每轮填死一个 rank，至多 R-1 轮收敛；
  3. 结果：每个 rank 拿到恰好 S·K 个 token——不多不少。

由此带来的静态 shape 三大好处：

  • 通信缓冲恒定 = S·K，而非最坏 S·K·R（见 fig_buffer_size 图）；
  • grouped GEMM 的 shape 在 launch 前已知（每个 rank 知道每专家多少 token），无需 host-device 同步；
  • 无 hot rank，负载零波动。

也就是说，MoonEP 不是预测 shape，而是用复制专家的余量空间，把任意不均衡的路由结果"摊平"到恒定 shape——这是把动态问题转化为静态问题的构造性解法。

TransFormer-封面

0xFF 参考

posted @ 2026-11-15 15:08  罗西的思考  阅读(6)  评论(0)    收藏  举报