6.2 张量并行推理

6.1 节说过，和 PP 不同，TP 能让多张卡同时算同一个 Token，所以它是降低单请求延迟的首选。代价是每层都要同步。本节先把一层的切法和两次 AllReduce 讲清，再算通信账，最后讨论 TP 度该取多大，以及为什么 TP 几乎从不跨节点。

第 3 章 3.3 节已经在 nano-vllm 的源码里读过 ColumnParallelLinear 和 RowParallelLinear。本节从部署者的角度，看这些切法在真实模型和硬件上意味着什么。

⏰ 信息时效：原理部分与框架无关；vLLM 相关描述基于 v0.32.0（2026-10-18 发布）。硬件带宽取自 NVIDIA 官网产品页。

📑 目录

• 1. 一层怎么切：列切和行切配对
• 2. 注意力按头切：GQA 带来的 KV 头复制
• 3. 通信账：两次 AllReduce 花多少时间
• 4. 为什么 TP 不出节点
• 5. TP 度：延迟和吞吐要分开看
• 6. vLLM 中的 TP 实现细节
• 7. 动手算：TP 度与每 Token 时间的粗略模型
• 总结
• 自我检验清单
• 参考资料

---

1. 一层怎么切：列切和行切配对

一个 Transformer 层有两个子层：注意力和 FFN。TP 的切法来自 Megatron-LM 论文，核心是让列切和行切成对出现。

拿 FFN 举例。设输入 $X\in\mathbb R^{n\times h}$，FFN 是 $Y=\text{act}(XA)\,B$，其中 $A\in\mathbb R^{h\times f}$，$B\in\mathbb R^{f\times h}$。

• 第一个矩阵按列切：$A=[A_1,\ldots,A_t]$，第 $i$ 张卡算 $\text{act}(XA_i)$。激活函数逐元素作用，每张卡自己算自己那一列，结果正确，不需要通信。
• 第二个矩阵按行切：$B$ 沿 $f$ 维切成 $B_1,\ldots,B_t$，第 $i$ 张卡算 $\text{act}(XA_i)B_i$，得到形状完整的 $n\times h$ 部分和。

$$Y=\sum_{i=1}^{t}\text{act}(XA_i)\,B_i$$

把 $t$ 份部分和加起来，就是一次 AllReduce。中间维 $f$ 上的激活一直留在各自的卡上，从不完整拼出来。

注意力的切法相同：QKV 投影按头列切，每张卡算自己那几个头的完整注意力；输出投影 o_proj 按行切，结束时再做一次 AllReduce。于是每层恰好 2 次 AllReduce。

下图是 Llama-3.1-70B 的一层在 TP=4 下的数据流，四种颜色是四张 GPU 各自的分片，紫色竖条是需要所有卡参与的 AllReduce：

几个数从模型配置来：64 个 Q 头除以 4 得 16，8 个 KV 头除以 4 得 2，FFN 中间维 28672 除以 4 得 7168。

📌 关键点：TP 的正确性只依赖一件事：行切之后的部分和加起来等于原结果。所以每个”列切 + 行切”的组合末尾必须有一次 AllReduce，少一次结果就错了。

---

2. 注意力按头切：GQA 带来的 KV 头复制

2.1 头数决定 TP 能取哪些值

注意力按头切分，所以 Q 头数必须能被 TP 大小整除。Llama-3.1-70B 有 64 个 Q 头，TP 可以取 1、2、4、8、16……；Qwen3-32B 也是 64 个 Q 头，同样如此。实际部署中 TP 还受单节点 GPU 数限制，常见取值是 2、4、8。

2.2 KV 头比 TP 少时怎么办

GQA 让多个 Q 头共享一组 K、V。Llama-3.1-70B 只有 8 个 KV 头，TP=16 时每张卡分不到一个完整的 KV 头。vLLM 的做法是复制：TP 大于 KV 头数时，每张卡存 1 个 KV 头，每个 KV 头被复制到 TP / KV 头数张卡上（这时要求 TP 能被 KV 头数整除）。

下图对比 TP=8 和 TP=16，蓝色格的数字是每卡的 Q 头数，绿色格的数字是这张卡存的 KV 头编号：

