vLLM Decode Context Parallelism for Long Context Workloads
TL;DR
vLLM 实现了 Decode Context Parallelism (DCP),这是一种沿着序列维度而非仅仅通过注意力头来分片 KV cache 的技术。这可以防止 Grouped-Query Attention (GQA) 和 Multi-head Latent Attention (MLA) 模型中的 KV cache 重复,从而在长上下文推理过程中实现显著更高的请求并发量和吞吐量。
The Memory Wall in Long-Context Inference
标准的 Tensor Parallelism (TP) 通过注意力头来划分 KV cache。当 GPU 数量超过 KV heads 的数量时,这种方法会造成内存瓶颈,因为系统必须在多个 GPU 之间重复 KV cache。
- Grouped-Query Attention (GQA): TP 只能将 KV cache 拆分到每个 GPU 一个头。一旦 TP 超过了 KV heads 的数量,cache 就会被复制。
- Multi-head Latent Attention (MLA): MLA 将 KV 压缩成一个在所有 query heads 之间共享的单一低秩潜向量 (low-rank latent vector)。由于它实际上只有一个 KV head,潜 KV cache 会在每个 TP rank 中完整地复制,严重限制了用于并发请求的可用 GPU 内存。
这种重复会消耗关键的 GPU 内存,限制了并发请求的数量并增加了每个 token 的成本。
Decode Context Parallelism (DCP) Explained
Decode Context Parallelism 通过按序列(上下文)维度在 GPU 之间拆分 KV cache 来解决内存瓶颈。与其让每个 GPU 为整个序列持有特定头部的 cache 的完整副本,不如让每个 GPU 负责同一序列的一块 token 位置。
例如,在一个 200K-token 的请求中,使用 DCP 的四个 GPU 可能会这样拆分工作负载:GPU 0 持有 tokens 0–50K,GPU 1 持有 50K–100K,依此类推。
The DCP Execution Process
DCP 遵循特定的通信节奏以在 decode 阶段保持正确性:
- AllGather Q: 由于注意力机制需要完整的 query vector,all-gather 操作会在每个 GPU 上组装出 query 的完整副本。 (对于 MLA 模型,这可以通过
VLLM_DCP_Q_REPLICATE=1进行优化,以便在加载时复制 query projections)。 - Compute: 每个 GPU 在其收集到的 query 与其本地的 KV cache 切片之间执行注意力计算。
- AllGather + ReduceScatter: 部分结果和 Log-Sum-Exp (LSE) 值通过 AllGather 共享。LSE 值用于使用 online-softmax 技巧进行重加权和合并部分结果,而 ReduceScatter 则在求和的同时仅向每个 GPU 返回相关的 head-slice。
Performance Benchmarks
vLLM 测试了 DCP 与基于 TP 的基准部署,使用 Kimi K2.6 在 NVFP4 模式下在 8×B200 节点上进行测试,利用了具有中位数输入约为 ~67k tokens 的 agentic long-context trace。
Throughput and Concurrency
DCP 与基准 TP 相比,能够维持显著更高的并发量和吞吐量:
- Baseline TP: 在并发量为 64 时达到 100% 内存使用率,吞吐量在 1,863 tok/s/GPU 附近达到平台期。
- DCP: 扩展到 512 的并发量,在仅使用 82% KV 内存的情况下达到 6,091 tok/s/GPU。
Stability Across Sequence Lengths
即使随着上下文长度的增加,性能也保持稳定。在 200k+ token 范围内,DCP 维持了高吞吐量-交互性前沿,而复制-KV 的基准线往往会因为内存不足而无法扩展。
Implementation and Usage
DCP 通过 --decode-context-parallel-size 参数启用。
Multi-head Latent Attention (MLA) Backend
用于 DeepSeek-V2, V3, R1, 和 Kimi K2.6 等模型。由于 MLA 实际上只有一个 KV head,序列可以拆分到完整的 TP degree。
- Constraint:
tensor_parallel_size >= decode_context_parallel_size且tensor_parallel_size % decode_context_parallel_size == 0。
Grouped-Query Attention (GQA) Backend
用于 Qwen3-235B 和 Llama-family 模型。TP 首先按 KV heads 拆分;DCP 随后将剩余的冗余副本进行分片。
- Constraint:
(tensor_parallel_size // num_key_value_heads) >= decode_context_parallel_size且(tensor_parallel_size // num_key_value_heads) % decode_context_parallel_size == 0。
Future Roadmap
vLLM 计划通过以下开发来扩展 DCP:
- Finer-grained parallelism: 更加精确地控制 TP 和 DCP size 从而减少过度配置。
- Communication Kernels: 为单节点和多节点设置开发更好的 all-to-all (A2A) kernels 以提高计算重叠度。
- Speculative Decoding: 将 DCP 与 MTP 和 speculative decoding 集成,以维持延迟优势。
- Disaggregated Serving: 强化对 prefill/decode (P/D) 分离式服务的支持。
- Expanded Model Support: 扩展支持到 GLM-5.2, Kimi K3, 以及开发 Prefill Context Parallelism (PCP)。
Sources
相关
- Dispatch
- 项目
- Dispatch
- 项目
- Dispatch