追踪一个 Token 在 MoE 推理中穿越 TP/CP/SP/EP/DP/PP 的路径

目录

• 六种切分
• The Setup: One Model, 32 GPUs, Two Meshes
  • 认识我们的 Token
• 第 1 跳：落到 GPU 11
• 第 2 跳：与 GPU 10 做 All-Gather
• 第 3 跳：与 GPU 9 交换 K/V
• 第 4 跳：分发到两个专家
• 一瞥单层
• 第 5 跳：PP 发送到 GPU 27
• 第 6 跳：把 K/V 留在 GPU 10、11、26、27 上
• 完整旅程：地图与流量
• Token 的来世
• 三个答案

让我们追踪一个提示词 token，看它如何在一个混合专家（MoE）模型于 32 块 GPU 上的 prefill 过程中穿行，其间涉及六种并行：数据并行（DP）、上下文并行（CP）、序列并行（SP）、张量并行（TP）、专家并行（EP）和流水线并行（PP）。Prefill 是在生成第一个输出 token 之前，对提示词执行的前向传播，用于填充 KV 缓存。

每一种单独看都很简单。难的是它们如何组合：哪些 GPU 构成哪个组，布局为何按这样的顺序排列，以及为什么一旦加入专家，数据并行副本就不再彼此独立。追踪一个 token 就能回答这三个问题。

没有人会同时部署全部六种；在 4 节点 × 8 GPU 上，DeepSeek 用 attention DP32 + EP32 做 prefill。这个网格把每个轴都放进了一条轨迹里。

上面的地图就是路线；其中六个编号步骤就是下面的六跳。Prefill 会丢弃它为这个 token 计算的一切，只保留其 K/V。

六种切分

每种并行都切掉工作的一个轴。只要后续步骤需要被切到另一块 GPU 上的东西，它就要付出通信代价。

All-gather：每块 GPU 都得到每个切片。Reduce-scatter：跨 GPU 求和，每块 GPU 保留自己的切片。All-to-all：每块 GPU 向其他每块 GPU 发送不同的切片。

SP 和 CP 都切分 token，但 SP 的切分（在一个 TP 对内）会在注意力前被 all-gather 撤销；CP 的切分永远不会（CP rank 之间交换 K/V）。这里的 SP 是 Megatron 风格；DeepSpeed-Ulysses 和若干 ring-attention 论文用这个名字指把一个长序列切分到多块 GPU 上，这更接近 CP。

配置：一个模型，32 块 GPU，两个网格

MoE 推理在同一批 GPU 上铺两个网格（Megatron 的 MoE Parallel Folding）：注意力使用 TP2 × CP2 × DP4 × PP2，专家使用 EP16 × PP2。DP 复制的是注意力，不是专家：四个副本的 16 块 GPU 共享一份 32 个专家的副本，每块 2 个，因此每个专家都能看到每个副本的 token。

编号：rank = tp + 2·cp + 4·dp + 16·pp。阶段 0 是 GPU 0–15（节点 0–1，层 1–4）；阶段 1 是 GPU 16–31（节点 2–3，层 5–8）。

为什么是这个顺序：TP 和 CP 伙伴在每一层都要通信，所以它们共享一个节点的 NVLink。DP 副本共享一个 EP 组，因此它们填满一个 stage。PP 只在每个 stage 边界通信一次，所以它位于最外层，跨 IB。EP 有意跨节点：跨越一个 stage 的两个节点后，每个 GPU 只留 2 个专家而不是 4 个（每个节点一份副本），从而为 KV cache 腾出内存。

认识我们的 token

我们的 token 位于 prompt 第三季度的位置 t，驻留在 GPU 11 上。

一份 行程日志 统计持有该 token 任意片段的 GPU、跨节点次数（其数据的 IB 传输）以及通信操作。

第 1 跳发生一次；第 2–4 跳是一层，在所有 8 层中重复；第 5 跳跨到 stage 1；第 6 跳结束 prefill。