复制带来的后果：

• 每卡 KV Cache 不再随 TP 缩小。TP 从 8 加到 16，每卡的权重减半，每个 Token 的 KV 却保持不变。6.1 节的估算里，TP=16 的 KV 容量只比 TP=8 多约 18%。
• 整体 KV 显存翻倍。同一份 KV 在两张卡上各存一次，16 张卡上 KV 的总量是不复制时的 2 倍。

所以对 GQA 模型，KV 头数是 TP 的自然上限。超过它之后，再要扩展就该换 PP 或 DP。vLLM 还提供 Decode Context Parallel（DCP），沿 Token 维度切 KV Cache 来消除这部分重复，6.6 节会介绍。

⚠️ 注意：MLA（如 DeepSeek-V3）只缓存一份压缩后的潜向量，没有”KV 头”可以切，TP 下每张卡都要存完整的潜向量 KV。这也是 DeepSeek-V3 这类模型通常让注意力走 DP 而非大 TP 的原因之一（6.4 节）。

---

3. 通信账：两次 AllReduce 花多少时间

3.1 一次 AllReduce 的代价

nccl-tests 的性能文档推导了 AllReduce 的传输时间：大小为 $S$ 的数据在 $t$ 张卡上做 AllReduce，每个元素需要 $2(t-1)$ 次传输，由 $t$ 条带宽为 $B$ 的链路分担，理想时间是 $\frac{2(t-1)}{t}\cdot\frac{S}{B}$。再加上固定开销，总时间约为：

$$T_{\text{AR}}\approx\alpha+\frac{2(t-1)}{t}\cdot\frac{S}{B}$$

$B$ 是每张卡的链路带宽，$\alpha$ 是固定开销（Kernel 启动、各卡对齐、链路延迟）。前一项和数据量无关，后一项和数据量成正比。

3.2 Decode 的消息很小

Decode 时 $S=n\times h\times b$，$n$ 是这一步的 Token 数。Llama-3.1-70B（$h=8192$，BF16）：

Decode 时带宽项只有几微秒，和固定开销同一量级甚至更小，所以 Decode 阶段的 TP 通信主要受延迟限制。每步 160 次 AllReduce，每次多 10 µs，TPOT 就多 1.6 ms。

Prefill 一次处理成千上万个 Token，消息大到上百 MiB，变成带宽限制，此时 NVLink 和 PCIe 的差距会被放大好几倍。

💡 提示：用 nccl-tests 的 all_reduce_perf 在自己的机器上扫一遍 8 B 到几百 MiB 的消息大小，能直接看到延迟区和带宽区的分界。它报告的 busbw 已经乘了 $2(t-1)/t$，可以和链路带宽直接比。

---

4. 为什么 TP 不出节点

一台 8 卡 H100 SXM 服务器里，GPU 之间通过 NVLink 和 NVSwitch 全互联，NVIDIA 给出的 NVLink 带宽是每卡 900 GB/s（两个方向合计）。跨节点要走网卡：以 DGX B200 为例，8 张 GPU 配 8 个单端口 ConnectX-7，每个 400 Gb/s，平均每张 GPU 约 50 GB/s。

把 TP=16 摊到两台 8 卡机器上，每次 AllReduce 都必须有数据跨过网卡。Ring 算法下整个环的速度由最慢的一段决定，带宽从 450 GB/s 掉到 50 GB/s，跨机的延迟也比机内高。每层两次、每步上百次，代价叠加起来就很明显。

下图左边是 TP=16 横跨两台机器，右边是 TP=8 留在机内、两台机器之间用 PP=2 连接：

这就是 vLLM 文档的建议：TP 大小取每个节点的 GPU 数，PP 大小取节点数。PP 每步只在段边界传一次激活，次数从上百次降到个位数。

例外是机柜级 NVLink 系统，比如 GB200 NVL72，72 张 GPU 在同一个 NVLink 域内。在这类系统上”节点”的边界变了，TP 可以跨物理服务器，判断标准应换成”是否在同一个 NVLink 域”。

---

5. TP 度：延迟和吞吐要分开看

TP 度加大时，几件事同时发生：

