Skip to content

上下文并行部署

上下文并行主要解决长上下文请求的服务问题。由于预填充（prefill）和解码（decode）具有截然不同的特性，并且有着不同的服务等级目标（SLO），因此我们需要分别为它们实现上下文并行。主要的考虑因素包括：

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

预填充上下文并行

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

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

  1. 部分 query，完整 key/value：如果请求的 token 长度适中（我们可以承受保存完整的 key/value 张量），且目标是加速预填充（并将预填充的计算时间分摊到查询 token 上），那么我们可以从所有 GPU 收集 key/value 张量，并让每个 GPU 计算与其 query token 块对应的注意力输出。
  2. 部分 query，部分 key/value：如果请求的 token 长度过长，我们无法再承受保存完整的 key/value 张量，那么每个 GPU 只能计算一个块的 query/key/value 张量，并使用如 ring-attention 等技术来逐块发送/接收 key/value 张量。

这两种方法目前正在积极开发中。

解码上下文并行

由于解码的自回归特性，每一步解码都需要针对存储在分页 KV 缓存中的大量 key/value token 计算少量 query token。解码上下文并行的核心在于如何将 KV 缓存分片到多个 GPU 上。

对于一个具有 H 个 kv-head 的模型，一个包含 T 个 token 的请求需要在 KV 缓存中存储 H * T 个 key/value 张量。

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

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

请注意，KV 缓存在解码过程中会增长，因此分片策略需要仔细实现。我们使用交错策略沿 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 缓存复制。

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

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

技术讨论

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