跳到内容

上下文并行部署

上下文并行主要解决长上下文请求的服务问题。由于 prefill 和 decode 具有截然不同的特性和非常不同的 SLO（服务水平目标），我们需要为它们分别实现上下文并行。主要考虑因素是：

  • 对于长上下文 prefill，我们需要通过分摊 prefill 计算时间到查询（query）tokens 来控制 TTFT（首次 token 的时间）。
  • 对于长上下文 decode，我们需要更多的 KV 缓存空间来增加 batchsize（从而提高吞吐量）。

Prefill 上下文并行

在 prefill 过程中，对于一个包含 T 个新 token 的长请求，我们需要计算这些新 token 的查询/键/值（query/key/value）张量。假设我们有 N 个 GPU，我们可以将请求分割成 N 块，每个 GPU 计算其负责的查询/键/值张量的其中一块。

根据用例的不同，有两种可能的策略：

  1. 部分查询，完整键/值：如果请求 token 长度适中（我们可以负担存储完整的键/值张量），目标是加速 prefill（并分摊 prefill 计算时间到查询 tokens），那么我们可以从所有 GPU 收集键/值张量，让每个 GPU 计算对应于其分块的查询 tokens 的注意力输出。
  2. 部分查询，部分键/值：如果请求 token 长度太长，我们无法负担存储完整的键/值张量，那么我们只能为每个 GPU 计算查询/键/值张量的一个分块，并使用 ring-attention 等技术逐块发送/接收（send/recv）键/值张量。

这两种方法都在积极开发中。

Decode 上下文并行

由于解码的自回归性质，每个解码步骤都需要计算少量查询 token 相对于分页 KV 缓存中大量键/值 token 的结果。 Decode 上下文并行的核心是如何在 GPU 之间分片（shard）KV 缓存。

对于一个具有 H 个 kv-head 的模型，一个上下文长度为 T 的请求需要在 KV 缓存中存储 H * T 个键/值张量。

  1. 如果一个 GPU 可以容纳所有这些，并且性能足够好，那么就不需要并行化。
  2. 如果一个 GPU 无法容纳所有这些，或者我们希望在 KV 缓存中容纳更多请求，我们可以首先沿着 H 维度对 KV 缓存进行分片，这就是普通的张量并行分片。这很简单，只需在命令行中添加 -tp <num_gpus>。
  3. 由于 H 是有限的（由模型架构决定），当我们继续增加张量并行大小（tensor parallel size）时，每个 GPU 的 KV 缓存会重复 tp_size / H 次。当然，重复不利于效率。然后我们需要添加 decode 上下文并行，沿着 T 维度进一步分片 KV 缓存。这很简单，只需在命令行中添加 -dcp <size>。请注意，size 不会增加我们需要的 GPU 数量，而只是减少 KV 缓存的重复。 dcp 大小应在 [1, tp_size/H] 的范围内。 dcp 越大，KV 缓存的重复次数越少，但通信开销会增加。

理论上，可以将 dcp 大小扩展到 tp_size / H 以上，以进一步分片 KV 缓存并加速解码阶段。然而，由于解码过程中查询 token 的数量有限，对于剩余的 dcp_size - tp_size / H 个 GPU 的非注意力层该怎么处理尚不清楚。为了简单起见，dcp 大小上限为 tp_size / H。如果您想进一步加速解码阶段，可以考虑先增加 tp_size，然后再增加 dcp 大小。

请注意，KV 缓存会在解码过程中增长，分片策略需要仔细实现。我们采用交错策略（interleaving strategy）沿 T 维度分片 KV 缓存，这样未来 token 的 KV 缓存就可以自然地沿 T 维度分片。这由 Moonshot 的 Chao Hong 提出，并在 这篇论文 中得到了详细解释。

案例研究

对于 DeepSeek-R1，当启用 MLA 时，我们有 1 个 kv-head。典型的单节点部署 -tp 8 会导致 8 倍的 KV 缓存重复。我们可以考虑添加 -dcp 8 来减少 KV 缓存的重复。

对于 Kimi-K2，其架构与 DeepSeek-R1 类似，但参数更多。当以 -tp 16 部署时，KV 缓存重复为 16 倍。我们可以添加 -dcp 16 来完全消除 KV 缓存的重复，但这会增加通信开销。我们也可以添加 -dcp 8 来将 KV 缓存的重复减少到 2 倍。虽然仍然会复制两次 KV 缓存，但通信开销较小，因为 DCP 通信仅发生在节点内部。

对于 Qwen3-235B-A22B，我们有 4 个 kv-head。当以 -tp 8 部署时，KV 缓存重复为 2 倍。然后我们可以添加 -dcp 2 来消除 KV 缓存的重复。

总之，对于 decode 上下文并行，尝试增加 -tp 大小直到获得满意的性能，然后添加 -dcp 来减少 KV 缓存的重复。

vLLM 支持 decode 上下文并行，支持 MLA 和 GQA 模型。一些注意力后端也支持 decode 上下文并行与 MTP（多 token 预测）的组合，以进一步加速解码阶段。

技术讨论

主要讨论发生在 vLLM Slack 的 #sig-context-parallel 频道。