Skip to content
Self-Knowing
selfknoww

分布式推理优化

约 3407 个字 预计阅读时间 11 分钟

Outline

  • 视频定位与本讲边界
  • 为什么需要分布式推理
  • 通信基础：硬件、通信库、集体通信
  • 主流并行方式：TP / PP / EP / DP
  • 补充：Sequence Parallel vs Context Parallel
  • 为什么后面还会出现 PD / Prefix Caching / Speculative Decoding / Benchmarking
  • 从代码视角看这些概念
  • 小结

为什么需要分布式推理

分布式推理不是简单地“把模型切到多张卡上”，它本质上是在解决下面几类问题：

  • 单卡放不下模型参数。
  • 单卡放得下参数，但放不下高并发场景下的 KV Cache。
  • 单卡能跑，但吞吐不够，无法支撑更高的 QPS 或更大的 batch。
  • 同一个模型在不同阶段的资源需求并不一样，适合拆开调度。
  • prefill 阶段更偏算力密集。
  • decode 阶段更偏显存带宽和时延敏感。

一个很重要的认识是：推理优化不只是算子优化，也不只是模型压缩。规模一旦上来，系统架构、调度策略、通信开销、缓存命中率，往往比单个 kernel 的优化更影响最终效果。

分布式推理真正难的地方

随着 GPU 数量增多、节点数增多，系统难点会从“算不动”逐渐变成“协同不好”：

  • 通信开销会上升，尤其是跨卡、跨机同步。
  • 最慢的那个 worker 会拖累整个请求的尾延迟。
  • 某些并行策略非常依赖互联质量，硬件一变，收益可能立刻下降。
  • 多租户场景下，长请求、短请求、cache hit、cache miss 混在一起，调度复杂度会很高。
  • 集群规模越大，偶发故障、超时、重试、热点倾斜的问题越常见，因此健壮性会变得很重要。

可以粗略记一个公式：

总时延 = 计算时间 + 通信时间 + 调度等待时间 + 排队时间

很多时候，优化前两项很直观，但真正难的是后两项。

通信基础

CPU 通信 vs GPU Direct

CPU 中转通信的优点和缺点都很明显：

  • 优点：控制简单，做同步、阻塞、状态管理更容易。
  • 缺点：数据路径会变成 GPU -> CPU -> GPU，多了一次中转，延迟和带宽都不理想。

因此高性能推理系统更希望走 GPU 直连或尽量减少主机参与。

常见通信硬件

  • PCIe：最常见的主机与设备互联，通用但带宽和延迟通常不如专用 GPU 互联。
  • NVLink：NVIDIA GPU 之间的高速互联，适合单机多卡高频通信。
  • InfiniBand：常见的跨节点高速网络，很多大规模训练/推理集群都会用。
  • RDMA：Remote Direct Memory Access，允许一台机器直接访问另一台机器的内存，减少内核参与和数据拷贝。
  • 核心优势可以概括为：bypass kernel + zero copy。
  • 在 GPU 场景里，通常还会讨论 GPUDirect RDMA，即尽量避免 GPU 数据绕回 CPU 内存。

通信库

在 PyTorch / vLLM 这类系统里，常见通信栈大致是：

  • torch.distributed：统一的分布式接口层。
  • NCCL / PyNccl：NVIDIA GPU 常见通信实现。
  • shared memory：同机不同进程之间共享数据。
  • custom all-reduce / all-gather kernel：针对特定场景自定义通信 kernel，减少通用实现带来的额外开销。

在 vLLM 里，和设备通信相关的逻辑可以从 vllm/distributed/device_communicators 一路往下读。

常见集体通信操作

all-reduce

  • 含义：先做规约，再把结果发给所有参与方。
  • 常见规约：sum / mean / max
  • 最后每台机器都拿到同一个结果。

例子：4 台机器分别有 [0] [1] [2] [3]

  • all-reduce(sum) 之后：
  • 0 号机器拿到 [6]
  • 1 号机器拿到 [6]
  • 2 号机器拿到 [6]
  • 3 号机器拿到 [6]

all-gather

  • 含义：收集每个机器上的原始数据并拼接。
  • 不做求和，不做平均，只是“收齐”。
  • 最后每台机器都拿到完整拼接后的结果。

例子：4 台机器分别有 [0] [1] [2] [3]

  • all-gather 之后：
  • 每台机器都拿到 [0, 1, 2, 3]

reduce-scatter

  • 可以看成 all-reduce 的前半段和 scatter 的组合。
  • 先规约，再把结果拆分给不同机器。
  • 在张量并行里很常见，因为它比“先 all-reduce 再手动切片”更高效。

all-to-all

  • 每个 worker 都会给其他 worker 发一部分数据，同时也从其他 worker 收一部分数据。
  • MoE + Expert Parallel 里非常常见，因为 token 会被路由到不同 expert 所在的 GPU。

