[Docs] Explain Kimi-K3 EP32 and EP4 execution

2026-08-21 15:19:12 +08:00 · 2026-08-21 15:19:12 +08:00 · d790df39b2
commit d790df39b2
parent 96ffea6d37
1 changed files with 390 additions and 0 deletions
--- a/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/EP32_VS_EP4.md
+++ b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/EP32_VS_EP4.md
@ -0,0 +1,390 @@
+# Kimi-K3 EP32 与 EP4：实现方式及 TTFT 差异分析
+
+## 1. 先给结论
+
+在当前四机 32 卡、`TP=32`、`DP=1`、`moe_a2a_backend=none` 的 Kimi-K3 部署里：
+
+- `EP32` 不是 32 份模型副本，而是 896 个专家分散到 32 张卡；每张卡持有 28 个完整专家。
+- `EP4` 也不是只用 4 张卡跑 MoE，而是把 32 张卡分解成 `EP4 × MoE-TP8`。每个节点持有 224 个专家，每个专家的中间维再由节点内 8 张卡切分。
+- 两种布局每卡承载的专家参数量和理论 MoE FLOPs 接近，但计算颗粒度不同。
+- 当前没有启用 A2A。每张卡仍看到同一批 token，只计算自己负责的专家或专家分片，随后在 TP32 上汇总结果。
+- 同后端 Marlin、同请求口径下，EP4 相比 EP32 的 Input TPS 提高约 **15.9%**，TTFT P50 降低约 **13.9%**。
+- 这次收益不是“把 AllReduce 从 32 卡缩成 4 卡”。Kimi-K3 当前 `A2A=none` 路径仍在完整 TP32 组上归约。更合理的解释是：EP4 把每个专家的计算切到 8 卡，改善了路由偏斜、热点专家关键路径和 grouped MoE kernel 的并行形状。
+
+## 2. 审计对象
+
+### 2.1 模型参数
+
+来源：`/data/hf_models/Kimi-K3/config.json` 的 `text_config`。
+
+| 参数 | 数值 |
+|---|---:|
+| hidden size | 7168 |
+| Latent MoE hidden size | 3584 |
+| 单专家 intermediate size | 3072 |
+| routed experts | 896 |
+| 每 token 激活专家数 | 16 |
+| Transformer 层数 | 93 |
+
+### 2.2 源码版本
+
+审计仓库：
+
+```text
+/data/hzy/src/sglang-kimi-sm120
+commit fb929bbccb7a34640f3d4904174767eb7ae95509
+```
+
+实验镜像：
+
+```text
+local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5
+```
+
+EP32 和 EP4 的 Marlin 实验使用同一镜像、同一模型、同一请求、同一 chunk 和同一服务参数。启动命令中唯一关键差异是：
+
+```text
+EP32: --tp-size 32 --ep-size 32
+EP4 : --tp-size 32 --ep-size 4
+```
+
+## 3. EP、MoE-TP 和外层 TP 的关系
+
+SGLang 不是把 `TP` 和 `EP` 相乘申请 GPU，而是在外层 TP 组内部进一步分解 MoE：
+
+```python
+moe_tp_size = tensor_model_parallel_size // moe_ep_size // moe_dp_size
+```
+
+源码：
+
+```text
+python/sglang/srt/distributed/parallel_state.py:2526-2528
+python/sglang/srt/model_executor/model_runner_components/moe_ep_setup.py:123-137
+```
+
+本实验 `TP=32`、`MoE-DP=1`，因此：
+
+```text
+TP32 = EP32 × MoE-TP1
+TP32 = EP4  × MoE-TP8
+```
+
+| 配置 | EP size | 派生 MoE-TP size | 含义 |
+|---|---:|---:|---|
+| EP32 | 32 | 1 | 每个专家只落在 1 张 GPU 上 |
+| EP4 | 4 | 8 | 每个专家落在一个 8 GPU MoE-TP 组上 |
+
+所以 `EP4` 的准确叫法是“较小 EP + 较大 MoE-TP”的混合专家布局。
+
+## 4. 32 张卡具体怎么分组
+
+假设全局 rank `0-7 / 8-15 / 16-23 / 24-31` 分别属于四个节点。
+
+### 4.1 EP32
+
+```text
+MoE-EP group: [0, 1, 2, ..., 31]
+MoE-TP group: 每个 rank 自己，size=1
+```
+
+896 个专家平均切到 32 个 EP rank：
+
+```text
+每 rank 专家数 = 896 / 32 = 28
+```
+
+例如：
+
+```text
+rank 0 负责专家   0-27
+rank 1 负责专家  28-55
+...
+rank 31 负责专家 868-895
+```
+
+每个专家在其所属 GPU 上保留完整的 intermediate `3072`。
+
+### 4.2 EP4
+
+SGLang 构造的 MoE-TP 组是连续 8 rank：
+
+```text
+[0-7], [8-15], [16-23], [24-31]
+```
+
+每组正好对应一个 8 卡节点。MoE-EP 组则按 MoE-TP rank 跨节点取 rank：
+
+```text
+[0,8,16,24], [1,9,17,25], ..., [7,15,23,31]
+```
+
+源码：
+
+```text
+python/sglang/srt/distributed/parallel_state.py:2559-2615
+```
+
+896 个专家先切为 4 份：
+
+```text
+每个 EP rank 专家数 = 896 / 4 = 224
+```
+
+每个专家的 intermediate 再切给 MoE-TP8：
+
+```text
+每卡专家分片宽度 = 3072 / 8 = 384
+```
+
+因此：
+
+```text
+节点 0 的 8 张卡共同持有专家   0-223
+节点 1 的 8 张卡共同持有专家 224-447
+节点 2 的 8 张卡共同持有专家 448-671
+节点 3 的 8 张卡共同持有专家 672-895
+```
+
+节点内 8 张卡不是各存一份 224 个完整专家，而是各存这些专家的 `1/8` intermediate 分片。
+
+## 5. 为什么 EP4 没有让每卡专家权重暴涨 8 倍
+
+SGLang 先按 EP 切专家数量，再按 MoE-TP 切专家中间维：
+
+```python
+self._num_local_routed = self._num_global_routed // storage_ep_size
+self.intermediate_size_per_partition = intermediate_size // self.moe_tp_size
+```
+
+源码：
+
+```text
+python/sglang/srt/layers/moe/fused_moe_triton/layer.py:291-309
+```
+
+权重加载时，`w1/w3` 沿输出维切分，`w2` 沿输入维切分：
+
+```text
+python/sglang/srt/layers/moe/fused_moe_triton/layer.py:618-682
+python/sglang/srt/layers/moe/fused_moe_triton/layer.py:702-770
+```
+
+只看与 intermediate 成正比的部分，每卡工作量满足：
+
+```text
+EP32: 28  × 3072 = 86016
+EP4 : 224 × 384  = 86016
+```
+
+真实服务也印证了每卡权重显存接近：
+
+| 配置 | 后端 | 每卡 weight memory |
+|---|---|---:|
+| EP32 | Marlin | 59.63 GiB |
+| EP4 | Marlin | 59.28 GiB |
+| EP4 | FlashInfer MXFP4 | 59.65 GiB |
+
+小幅差异来自权重布局、padding、workspace 和后端转换，不是模型参数量发生了数量级变化。
+
+## 6. 当前 `A2A=none` 时，一个 token 怎么执行
+
+### 6.1 路由
+
+Kimi-K3 的 gate 为每个 token 计算 896 个专家分数，再选 top-16：
+
+```text
+hidden states
+  -> gate / grouped top-k
+  -> 16 个全局 expert id + routing weight
+```
+
+模型入口：
+
+```text
+python/sglang/srt/models/kimi_k3.py:386-469
+python/sglang/srt/models/kimi_k3.py:1020-1059
+```
+
+### 6.2 没有 token A2A 搬运
+
+`moe_a2a_backend=none` 会创建 `StandardDispatcher`：
+
+```text
+python/sglang/srt/layers/moe/fused_moe_triton/layer.py:137-178
+```
+
+StandardDispatcher 不会把 token 发到专家所在 GPU。它保留 hidden states，只处理专家 ID：
+
+```python
+hidden_states = hidden_states
+```
+
+Marlin 路径把不属于本 EP rank 的全局专家映射为 `-1`，本地 kernel 只计算自己负责的专家：
+
+```text
+python/sglang/srt/layers/moe/token_dispatcher/standard.py:182-246
+python/sglang/srt/layers/moe/fused_moe_triton/layer.py:903-911
+```
+
+FlashInfer 路径不在 dispatcher 中改 ID，而是把 `moe_ep_size/rank` 与 `moe_tp_size/rank` 传给 CUTLASS runner，由 runner 解释全局专家 ID：
+
+```text
+python/sglang/srt/layers/quantization/mxfp4.py:1378-1401
+```
+
+两者实现位置不同，但语义一致：没有 A2A token dispatch，每个 rank 从同一批 token 中计算自己的专家贡献。
+
+### 6.3 汇总部分结果
+
+Kimi-K3 的 routed expert 输出位于 3584 维 latent 空间。A2A 未启用时，源码要求在 RMSNorm 前完成归约：
+
+```python
+def _routed_needs_reduce(self):
+    return self.tp_size > 1 and get_moe_a2a_backend().is_none()
+
+return self._latent_norm(tensor_model_parallel_all_reduce(latent))
+```
+
+源码：
+
+```text
+python/sglang/srt/models/kimi_k3.py:699-702
+python/sglang/srt/models/kimi_k3.py:955-960
+```
+
+这里调用的是完整的 `get_tp_group()`，即 TP32。某些优化路径会把这次归约与相邻算子融合，但数学上的通信域仍是 TP32。
+
+因此，本次 EP4 收益不能解释成“MoE AllReduce 从 32 rank 缩成 4 rank”。
+
+## 7. EP4 为什么更快
+
+### 7.1 同口径实验结果
+
+唯一能隔离 EP 变量的比较是 Marlin 对 Marlin：
+
+```text
+模型 Kimi-K3
+4 节点 / 32 GPU
+TP32 / PP1 / DP1
+ISL=16K / OSL=1 / C=8
+chunked_prefill_size=8K
+moe_a2a_backend=none
+moe_runner_backend=marlin
+```
+
+| 配置 | repeats | Input TPS | TTFT P50 | TTFT P95 |
+|---|---:|---:|---:|---:|
+| EP32 | 2 | 2531.43 | 50.52 s | 53.64 s |
+| EP4 | 3 | 2935.02 | 43.51 s | 46.22 s |
+| EP4 vs EP32 | | **+15.9%** | **-13.9%** | **-13.8%** |
+
+结果路径：
+
+```text
+EP32:
+/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/results/kimi3-moe-prefill-20260818-130900/raw/
+
+EP4:
+/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/results/kimi3-ep4-moe-full-20260818-151349/raw/
+```
+
+EP4 + FlashInfer 的同场景结果为 3257.96 Input TPS、39.19 s TTFT P50。它相对 EP4 + Marlin 还包含约 11% 的 MoE backend 收益，不能全部归因于 EP。
+
+### 7.2 主要机制：把单专家关键路径从 1 卡摊到 8 卡
+
+Kimi-K3 每个 token 激活 16 个专家。路由并不保证所有专家收到完全一样数量的 token，总会出现热点专家和冷门专家。
+
+EP32 中：
+
+```text
+一个专家只属于一张 GPU
+热点专家的完整 3584 -> 3072 -> 3584 计算由这张 GPU 独自完成
+其他 GPU 即使更空，也不能帮助它完成该专家
+```
+
+EP4 中：
+
+```text
+一个专家属于一个节点的 8 卡 MoE-TP 组
+同一专家的 intermediate 3072 被切成 8 份，每卡计算 384
+热点专家的矩阵计算由 8 张 GPU 共同承担
+```
+
+这不会减少总 FLOPs，但会降低单个热门专家压在单卡上的关键路径，并让节点内 8 卡对同一专家协同计算。
+
+### 7.3 grouped MoE kernel 的执行形状改变
+
+每卡的理论 FLOPs 接近，但 kernel 看到的形状不同：
+
+| 配置 | 本地专家槽位 | 单专家 intermediate 分片 | 平均每 token 本 rank 参与的专家分片 |
+|---|---:|---:|---:|
+| EP32 | 28 | 3072 | 16/32 = 0.5 |
+| EP4 | 224 | 384 | 16/4 = 4 |
+
+EP4 让更多专家分片同时出现在一个 rank 的 grouped MoE 调度中，并把单个专家的大矩阵切成更多可并行的小分片。在本次 SM120 + Marlin 实测中，这个形状明显更合适。
+
+这属于“同样总计算，不同并行颗粒度”的收益，不是模型数学变化。
+
+### 7.4 通信没有消失，但也没有因 MoE-TP8 再增加一次完整归约
+
+当前 Kimi `A2A=none` 路径最终直接在 TP32 汇总 latent 输出。这一次 TP32 AllReduce 同时合并了：
+
+- 不同 EP rank 持有的不同专家贡献；
+- EP4 中同一专家的 8 个 MoE-TP 分片贡献。
+
+所以 EP4 没有额外增加一轮独立的“EP4 AllReduce + TP8 AllReduce”。这使它可以获得更好的计算分摊，而不在当前实现中付出第二次完整 collective 的代价。
+
+## 8. EP4 与真正 A2A EP 的区别
+
+当前实验：
+
+```text
+--ep-size 4
+--moe-a2a-backend none
+```
+
+它只是改变专家权重的 EP/MoE-TP 布局，不会按路由结果搬 token。
+
+真正的 DeepEP 或 FlashInfer A2A 会执行：
+
+```text
+dispatch: token 按 top-k 发往专家所在 rank
+compute : rank 只计算收到的 token
+combine : 专家结果返回 token 原所属 rank 并加权合并
+```
+
+SGLang 的 dispatcher 选择源码：
+
+```text
+Standard/A2A 分支：python/sglang/srt/layers/moe/fused_moe_triton/layer.py:137-178
+DeepEP dispatcher：python/sglang/srt/layers/moe/token_dispatcher/
+```
+
+Kimi-K3 源码也明确说明，A2A 后 combine 已返回完整 routed sum，因此不能再做 TP AllReduce：
+
+```text
+python/sglang/srt/models/kimi_k3.py:506-516
+python/sglang/srt/models/kimi_k3.py:699-709
+```
+
+所以“EP4 比 EP32 快”不能直接推出“DeepEP EP4 一定更快”。A2A 会引入完全不同的 token dispatch/combine、buffer 和共享专家约束，需要单独实验。
+
+## 9. 证据边界
+
+已经由源码和实验确认：
+
+1. `TP32/EP32` 派生 `MoE-TP1`，`TP32/EP4` 派生 `MoE-TP8`。
+2. EP32 每卡 28 个完整专家；EP4 每个节点 224 个专家、每卡持有其 1/8 intermediate 分片。
+3. 两种 Marlin 实验只改变 `--ep-size`，EP4 的 TTFT P50 低 13.9%。
+4. 当前 `A2A=none` 不搬 token，Kimi routed latent 仍在完整 TP32 上归约。
+5. 每卡权重显存基本相同，收益不是因为 EP4 少加载了大量权重。
+
+目前还没有由成对 Nsight Trace 直接确认：
+
+1. 13.9% TTFT 收益中，Marlin grouped MoE kernel 本身占多少；
+2. 路由负载偏斜改善占多少；
+3. padding、排序和 expert mapping 开销分别变化多少。
+
+若需要把原因从“源码约束下的机制解释”升级为“算子级定量归因”，下一步只需补一组相同 16K、C8、chunk8K 的 EP32/EP4 Nsight 对照，比较每层 Marlin MoE kernel、top-k/sort 与 TP32 AllReduce 时间。现有结果已经足够说明 EP4 是当前部署更好的基线，但不应虚构尚未采集的算子时间占比。
+