• 每卡读的权重变少：Decode 每步的访存时间约按 $1/t$ 下降，这是延迟收益的主要来源。
• 通信次数不变、每次参与的卡变多：每步仍是 $2L$ 次 AllReduce，固定开销不随 $t$ 下降，在总时间里的占比越来越高。
• 每卡的矩阵变小：GEMM 的 N 维或 K 维被切成 $1/t$，小矩阵更难用满 Tensor Core，Prefill 阶段尤其明显。
• 每卡空出更多显存：权重和 KV（在 KV 头数以内）都按 $1/t$ 缩小，能放下更多并发请求。

前三条说明，TP 度翻倍，延迟下降不到一半，单卡吞吐往往还会下降。第四条又说明，当系统被 KV 容量卡住时，增大 TP 能让每张卡多服务几个请求，单卡吞吐反而可能上升。

Pope 等人在 PaLM 推理的研究里得出过类似的结论：没有一种切分布局在所有 Batch 大小下都最优，低延迟和高吞吐需要不同的布局。落到 vLLM 部署上，可以按下面的思路选：

⚠️ 注意：比较不同 TP 度的吞吐时，一定要按卡数归一化，并在相同的延迟约束下比较。TP=8 的单副本吞吐比 TP=4 高，不代表 8 张卡上开一个 TP=8 比开两个 TP=4 更划算。

---

6. vLLM 中的 TP 实现细节

6.1 启动方式

单节点只需一个参数：

vllm serve meta-llama/Llama-3.1-70B-Instruct --tensor-parallel-size 4

vLLM 为每张卡起一个 Worker 进程。单节点默认用 Python 多进程（mp）管理这些 Worker，多节点时可以用 Ray（6.5 节）。Worker 之间用 NCCL 做集合通信。

6.2 哪些层被切

Embedding 和 LM Head 的词表通常很大（Llama-3.1 是 128256），按词表切能省下可观的显存。3.3 节读过 nano-vllm 里对应的 VocabParallelEmbedding 和 ParallelLMHead。

6.3 Custom AllReduce

NCCL 的通用算法针对各种拓扑设计，在节点内的小消息上并不是最快的。vLLM 实现了自己的 AllReduce Kernel：节点内的 GPU 能直接读写彼此的显存时，各卡直接从对方的缓冲区读数据求和，少走几轮环形传输，对 Decode 这种延迟受限的小消息更友好。条件不满足或消息较大时，退回 NCCL。

遇到这部分相关的错误（例如某些虚拟化环境里 P2P 访问不可用），可以用 --disable-custom-all-reduce 关掉它，让所有 AllReduce 都走 NCCL，用来定位问题。

💡 提示：Custom AllReduce 没有启用时，启动日志会打印原因，例如 P2P 检测失败，或超过两张只有 PCIe 互联的 GPU。TP 性能异常时，先看这一段日志。

---

7. 动手算：TP 度与每 Token 时间的粗略模型

下面是一个假设成本模型：只算 Decode 一步读权重、读 KV 和 TP 通信的时间，忽略计算、Kernel 启动和调度。$\alpha=10$ µs 是假设值，不对应任何实测。TP=1 和 TP=2 时 70B 模型实际装不下，这里只为了看趋势。

# 假设成本模型：只估算 Decode 一步的访存时间和 TP 通信时间，忽略计算与 Kernel 启动
HBM = 3.35e12          # H100 SXM 显存带宽，B/s
NVLINK = 450e9         # NVLink 4 单向带宽，B/s（900 GB/s 为双向合计）
ALPHA = 10e-6          # 每次 AllReduce 的固定开销，假设值
LAYERS, HIDDEN = 80, 8192
WEIGHT_BYTES = 70.6e9 * 2
KV_PER_TOKEN = 320 * 1024        # 整个模型每 Token 的 KV，B
CTX = 4096                        # 每个请求的平均上下文长度


def step_ms(tp, batch):
    weight = WEIGHT_BYTES / tp / HBM
    kv = batch * CTX * KV_PER_TOKEN / min(tp, 8) / HBM   # TP>8 后 KV 不再缩小
    msg = batch * HIDDEN * 2
    comm = 0 if tp == 1 else 2 * LAYERS * (ALPHA + 2 * (tp - 1) / tp * msg / NVLINK)
    return (weight + kv + comm) * 1e3, comm * 1e3