这几个操作里：

  • TP 很依赖 all-reduce / all-gather / reduce-scatter
  • EP 很依赖 all-to-all
  • PP 更像是相邻 stage 之间传中间激活

主流并行方式

一个重要结论是：实际系统里几乎从来不是只用一种并行方式，而是混合并行。比如：

  • 单机多卡 dense 模型：TP + DP
  • 超大 dense 模型：PP + TP + DP
  • MoE 模型：EP + DP，有时再叠一点 TP

1. Tensor Parallel（TP）

TP 的核心思想是：把同一层里的张量切到多张卡上算。

常见做法包括：

  • 按 hidden dimension 切分线性层权重。
  • 在 attention 里按 head 或投影矩阵切分。
  • 在 MLP 里把 up/down projection 分到不同 GPU。

所以更准确地说，TP 不只是“给 attention 用的”，而是对一层内部的矩阵计算做并行；attention 是最常见场景之一，但 MLP 也一样能做 TP。

TP 的优点：

  • 能直接降低单个请求的计算压力。
  • 对超大 dense 模型很关键，不然模型根本放不下。
  • 在单机高带宽互联下，通常能显著改善延迟。

TP 的代价：

  • 通信频率高，几乎每层都要同步。
  • 很依赖卡间互联质量，跨机 TP 成本通常明显更高。
  • 扩展性不是无限的，切太细之后，通信会吃掉收益。

一个常见经验是：如果卡间互联一般，TP 的收益很容易被通信抵消；如果是 NVLink 这种高带宽环境，TP 会更“值”。

2. Pipeline Parallel（PP）

PP 的核心思想是：按层切模型，每个 worker 负责一段连续的层。

例如：

  • worker 0 负责 layer 0 ~ 15
  • worker 1 负责 layer 16 ~ 31
  • worker 2 负责 layer 32 ~ 47

这样做的优点是：

  • 对卡间带宽的要求通常低于 TP。
  • 当单卡放不下整个模型时，PP 是很自然的拆法。
  • 跨机时往往比细粒度 TP 更稳。

PP 的问题也很典型：

  • 单请求延迟不一定变好，因为本质上还是按 stage 串行流过。
  • 如果 batch 不够大，会出现 pipeline bubble，也就是有的 stage 在等活干。
  • stage 之间负载不均衡时，最慢的 stage 会卡住整条流水线。

所以 PP 更像是一个“容量扩展”和“硬件适配”手段，而不是天然的低时延利器。推理场景下，它常常是“能跑起来”的方案，但未必是“时延最优”的方案。

如果从代码视角看，可以重点看模型层的切分边界。例如在 vLLM 的模型执行路径里，经常能看到类似 start_layer -> end_layer 这样的划分思路。

3. Expert Parallel（EP）

EP 主要服务于 MoE（Mixture of Experts） 模型。

MoE 和普通 dense 模型的最大区别在于：

  • dense 模型里，一层的所有参数通常都会参与本次前向。
  • MoE 里，会先经过 router / gate，再只激活少数几个 expert。

这意味着 expert 可以天然地拆到不同 GPU 上，于是就有了 EP。

典型流程大概是：

  1. 对 token 做路由，决定它该去哪个 expert。
  2. 把 token 按 expert 归类并发到对应 GPU，通常对应一次 all-to-all 或类似 shuffle。
  3. 每个 expert 在本地完成前向。
  4. 把结果再聚合回原来的 token 顺序。

EP 的优势：

  • 对 MoE 很自然，因为不是所有参数都要同时参与计算。
  • 相比 dense 模型（一次计算时，网络中这一层的大部分参数都会参与），MoE 更容易在参数规模上继续扩张。

EP 的难点：

  • expert 负载不均衡，热门 expert 容易变热点。
  • 共享 expert 或高频 expert 可能成为瓶颈，因此有时会复制一份或多份。
  • all-to-all 通信代价大，尤其是跨机时会很痛。
  • token 路由是动态的，不同请求、不同 step 的通信模式可能都不一样。

可以把它记成一句话：

  • TP 更像是在“切算子”
  • EP 更像是在“切专家”

4. Data Parallel（DP）

  • 单靠 TP 能扩展到的规模，远小于 EP 真正需要覆盖的规模。
  • TP 的大小不能超过 attention head 的数量。
  • linear layer 的并行度通常比 attention layer 更高；attention 更容易成为并行瓶颈，所以需要通过并行处理更多 request，来把 attention 这一侧的并行度“抬起来”。

DP 的核心思想最直接：每张卡上放一份完整模型副本，不同请求分发到不同副本处理。

DP 的特点是：

  • 单个请求内部几乎不需要频繁跨卡同步。
  • 更擅长提升整体吞吐，而不是优化单请求时延。
  • 非常适合和其他并行方式叠加，用来横向扩容服务能力。

