跳到内容

上下文并行部署

上下文并行主要用于解决长上下文请求的服务问题。由于预填充（Prefill）和解码（Decode）阶段表现出截然不同的特性，且具有不同的服务水平目标（SLO），我们需要分别为它们实现上下文并行。主要考虑因素如下：

  • 对于长上下文预填充，我们需要通过将计算时间分摊到各个查询（query）token 上，从而控制首字延迟（TTFT）。
  • 对于长上下文解码，我们需要更多的 KV cache 空间来增加批处理大小（从而提高吞吐量）。

预填充上下文并行

在预填充阶段，对于一个包含 T 个新 token 的长请求，我们需要计算这些新 token 的查询/键/值（query/key/value）张量。假设我们有 N 个 GPU，我们可以将请求拆分为 N 个块，每个 GPU 计算其中一块查询/键/值张量。

根据使用场景，有两种可能的策略：

  1. 部分查询，完整键/值：如果请求的 token 长度中等（我们可以负担得起存储完整的键/值张量），且目标是加速预填充（并将预填充的计算时间分摊到查询 token 上），那么我们可以从所有 GPU 收集键/值张量，并让每个 GPU 计算其自身查询块对应的注意力输出。
  2. 部分查询，部分键/值：如果请求的 token 长度过长，以至于无法容纳完整的键/值张量，那么每个 GPU 只能计算一块查询/键/值张量，并使用诸如 ring-attention 的技术，以分块方式发送/接收键/值张量。

以上两种方法目前都在积极开发中。

解码上下文并行

由于解码的自回归特性，每一个解码步骤都需要针对存储在分页 KV cache 中的大量键/值 token 计算少量的查询 token。解码上下文并行的核心在于如何跨 GPU 对 KV cache 进行分片。

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

  1. 如果单个 GPU 能完全装下这些数据，且性能足够好，则无需并行化。
  2. 如果单个 GPU 装不下，或者我们希望在 KV cache 中存储更多请求，我们可以首先沿 H 维度对 KV cache 进行分片，这就是标准的张量并行（Tensor Parallel）分片。只需在命令行添加 -tp <num_gpus> 即可。
  3. 由于 H 是有限的（由模型架构决定），当我们继续增加张量并行大小时，每个 GPU 上的 KV cache 将会重复 tp_size / H 次。当然，重复对效率不利。因此，我们需要引入解码上下文并行（DCP）来进一步沿 T 维度对 KV cache 进行分片。这只需在命令行添加 -dcp <size> 即可。请注意，size 不会增加我们需要启动的 GPU 数量，而只是减少了 KV cache 的重复。DCP 大小应在 [1, tp_size/H] 范围内。DCP 大小越大，KV cache 的重复就越少，但通信开销会增加。

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

请注意，KV cache 在解码过程中会增长，因此需要仔细实现分片策略。我们使用一种交错策略（interleaving strategy）沿 T 维度对 KV cache 进行分片，使得未来 token 的 KV cache 可以自然地沿 T 维度进行分片。该方案由 月之暗面（Moonshot）的 Chao Hong 提出，并在 这篇论文 中进行了详细解释。

案例研究

对于 DeepSeek-R1，当开启 MLA 时，它有 1 个 kv-head。典型的单节点 -tp 8 部署会导致 8 倍的 KV cache 重复。我们可以考虑添加 -dcp 8 来消除这种重复。

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

对于 Qwen3-235B-A22B，它有 4 个 kv-head。当我们使用 -tp 8 部署时，KV cache 重复为 2 倍。此时，我们可以添加 -dcp 2 来消除 KV cache 重复。

简而言之，对于解码上下文并行，请先尝试增加 -tp 大小直到获得满意的性能，然后再添加 -dcp 来减少 KV cache 的重复。

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

技术讨论

主要讨论在 vLLM Slack 的 #sig-context-parallel 频道中进行。