第 1 跳：落在 GPU 11 上

1. DP：负载均衡器将请求发送到副本 2。
2. CP：zigzag 将 prompt 切成 4 个块，并把 CP rank r 的块 r 和 3 − r 分给它。
3. SP：TP rank 0 取块 r，TP rank 1 取块 3 − r。

这就是 SP 布局：每个 GPU 拥有 T/4 个完整 token。

Zigzag 平衡了因果掩码（Llama 3，§3.3.2）。在半块块中，块 i 关注 i 个完整的更早块加上自身的一半，因此各块代价为 1、3、5、7，每个 rank 得到 8。有 2·CP 个块时，rank r 得到 r 和 2·CP − 1 − r。

嵌入查找是局部的：这个玩具给每个 stage-0 GPU 完整的表，尽管真实引擎会像 LM head（第 6 跳）那样按 TP 对在词表维度上拆分它。

陷阱：RoPE（旋转位置嵌入）和因果掩码必须使用全局位置 t，而不是 GPU 11 上的局部索引。

第 2 跳：与 GPU 10 进行 All-Gather

• RMSNorm：按 token 进行，所以 GPU 11 独自运行它。
• All-gather：GPU 10 和 11 交换它们的 T/4 个 token；注意力需要全部 T/2。
• 列并行 QKV：每个 GPU 持有半数头的 QKV 权重列，并对所有 T/2 个 token 做投影。我们 token 的前半头落在 GPU 10 上，后半头落在 GPU 11 上。
• KV cache 写入：我们 token 的 K/V 以相同方式拆分，每个 GPU n_kv/2 个 KV 头。

约束：TP 必须整除 n_h。它还必须整除 n_kv，除非 n_kv < TP；那么 TP 必须是 n_kv 的倍数，并且每个 KV 头及其 cache 被复制 TP/n_kv 次。这里两者都是偶数，所以没有复制。

行程日志：2/32 GPU（+10）· 0 次跨节点 · 1 次通信操作（+all-gather）

第 3 跳：与 GPU 9 交换 K/V

GPU 10 和 11 持有块 1 和 2，但我们的 query 还需要块 0。GPU 8 和 9 在它们的 all-gather 之后持有块 0 和 3（在此之前，块 0 只在 GPU 8 上）。

Ring attention 保持 query 不动并移动 K/V。在 CP = 2 时，环是相同头伙伴之间的一次交换，GPU 11 ↔ GPU 9 和 GPU 10 ↔ GPU 8，等同于 Llama 3 的 K/V all-gather；ring 仅在更大的 CP 时不同。我们的 K/V 前往 GPU 8 和 9 以服务块 3，其 query 全部在它之后。

GPU 11 现在为我们的 query 持有两个部分输出：$O_1$ 基于其本地 K/V，$O_2$ 基于伙伴的块 0。每个都已经除以其 softmax 和 $\ell_i = \sum_j e^{s_j - m_i}$ 并携带其最大分数 $m_i$，因此 GPU 11 在本地合并它们：

\[\begin{aligned} O &= \frac{w_1 O_1 + w_2 O_2}{w_1 + w_2} \\ w_i &= e^{m_i - m}\,\ell_i, \quad m = \max(m_1, m_2) \end{aligned}\]

• GQA 缩小了交换量：K 和 V 每 GPU 覆盖 n_kv/2 个头，而不是 n_h/2。
• Reduce-scatter：行并行的输出投影（每个 GPU 持有其头对应的行）让每个 GPU 只得到总和的 [T/2, d] 份额；reduce-scatter 将它们求和，并把各自的 T/4 个 token 交还给每个 GPU，于是我们的 token 在 GPU 11 上又完整了。

添加 SP 不会给 TP 增加额外字节（Korthikanti et al.）。普通 TP 会以 all-reduce 结束注意力，让每个 token 都留在两个 GPU 上。环形 all-reduce 是先 reduce-scatter 再 all-gather；SP 在两者之间执行逐 token 的工作。