DP 的优势：

  • 实现思路最直观。
  • 对网络互联的要求通常低于 TP/EP。
  • 很适合多副本服务、弹性扩缩容和故障隔离。

在做 DP（Data Parallel）时，不同副本上的 request 数量或执行进度可能不一致。为了避免某些通信操作卡死，需要用 padding request 把各个副本“对齐”。

DP 的难点：

  • 调度器要尽量把请求分发到“合适的副本”上，而不是盲目轮询。
  • 如果用了 Prefix Cache，最好把相似请求路由到同一个副本，提升 cache locality。
  • 某些严格同步的 SPMD 推理实现里，还需要做请求 padding 或 step 对齐，避免部分 rank 提前结束造成死锁。

一句话总结就是：DP 主要解决的是“吞吐不够”的问题，不直接解决“单请求太慢”的问题。

5. Sequence Parallel vs Context Parallel（补充）

这两个概念经常一起出现，但最好不要混为一谈。

Sequence Parallel（SP）

在 Megatron 这套术语里，Sequence Parallel 往往是 Tensor Parallel 的补充优化。

它主要做的事情不是把整段上下文完整切开去独立跑 attention，而是把一部分沿 sequence 维度天然可分的计算分摊出去，比如：

  • LayerNorm
  • Dropout
  • 一些不需要跨 token 交互的逐位置操作

它的主要价值是：

  • 降低 TP 场景下的激活内存
  • 让 TP 更省显存
  • 通常和 TP 搭配出现，而不是独立作为主并行策略

所以如果严格一点说，SP 更像是 TP 体系里的一个辅助优化。

Context Parallel（CP）

Context Parallel 才更接近“把长序列本身切到多张卡上处理”。

它特别适合下面的场景：

  • 上下文很长
  • prefill 阶段激活和 attention 代价很高
  • 单卡显存对长序列支撑不住

可以把它理解为：
不是把“模型层内部的权重”切开，而是把“很长的 token 序列”分摊到多张卡上处理。

它的价值主要在于：

  • 降低长上下文场景下的激活和中间状态压力
  • 在超长序列 prefill 时更容易扩展
  • 常常和 TP 一起配合使用

它的难点在于：

  • attention 的依赖范围天然覆盖整个上下文，因此需要额外通信来拿到完整的 KV 视图
  • 实现复杂度高于纯 DP
  • 如果序列不够长，收益不一定明显

所以在长上下文模型里，除了 TP / PP / EP / DP，通常也会把 CP 纳入考虑；而 SP 更像是 TP 体系下的配套优化。

并行策略怎么选

可以用一个非常实用的判断框架：

  • 模型放不下：先考虑 TP 或 PP
  • 单请求太慢：优先看 TP，再看算子优化
  • 吞吐不够：优先加 DP
  • 模型是 MoE：优先考虑 EP
  • 长上下文、多轮对话、长 prompt 很多：重点看 PD 分离 + Prefix Cache

不要把分布式推理理解成“卡越多越快”。很多时候，加卡只是在加通信、加调度复杂度，未必在加有效性能。

从代码视角看这些概念

如果继续结合 vLLM 看代码，最好先区分一下 V0 和 V1 视角。

偏模型实现、跨版本都值得看

  • vllm/model_executor/models/llama.py
  • 看模型层是怎么组织的
  • 看 attention / MLP 的前向路径
  • vllm/distributed/...
  • 看通信组、通信算子、device communicator 的抽象

更偏 V0 的阅读入口

  • vllm/attention/backends/flash_attn.py
  • vllm/worker/model_runner.py

如果你这篇笔记本来就混着一些 V0 视角，这两个路径继续保留是合理的。

更偏 V1 的阅读入口

  • vllm/v1/attention/backends/...
  • vllm/v1/core/sched/...

如果是看当前版本的主调度路径、attention backend 和 engine 行为，通常更应该往 v1/ 目录里读。

比较好的阅读顺序是：

  1. 先看单机单卡的前向主路径。
  2. 再看 TP / PP / EP 是在哪些边界插入通信。
  3. 最后再看调度、KV Cache、Prefix Cache 这些系统级优化。

小结

分布式推理优化不是某个单点技巧，而是一整套系统工程：

  • 模型放不下，用并行策略解决
  • 算得慢，用算子和并行解决
  • prompt 太长，用 Prefix Cache 和 PD 分离解决
  • decode 太贵，用 Speculative Decoding 解决
  • 系统到底有没有变好，用严谨的 Benchmark 验证

真正高水平的推理系统，拼的不是某一个技巧，而是：

  • 通信是否足够高效
  • 调度是否足够聪明
  • 缓存是否足够可复用
  • 指标是否真的贴近线上业务
Created: April 29, 2026
Last update: April 29, 2026

Discussion