《vLLM 内核探秘》完整目录
- 前言
- 第1章 架构总览
- 第2章 EngineCore：引擎的心脏
- 第3章 调度器：Token 的交通指挥
- 第4章 PagedAttention：虚拟内存的启示
- 第5章 KV Cache 管理：寸土寸金的显存
- 第6章 Worker 与 Executor：GPU 军团
- 第7章 模型加载与权重管理
- 第8章 前向计算与 CUDA Graph
- 第9章 采样与输出处理
- 第10章 前缀缓存：零开销的加速
- 第11章 分块预填充与混合批处理
- 第12章 投机解码：以小博大
- 第13章 量化引擎：精度与速度的平衡
- 第14章 张量并行与流水线并行（当前）
- 第15章 多模态推理
- 第16章 LoRA 适配器热切换
- 第17章 API 服务器与生产部署
- 第18章 设计模式与架构哲学
第14章 张量并行与流水线并行
"Divide and conquer, but make sure the pieces can talk to each other."
:::tip 本章要点
- 理解张量并行（TP）的原理：如何在不改变数学结果的前提下切分矩阵乘法
- 掌握 Megatron-LM 张量并行方案在 vLLM 中的应用
- 理解流水线并行（PP）的基本原理与调度策略
- 深入 NCCL 通信原语在 vLLM 中的使用
- 认识 TP vs PP 的选择策略：什么时候用哪种 :::
14.1 大模型的困境
Llama-2-70B 在 FP16 下需要 140 GB 显存。一张 A100-80GB 放不下。即使量化到 FP8（70 GB），也需要考虑 KV Cache 的额外显存需求------实际上至少需要两张卡。
更大的模型（如 405B）需要 8 张甚至 16 张卡。如何将计算分布到多张卡上，且保持高效的 GPU 利用率？
分布式推理面临的核心矛盾是：GPU 间的通信带宽远低于 GPU 内部的计算带宽。 一张 A100 的显存带宽是 2 TB/s，而 NVLink 的双向带宽是 600 GB/s（差 3 倍），以太网/InfiniBand 更是只有 25-100 GB/s（差 20-80 倍）。这意味着跨 GPU 的数据传输是分布式推理的主要瓶颈，而非计算本身。
因此，分布式策略的核心目标是：最小化 GPU 间的通信量和通信次数。 两种并行策略在这个目标上的取舍截然不同：
- 张量并行（Tensor Parallelism, TP）------每层都需要通信（AllReduce），但每次通信量小。适合高带宽互联（NVLink）。
- 流水线并行（Pipeline Parallelism, PP）------只在阶段边界通信（点对点传输），通信次数极少。适合低带宽互联（跨机器）。
两种并行策略各有适用场景：
GPU 0"] --> PPB["Layer 21-40
GPU 1"] end style TPA fill:#3b82f6,color:#fff,stroke:none style TPA2 fill:#10b981,color:#fff,stroke:none style TPB fill:#3b82f6,color:#fff,stroke:none style TPB2 fill:#10b981,color:#fff,stroke:none style PPA fill:#3b82f6,color:#fff,stroke:none style PPB fill:#10b981,color:#fff,stroke:none
14.2 张量并行（Tensor Parallelism）
张量并行将每一层的计算切分到多张卡上。核心思想来自 Megatron-LM：利用矩阵乘法的可分性。
列并行（Column Parallel）
源码 ：
vllm/model_executor/layers/linear.py:345
vLLM 的 ColumnParallelLinear 实现了列切分：
python
# vllm/model_executor/layers/linear.py:345-384 (简化)
class ColumnParallelLinear(LinearBase):
"""Linear layer with column parallelism.
The linear layer Y = XA is parallelized along A's second dimension:
A = [A_1, ..., A_p].
"""
def __init__(self, input_size, output_size, ...):
self.tp_size = get_tensor_model_parallel_world_size()
# 每张卡只持有 output_size / tp_size 列
self.output_size_per_partition = output_size // self.tp_size对于线性层 Y = XW，如果将 W 按列切分为 [W₁ | W₂]：
css
Y = X × [W₁ | W₂] = [X×W₁ | X×W₂]GPU 0 计算 X×W₁，GPU 1 计算 X×W₂。输入 X 在两张卡上完全相同（广播），输出 Y 在两张卡上各持有一半。
适用于：QKV 投影（按注意力头切分）、MLP 的 gate/up 投影（按中间维度切分）。
行并行（Row Parallel）
对于线性层 Y = XW，如果将 W 按行切分：
css
Y = [X₁ X₂] × [W₁; W₂] = X₁×W₁ + X₂×W₂GPU 0 计算 X₁×W₁，GPU 1 计算 X₂×W₂，然后做 AllReduce（求和）得到最终结果。
适用于：注意力的输出投影、MLP 的 down 投影。
通信开销
张量并行在每个 Transformer 层中引入两次 AllReduce 通信（一次在注意力后，一次在 MLP 后）。在 NVLink 连接的 GPU 之间（如 A100 DGX 的 600 GB/s 双向带宽），这个开销很小。但如果跨机器（通常只有 100 Gbps ≈ 12.5 GB/s），AllReduce 的延迟就会成为显著瓶颈。
经验法则：张量并行只在 NVLink 连接的 GPU 之间使用，不要跨机器。
14.3 流水线并行（Pipeline Parallelism）
流水线并行将模型的不同层放在不同的 GPU 上。数据"流"过各个阶段：
流水线并行只在阶段边界传输一次激活值（而非每层两次 AllReduce），通信量远小于张量并行。因此适合跨机器部署。
但流水线并行有一个固有问题：气泡（Pipeline Bubble）。当 GPU 0 在处理第一批数据时，GPU 1 在空闲；反之亦然。微批次化（将一个批次切成多个微批次）可以减少气泡，但无法完全消除。
vLLM 通过 Ray Compiled DAG 实现跨机器的流水线并行，将多个通信步骤编译为一个固定的执行图，减少了每步的调度开销。
14.4 TP + PP 的组合
实践中，两种并行经常组合使用。例如 8 卡 2 机部署 405B 模型：
ini
机器 A (4 × A100): TP=4, PP stage 0 (Layer 1-40)
机器 B (4 × A100): TP=4, PP stage 1 (Layer 41-80)机内 TP=4（利用 NVLink 的高带宽），机间 PP=2（跨机通信量最小化）。
通信量对比
理解为什么这样组合，需要比较两种并行的通信模式：
张量并行的通信 ：每个 Transformer 层有 2 次 AllReduce。假设隐藏维度 d=8192，batch_size=1，每次 AllReduce 传输的数据量 = batch_tokens × d × dtype_size。对于 1 个 Token 的解码步：1 × 8192 × 2 bytes = 16 KB。看起来很小，但 AllReduce 的延迟主要来自启动开销（几微秒），而非数据传输。80 层模型 × 2 次/层 = 160 次 AllReduce/步------NVLink 的低延迟（1-2 μs）可以承受，但跨机器的 RDMA（10-50 μs）会累积到毫秒级。
流水线并行的通信 ：只在 PP 阶段边界传输一次激活值（batch_tokens × d × dtype_size）。2 个 PP 阶段只有 1 次点对点通信/步，延迟完全可以被 GPU 计算覆盖。
每层 2 次 AllReduce
总计 160 次/步
适合 NVLink"] PP["流水线并行
阶段间 1 次 P2P
总计 1 次/步
适合跨机"] end style TP fill:#3b82f6,color:#fff,stroke:none style PP fill:#10b981,color:#fff,stroke:none
如何选择并行策略
| 场景 | 推荐策略 | 原因 |
|---|---|---|
| 单机 2-8 卡 | 纯 TP | NVLink 带宽足够，PP 有气泡浪费 |
| 2 机 8-16 卡 | 机内 TP + 机间 PP | 最小化跨机通信量 |
| 模型能放进单卡 | 不并行 | 并行有通信开销 |
| 延迟敏感 | 优先 TP | PP 有流水线气泡增加延迟 |
| 吞吐优先 | TP + PP | 更多卡 = 更多并发 |
气泡问题的缓解
流水线并行的气泡（bubble）是固有的------当第一个阶段在计算时，最后一个阶段在等待。微批次化（micro-batching）可以减少气泡比例：
yaml
微批次数 = 1: 气泡率 ≈ 50%（2 个阶段）
微批次数 = 4: 气泡率 ≈ 20%
微批次数 = 8: 气泡率 ≈ 11%vLLM 通过连续批处理自然地实现了微批次化------每一步的批次就是一个"微批次"，连续的步骤形成了流水线。
14.5 KV Cache 的分布式传输
当模型采用流水线并行时，一个有趣的问题出现了：KV Cache 需要跟着数据流动吗？
答案是不需要 。每个 PP 阶段只负责自己那些层的 KV Cache。Layer 1-20 的 KV Cache 在 GPU 0 上，Layer 21-40 的在 GPU 1 上。数据在阶段间传输的是激活值（隐藏状态），不是 KV Cache。
但在一种新兴的架构------**Prefill-Decode 分离（Disaggregated Serving）**中，KV Cache 确实需要在 GPU 之间传输。这种架构的动机是：预填充和解码对 GPU 资源的需求截然不同。
- 预填充是计算密集型（大量 Token 的自注意力），适合高算力 GPU
- 解码是内存带宽密集型（每步只处理 1 个 Token），适合高带宽 GPU
将两者分离到不同的 GPU 池，可以独立扩缩容------预填充 GPU 可以处理完一个请求后立即服务下一个请求，不被慢速的解码过程拖住。
传输的数据量有多大？一个 Llama-2-70B 的请求，1000 Token 的 KV Cache 约需传输 num_layers × 2 × num_kv_heads × head_dim × 1000 × 2 bytes ≈ 80 × 2 × 8 × 128 × 1000 × 2 = 328 MB。在 100 Gbps 的 RDMA 网络下，传输时间约 26ms------与预填充本身的计算时间（~50-100ms）相比是可接受的。vLLM 在 vllm/distributed/kv_transfer/ 中实现了这个 KV Cache 传输机制。
14.6 vLLM 的并行状态管理
源码 ：
vllm/distributed/parallel_state.py
init_distributed_environment()（parallel_state.py:860）是分布式初始化的入口：
python
# vllm/distributed/parallel_state.py:860-895 (简化)
def init_distributed_environment(
world_size, rank, distributed_init_method="env://",
local_rank=-1, backend="nccl",
):
# 如果启用了数据并行(DP)，调整 rank 和 world_size
if config.parallel_config.data_parallel_size > 1:
rank = dp_rank * world_size + rank
world_size = world_size_across_dp
# 初始化 PyTorch 分布式进程组
if not torch.distributed.is_initialized():
torch.distributed.init_process_group(
backend=backend, # NCCL for GPU
init_method=distributed_init_method,
world_size=world_size,
rank=rank)初始化后，每个 Worker 通过 get_tensor_model_parallel_rank() 和 get_pipeline_model_parallel_rank() 获取自己在 TP/PP 组中的位置，决定加载哪些层、切分哪些权重。
vllm/distributed/parallel_state.py 管理所有并行相关的全局状态：
- TP 组------同一 PP 阶段内做张量并行的 GPU 组
- PP 组------不同 PP 阶段之间的 GPU 组
- World 信息------TP rank、PP rank、总 rank 等
Worker 在初始化时根据自己的 rank 加入正确的通信组，之后的 AllReduce、Send/Recv 都在组内完成。
14.7 本章小结
- 张量并行------切分层内计算，每层 2 次 AllReduce，适合 NVLink 连接的卡
- 流水线并行------切分层间数据流，只在阶段边界通信，适合跨机器
- 组合使用------机内 TP + 机间 PP 是大模型部署的标准范式
- 通信优化------NCCL 原语 + Ray Compiled DAG 减少通信开销
14.8 实际部署案例
案例一：单机 8 卡部署 Llama-3-405B
bash
# 8 张 A100-80GB，纯 TP
vllm serve meta-llama/Llama-3-405B \
--tensor-parallel-size 8 \
--max-model-len 8192 \
--gpu-memory-utilization 0.92405B FP16 = 810 GB，8 张 A100 × 80 GB = 640 GB。放不下！需要 FP8 量化：
bash
# FP8 量化后 405 GB，8 卡可以放下
vllm serve meta-llama/Llama-3-405B \
--tensor-parallel-size 8 \
--quantization fp8 \
--max-model-len 8192案例二：2 机 16 卡部署 405B（高并发）
bash
# 机内 TP=8, 机间 PP=2
# 机器 A (GPU 0-7): PP stage 0, TP=8
# 机器 B (GPU 8-15): PP stage 1, TP=8
vllm serve meta-llama/Llama-3-405B \
--tensor-parallel-size 8 \
--pipeline-parallel-size 2 \
--distributed-executor-backend rayPP=2 将模型分成两段（各 ~200 层），机内 TP=8 利用 NVLink 高带宽，机间 PP 只传输激活值（通信量小）。这种配置下，两台机器的 KV Cache 总容量翻倍，支持更高的并发。
案例三：边缘部署 7B 模型
bash
# 单卡部署，无需并行
vllm serve meta-llama/Llama-3-8B \
--max-model-len 4096 \
--gpu-memory-utilization 0.95小模型单卡即可，不需要任何并行策略。vLLM 自动选择 UniProcExecutor，零通信开销。
部署决策的一般原则
分布式部署的核心决策是如何分配 TP 和 PP 的比例。以下是实践中总结的经验：
原则一：先用 TP 填满一台机器内的 GPU。 NVLink 的带宽（600 GB/s 双向）比跨机网络（100-200 Gbps ≈ 12-25 GB/s）高一个数量级。TP 的 AllReduce 通信每层发生两次，需要高带宽------因此 TP 必须限制在 NVLink 连接的 GPU 之间。
原则二：跨机器只用 PP。 PP 的点对点传输只在阶段边界发生一次（而非每层两次），通信量和频率都远低于 TP。即使在较慢的网络上，PP 的开销也可以被 GPU 计算时间覆盖。
原则三：并行度不要超过需要。 每增加一级并行都引入额外的通信开销和代码复杂度。如果模型能放进 2 张卡（TP=2），就不要用 4 张卡（TP=4）------虽然理论上 4 卡更快，但通信开销可能抵消算力增益。
原则四：监控 GPU 利用率。 如果 nvidia-smi 显示 GPU 利用率不均衡（一些卡 90%，另一些 50%），说明并行策略不平衡------可能需要调整 TP/PP 的分配或检查层的划分是否均匀。
源码导航
- 并行状态：
vllm/distributed/parallel_state.py- 通信操作：
vllm/distributed/communication_op.py- 设备通信器：
vllm/distributed/device_communicators/- Ray Executor：
vllm/v1/executor/ray_distributed_executor.py