行程日志：4/32 GPU（+8、9）· 0 次跨节点 · 3 次通信操作（+K/V 交换、reduce-scatter）

第 4 跳：分发到两个专家

reduce-scatter 之后，每个 token 只存在于一个 GPU 上，并且只被分发一次，而不是两次：这就是 SP 在 MoE 推理中的真正作用。vLLM 的 sequence-parallel MoE 是一个更简单的版本：它保留 all-reduce，只在分发前拆分 token。

1. 路由：GPU 11 的路由器选择专家 5 和 27（示例），权重为 $p_5$ 和 $p_{27}$。
2. 分发：在 EP 组上进行 all-to-all。专家 27 在 GPU 13 上，相距一跳 NVLink。专家 5 在 GPU 2 上，即节点 0：该 token 跨 IB 一次，到达 GPU 3，即中继，它拥有 GPU 11 的本地索引，因此也使用相同的 IB 通道；GPU 3 通过 NVLink 转发它。DeepSeek-V3 就是这样路由其 all-to-all 的，DeepEP 用于预填充的高吞吐内核也是如此。
3. 合并：GPU 2 和 13 运行各自的专家，并通过反向路径返回结果，每个中继先对其节点内加权输出求和，这样只有一份副本跨 IB。GPU 11 将各部分相加：

\[y = p_5\,E_5(x) + p_{27}\,E_{27}(x)\]

这就是 DP 副本不再彼此独立的地方：每个副本有自己的注意力，但四者共享一份专家副本，因此每个 MoE 层都是一次会合。空闲副本仍会运行一个虚拟批次，以免其他副本挂起。最慢的副本决定全部 16 个 GPU 的节奏，因此负载均衡器必须把注意力工作均匀分配。

行程日志：7/32 GPU（+13、3、2）· 2 次跨节点（+2）· 5 次通信操作（+分发、合并）

一层概览

每一层都执行虚线上方的步骤：

• 两种布局交替：在注意力之外，SP 布局让每个 GPU 获得 T/(CP·TP) 个完整 token；在注意力之内，每个 TP rank 看到 T/CP 个 token 和 1/TP 个头。
• TP 的 all-gather 止步于 [T/CP, d]：CP 交换 K/V 而不是收集 token，因此没有注意力步骤会具体化 [T, d]。
• 每层五次通信操作：all-gather、K/V 交换、reduce-scatter、分发、合并。

第 5 跳：PP 发送到 GPU 27

• 第 2–4 层：专家位于 GPU 11（自身）和 4（节点 0，经由中继 GPU 3：每个方向各跨 IB 一次）；然后是 15 和 8，都在节点 1 上，因此不跨节点（GPU 8 已经持有我们的 K/V）；然后是 1 和 6，都在节点 0 上，仍然每个方向各跨一次。
• PP 发送：在第 4 层之后，GPU 11 将其 [T/4, d] 切片发送到节点 3 上的 GPU 27。
• 第 5–8 层：在节点 2 和 3 上同样如此（TP 伙伴 26，CP 伙伴 24 和 25，中继 GPU 19）。每一层将 token 发送到节点 2 上的一个专家和节点 3 上的一个专家：GPU 23 和 28、16 和 25（已经是 CP 伙伴）、20 和 29、21 和 31。

只有当多个微批次同时在途时，流水线才能带来收益：其他请求，或一个长提示词的连续块（分块流水线并行）。若只有一个，则每个阶段在另一个工作时处于空闲。由于 DP 副本在每个 MoE 层相遇，每个微批次（下文的 A–C）都跨越全部四个副本，它们同步推进：

行程日志：第 4 层后，11/32 GPU · 6 次节点穿越 · 20 次通信操作；PP 发送后，12/32 · 7 · 21。

跳 6：将 K/V 留在 GPU 10、11、26、27 上