print(f"{'TP':>3} {'batch':>5} {'每步 ms':>8} {'其中通信':>8} {'Token/s':>9} {'Token/s/卡':>10}")
for batch in [1, 64]:
    for tp in [1, 2, 4, 8]:
        t, c = step_ms(tp, batch)
        print(f"{tp:>3} {batch:>5} {t:>8.2f} {c:>8.2f} {batch / t * 1e3:>9.0f} {batch / t * 1e3 / tp:>10.0f}")

输出：

 TP batch    每步 ms     其中通信   Token/s  Token/s/卡
  1     1    42.55     0.00        24         24
  2     1    22.88     1.61        44         22
  4     1    12.25     1.61        82         20
  8     1     6.93     1.61       144         18
  1    64    67.79     0.00       944        944
  2    64    35.87     1.97      1784        892
  4    64    19.11     2.16      3350        837
  8    64    10.73     2.25      5967        746

在这组假设下：

• TP 从 1 到 8，单请求的每步时间从 42.6 ms 降到 6.9 ms，约 6.1 倍，不到 8 倍。差距来自 1.6 ms 的通信，它不随 TP 变小。
• Batch=64 时，TP=8 的单卡吞吐比 TP=2 低约 16%。如果 KV 容量不是瓶颈，同样 8 张卡开 4 个 TP=2 副本，总吞吐更高。
• 把 ALPHA 改成 30 µs，再看 TP=8 的结果，就能感受到跨节点或无 NVLink 时 TP 为什么不划算。

真实系统里还有计算时间、Attention Kernel 的效率差异和调度开销。这个模型只用来理解趋势，部署决策以 vllm bench serve 的实测为准（6.6 节）。

---

📝 总结

• TP 用”列切 + 行切”配对，中间激活留在各卡，每层 2 次 AllReduce，80 层模型每步 160 次。
• 注意力按头切，Q 头数要能被 TP 整除；GQA 模型 TP 超过 KV 头数后，vLLM 复制 KV 头，每卡 KV 不再缩小。
• Decode 的 AllReduce 消息只有 KiB 到 MiB 级，受延迟限制；Prefill 消息大，受带宽限制。
• NVLink 4 单方向 450 GB/s，跨节点的 400 Gb/s 网卡约 50 GB/s，所以 TP 留在节点内，跨节点用 PP。
• 增大 TP 能降低延迟，但单卡吞吐通常下降；KV 容量是瓶颈时情况相反，要按卡数归一化后比较。
• vLLM 的 Custom AllReduce 在节点内加速小消息，出问题时用 --disable-custom-all-reduce 排查。

🎯 自我检验清单

• 能推导 FFN 列切 + 行切后的结果为什么等于原结果，并说明 AllReduce 放在哪里
• 能根据 Q 头数和 KV 头数判断一个模型可用的 TP 度，以及从哪个 TP 开始复制 KV 头
• 能用 $\alpha+\frac{2(t-1)}{t}\frac{S}{B}$ 估算一次 AllReduce 的时间，并判断它受延迟还是带宽限制
• 能对比 NVLink、PCIe、InfiniBand 的每卡带宽，解释 TP 不跨节点的原因
• 能解释为什么 TP 度翻倍时延迟下降不到一半、单卡吞吐可能下降
• 能为”最低延迟”和”最高吞吐”两个目标分别选出 TP 度和副本数
• 能说出 vLLM 中哪些层按词表切分，以及 --disable-custom-all-reduce 的用途

📚 参考资料

• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Efficiently Scaling Transformer Inference
• NVIDIA H100 Tensor Core GPU
• NVIDIA NVLink and NVLink Switch
• NVIDIA DGX B200 User Guide
• nccl-tests Performance（AllReduce 时间模型与 busbw）
• vLLM 文档：Parallelism and Scaling
• vLLM 源码：vllm/model_executor/layers/linear.py（v0.32.0）
• vLLM 源码：vllm/distributed/device_communicators/custom_all_reduce.py（v0.32.0）