vLLM Decode Context Parallelism for Long Context Workloads

TL;DR

vLLM은 KV 캐시를 단순히 어텐션 헤드(attention head)별로 나누는 것이 아니라 시퀀스 차원을 따라 GPU 간에 샤딩하는 기술인 Decode Context Parallelism (DCP)를 구현했습니다. 이를 통해 Grouped-Query Attention (GQA) 및 Multi-head Latent Attention (MLA) 모델에서 KV 캐시 중복을 방지하여, 긴 문맥(long-context) 추론 시 요청 동시성(concurrency)과 처리량(throughput)을 크게 향상시킬 수 있습니다.

The Memory Wall in Long-Context Inference

표준 텐서 병렬 처리(Tensor Parallelism, TP)는 KV 캐시를 어텐션 헤드별로 분할합니다. 이 방식은 GPU의 수가 KV 헤드의 수보다 많아질 경우, 시스템이 여러 GPU에 KV 캐시를 중복해서 복제해야 하므로 메모리 병목 현상이 발생합니다.

Grouped-Query Attention (GQA): TP는 GPU당 하나의 헤드까지만 KV 캐시를 분할할 수 있습니다. TP 크기가 KV 헤드 수를 초과하면 캐시가 중복됩니다.
Multi-head Latent Attention (MLA): MLA는 KV를 모든 쿼리 헤드와 공유되는 단일 저차원 잠재 벡터(low-rank latent vector)로 압축합니다. MLA는 사실상 하나의 KV 헤드만 가지므로, 잠재 KV 캐시는 모든 TP 랭크(rank)에 전체 복제되어 GPU 메모리를 심각하게 제한합니다.

이러한 중복은 중요한 GPU 메모리를 소모하여 동시 요청 수를 제한하고 토큰당 비용을 증가시킵니다.

Decode Context Parallelism (DCP) Explained

Decode Context Parallelism은 KV 캐시를 시퀀스(문맥) 차원을 따라 GPU 간에 분할하여 메모리 병목 현상을 해결합니다. 각 GPU가 전체 시퀀스에 대해 특정 헤드의 캐시 전체를 보유하는 대신, 동일한 시퀀스에 대해 각 GPU가 토큰 위치의 일부 청크(chunk)를 담당합니다.

예를 들어, 200K-토큰 요청의 경우, DCP를 사용하는 4개의 GPU는 GPU 0이 0–50K 토큰을, GPU 1이 50K–100K 토큰을 보유하는 방식으로 워크로드를 분할할 수 있습니다.

The DCP Execution Process

DCP는 디코딩 단계에서 정확성을 유지하기 위해 특정 통신 리듬을 따릅니다:

AllGather Q: 어텐션은 전체 쿼리 벡터를 필요로 하므로, all-gather 연산을 통해 모든 GPU에 쿼리의 완전한 복사본을 구성합니다. (MLA 모델의 경우, VLLM_DCP_Q_REPLICATE=1을 통해 로드 시점에 쿼리 프로젝션을 복제하도록 최적화할 수 있습니다).
Compute: 각 GPU는 수집된 쿼리와 로컬 KV 캐시 슬라이스(slice) 간의 어텐션을 수행합니다.
AllGather + ReduceScatter: 부분 결과와 Log-Sum-Exp (LSE) 값은 AllGather를 통해 공유됩니다. LSE 값은 online-softmax 트릭을 사용하여 부분 결과를 재가중치화하고 병합하는 데 사용되며, ReduceScatter는 이들을 합산하면서 각 GPU에 관련 헤드 슬라이스만 반환합니다.

Performance Benchmarks

vLLM은 8×B200 노드에서 NVFP4를 사용하는 Kimi K2.6 모델을 대상으로, 중간 입력값이 ~67k 토큰인 에이전트형 긴 문맥 트레이스(agentic long-context trace)를 사용하여 DCP와 기본 TP 배포를 비교 테스트했습니다.

Throughput and Concurrency

DCP는 기본 TP에 비해 훨씬 높은 동시성과 처리량을 유지합니다:

Baseline TP: 동시성 64에서 메모리 사용률 100%에 도달하며, 처리량은 GPU당 약 1,863 tok/s/GPU 근처에서 정체됩니다.
DCP: 동시성을 512까지 확장할 수 있으며, KV 메모리를 82%만 사용하면서 GPU당 6,091 tok/s/GPU에 도달합니다.

Stability Across Sequence Lengths

문맥 길이가 길어질수록 성능이 안정적으로 유지됩니다. 200k+ 토큰 범위에서는 DCP가 높은 처리량-상호작용성 경계(throughput-interactivity frontier)를 유지하는 반면, 복제된 KV 방식의 기본 TP는 메모리 부족으로 인해 확장이 불가능한 경우가 많습니다.

Implementation and Usage

DCP는 --decode-context-parallel-size 인자를 통해 활성화할 수 있습니다.

Multi-head Latent Attention (MLA) Backend

DeepSeek-V2, V3, R1, 및 Kimi K2.6과 같은 모델에서 사용됩니다. MLA는 사실상 하나의 KV 헤드를 가지므로, 시퀀스는 전체 TP 크기만큼 분할할 수 있습니다.

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 헤드별로 분할합니다. 그 다음 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 크기에 대한 더 정밀한 제어력을 제공하여 과잉 프로비저닝을 줄입니다.
Communication Kernels: 단일 및 다중 노드 설정에서 연산 중첩(compute overlap)을 개선하기 위해 더 나은 all-to-all (A2A) 커널을 개발합니다.
Speculative Decoding: 지연 시간(latency) 이점을 유지하기 위해 DCP를 MTP와 speculative decoding과 통합합니다.
Disaggregated Serving: prefill/decode (P/D) 분리 서빙(disaggregation) 지원을 강화합니다.
Expanded Model Support: GLM-5.2, Kimi K3 지원을 확장하고 Prefill Context Parallelism (PCP)을 개발합니다.