除非请求提示词对数概率，否则只有最后一个提示词位置 T − 1（在 GPU 25 上，块 3）会馈入 LM 头。阶段 1 的 TP 对各自持有按词表拆分的 LM 头，但只有 GPU 24 和 25 使用它们：GPU 25 将该行复制到 GPU 24，各自计算 V/2 个 logits，然后收集两半用于采样。

对于除 T − 1 之外的每个 token，只有 K/V 在预填充后存活。我们的是 2·n_kv·d_h·8 个元素：KV 头的前一半在 GPU 10（第 1–4 层）和 26（第 5–8 层）上，后一半在 GPU 11 和 27 上。

第 8 层在注意力之前计算其 K/V，且没有下一层，因此其注意力和 MoE 仅服务于 T − 1：我们 token 的第 8 层专家行程（GPU 21 和 31，其 15 次节点穿越中的 2 次）是无效工作。引擎无论如何都会计算它，而使用 EAGLE 或 MTP（多 token 预测）推测解码时它不会浪费：草稿头会读取每个最终隐藏状态。

全程：地图与流量

行程日志：23/32 GPU · 15 次节点穿越（14 次 EP，1 次 PP）· 41 次通信操作。

每 GPU 每层发送的元素数（g = n_h/n_kv；时间基于 450 GB/s NVLink、50 GB/s IB）：

在均匀路由下，1 − (16/32)·(15/31) ≈ 76% 的 token 在另一节点上至少有一个专家，每个往返各穿越一次：2·0.76·(T/4)·d ≈ 0.38·Td。该表还假设 BF16、峰值链路速率且无计算-通信重叠。偏斜路由会使 EP 更差；FP8 分发并将 all-to-all 与计算重叠会使其更好。无论哪种方式，EP 在最慢链路上移动的数据最多。

交换 PP 和 DP（rank = tp + 2·cp + 4·pp + 8·dp）会将每个阶段分散到所有四个节点。PP 发送将移到 NVLink 上，但一个 token 的候选专家中将有四分之三位于 IB 对面，而非一半，使 EP 的 IB 流量从每层约 0.38·Td 增加到 0.67·Td。四层这样的流量超过一次 (T/4)·d 发送。

Token 的来世

假设此网格也运行解码，且第一个输出 token（位置 T）落在 GPU 11 上。在每一层中，其查询都会读取我们 token 的 K/V。

预填充将 K/V 移动到查询；解码将查询移动到 K/V（Meta 的 CP 推理论文中的“pass-Q”）。在 all-gather 之后，GPU 10 和 11 各自持有查询头的一半，并将它们发送给其 CP 伙伴 GPU 8 和 9；部分输出返回以进行合并。新 token 的 K/V 在注意力之前写入，位于 GPU 10 和 11 上。将所有新 token 保留在那里会增长 CP rank 1 的缓存，因此实现以轮询方式将新 token 分配给 CP rank。

vLLM 的解码上下文并行也按 token 分片 KV 缓存，但从 TP 组中划分出分片；对于 GQA，它需要 TP > n_kv，而此网格不满足。其他轴也会变化：

由于瓶颈不同，大型部署将预填充和解码运行在具有不同网格的独立池上（DeepSeek：预填充 EP32，解码 EP144）。在这种设置下，我们 token 的 K/V 与其余 KV 缓存一起，从预填充池到解码池进行最后一次传输。

三个答案

• 哪些 GPU 组成哪个组：rank = tp + 2·cp + 4·dp + 16·pp。TP 和 CP 伙伴位于 token 的 4-GPU 副本内；其 EP 组跨越 16-GPU 阶段。
• 为什么是这个顺序：CP、SP 和 TP 每一层都通信，因此它们从不离开节点。EP 有意跨节点，并承载 15 次节点跨越中的 14 次。承载第 15 次的 PP 位于最外层，因为将 DP 放在其外面会使 EP 的 IB 流量几乎翻倍。
• 为什么 DP 副本不再独立：它们共享一组专家，因此每个 MoE 层都是一次会合。