Tensor Parallelism 执行路径#

第四章前三节解释了 manager 边界和 IPC 拓扑。这一节回答另一个问题：当 tp_size > 1 时，同一个 forward pass 是怎样被分配到多张 GPU 上的？

这不是"多 GPU 各跑各的请求"，而是"同一个请求的同一次前向，被拆开在多张 GPU 上协同计算"。

这一节解决什么问题#

  1. Tensor Parallelism 的权重切分方式——哪些矩阵被横切，哪些被纵切；
  2. 每次前向计算后，多个 rank 怎样通过 AllReduce 把结果合并；
  3. ModelRunner 在多 rank 场景下怎样让所有 rank 保持动作一致；
  4. 调试 TP 相关问题时先看哪里。

一张图先看 TP 的权重切分#

对 Transformer 模型，TP 主要切分两类矩阵：

单 GPU（tp_size=1）            4 GPU（tp_size=4）
─────────────────────          ──────────────────────────────────
Attention QKV projection       每 GPU 处理 1/4 的 attention heads
    [d_model, 3*d_model]  →    [d_model, 3*d_model/4] × 4 GPU

Attention output projection    每 GPU 处理 1/4 的行
    [d_model, d_model]    →    [d_model/4, d_model] × 4 GPU
                                + AllReduce
MLP gate/up projection         每 GPU 处理 1/4 的 hidden_dim
    [d_model, 4*d_model]  →    [d_model, d_model] × 4 GPU

MLP down projection            每 GPU 处理 1/4 的输入维度
    [4*d_model, d_model]  →    [d_model, d_model] × 4 GPU
                                + AllReduce

这种切分方式被称为 Megatron-LM 风格的 TP：attention heads 被均匀分配给各 rank（column parallel），输出矩阵按行切分（row parallel），每层结束时通过 AllReduce 合并各 rank 的部分结果。

前向执行的实际流程#

以 tp_size=4 为例，当 ModelRunner.forward(batch) 被调用时：

Step 1：所有 rank 收到相同的 ForwardBatch

rank 0 是 primary rank，负责从 scheduler 拿到 ForwardBatch，然后通过 NCCL broadcast 把 input_ids、positions、seq_lens 等输入张量发给 rank 1、2、3。这一步让所有 rank 开始时状态一致。

Step 2：各 rank 用自己的权重切片独立计算

在 attention 层：

  • rank 0 计算 head 0…(H/4-1) 的 Q、K、V；
  • rank 1 计算 head H/4…(H/2-1) 的 Q、K、V；
  • 以此类推。

每个 rank 只需要 GPU 显存里自己那部分权重，计算完全并行，没有通信。

Step 3：AllReduce 合并部分结果

在 attention output projection（row parallel）之后：

# 每个 rank 算出的是部分 output：output_partial [batch, seq, d_model]
# AllReduce 把 4 个 rank 的 partial output 相加，得到完整 output
output = all_reduce(output_partial)  # NCCL AllReduce

MLP down projection 之后同样有 AllReduce。

Step 4：rank 0 收集结果，继续后续处理

AllReduce 之后所有 rank 持有相同的 activation，继续下一层。到最后一层 logits 输出时，rank 0 负责把结果传回给 scheduler。

KV Cache 在 TP 下怎样分布#

KV cache 和 attention heads 一起被切分：rank i 只存储自己负责的 attention heads 对应的 K 和 V。

这意味着：

  • 每个 rank 的 token_to_kv_pool 只占总 KV 的 1/tp_size；
  • prefix reuse 时，每个 rank 各自查询自己的 RadixCache，各自决定哪些前缀可以复用；
  • 无需在 rank 之间同步 KV cache 状态（各 rank 的 cache 是对称的）。

这个设计让 KV 管理在 TP 下几乎无额外开销，但也意味着如果某一张 GPU 的 KV 显存先用完，整个 batch 就会被这张卡限制。

tp_rank == 0 的特殊职责#

从代码里可以看到多处 if self.tp_rank == 0: 判断。这不是因为 rank 0 “更重要”，而是系统需要一个主节点来：

  • 接收 scheduler 发来的 ForwardBatch；
  • 广播 input 给其他 rank；
  • 收集最后的 logits 输出；
  • 把结果返回给 DetokenizerManager。

其他 rank（1、2、3）在逻辑上是被动执行者：等待 rank 0 广播，计算本 rank 的权重切片，参与 AllReduce，然后继续等待下一轮。

从 IPC 拓扑看，rank 0 持有 recv_from_tokenizer 和 send_to_detokenizer 这两条通道，rank 1-3 不持有（上一节已经提到）。

Pipeline Parallelism 与 TP 的关系#

pp_size > 1 时（Pipeline Parallelism），模型的 Transformer 层被按组分配给不同的 PP rank：

  • rank 0 负责前 N/pp_size 层；
  • rank 1 负责接下来的 N/pp_size 层；
  • 以此类推，最后一个 PP rank 输出最终 logits。

PP 与 TP 可以同时开启，形成 3D 并行（TP × PP × DP）。但 PP 与 TP 的通信模式不同：

  • TP 的 AllReduce 是层内同步（同一层的多 GPU 相互通信）；
  • PP 的 P2P 是层间异步（每层计算完以后把 activation 送给下一个 PP rank）。

在 SGLang 的场景下，PP 通常在推理时引入 bubble（等待上游 rank 完成计算），对延迟有影响。相比之下，TP 在单节点内（NVLink 通信）的 AllReduce 开销很小，是默认的多 GPU 扩展策略。

调试 TP 相关问题时先看哪里#

如果看到的现象是：

  • 开了 TP 但显存占用不均衡（某张卡 OOM 另一张没满）；
  • 开了 TP 后输出结果和单 GPU 不一致；
  • 某种 batch size 下 TP 速度反而变慢；

更稳的顺序通常是：

  1. 输出不一致：先确认 AllReduce 是否完成（rank 0 收到的是否是 all_reduce 后的完整结果，而不是 partial result）；检查 NCCL 版本和 driver 兼容性。

  2. 显存不均衡：检查是否有 rank 持有额外状态（例如 rank 0 持有更多的输入缓冲）；检查 KV cache 是否按 tp_size 正确切分了。

  3. TP 下速度变慢：检查 AllReduce 是否是瓶颈（用 nsys 看 NCCL 内核的时间占比）；验证通信是否走 NVLink 而不是 PCIe。

小结#

这一节真正要建立的是一个具体执行图：

  • TP 把模型权重切分到多张 GPU，每张 GPU 只存一份切片；
  • 所有 rank 用相同的 ForwardBatch，独立计算后 AllReduce 合并；
  • rank 0 承担协调职责，其他 rank 是对称的执行者；
  • KV cache 随 attention heads 一起切分，无需跨 rank 同步。

理解了这张图，再看多 GPU 下的显存估算、延迟分析和 debug 路径时，就不会把 TP 当成"多 GPU 多跑几个请求"来理解。