컨텍스트 병렬 배포

컨텍스트 병렬 배포 (Context Parallel Deployment)

긴 컨텍스트 요청을 서빙할 때는 메모리와 연산이 턱없이 커지는 문제가 생겨요. 컨텍스트 병렬(context parallel)은 바로 이 긴 컨텍스트 요청을 여러 GPU에 나눠 처리하는 기법이에요. 이 페이지에서는 prefill과 decode 각각에서 컨텍스트 병렬이 어떻게 동작하는지 살펴볼게요.

출처: vLLM 공식 문서 — serving/context_parallel_deployment

컨텍스트 병렬은 주로 긴 컨텍스트 요청 서빙 문제를 해결해요. prefill과 decode는 특성이 상당히 다르고 SLO(서비스 수준 목표)도 달라서, 둘을 따로 컨텍스트 병렬로 구현해야 해요. 주요 고려사항은 이래요.

  • 긴 컨텍스트 prefill의 경우, prefill 연산 시간을 쿼리 토큰들 사이에 분산해 TTFT(time to first token)를 제어해야 해요.
  • 긴 컨텍스트 decode의 경우, KV cache 공간을 더 확보해 배치 크기(그리고 처리량)를 늘려야 해요.

Prefill 컨텍스트 병렬 (Prefill Context Parallel)

prefill 중에 새 토큰이 T개인 긴 요청이 있으면, 이 새 토큰들의 query/key/value 텐서를 계산해야 해요. GPU가 N개 있다면 요청을 N개 청크로 나눠서 각 GPU가 query/key/value 텐서의 한 청크를 계산할 수 있어요.

사용 사례에 따라 두 가지 전략이 가능해요.

  1. 부분 query, 전체 key/value: 요청 토큰 길이가 적당히 길고(전체 key/value 텐서를 유지할 여유가 있고), prefill을 가속화(그리고 prefill 연산 시간을 쿼리 토큰들 사이에 분산)하는 게 목표라면, 모든 GPU에서 key/value 텐서를 모아서 각 GPU가 자기 청크의 query 토큰에 해당하는 어텐션 출력을 계산하게 해요.
  2. 부분 query, 부분 key/value: 요청 토큰 길이가 너무 길어서 전체 key/value 텐서를 유지할 수 없다면, 각 GPU가 query/key/value 텐서의 한 청크만 계산하고 ring-attention 같은 기법으로 key/value 텐서를 청크 단위로 주고받아요.

두 접근 모두 활발히 개발 중이에요.

Decode 컨텍스트 병렬 (Decode Context Parallel)

디코딩은 자동회귀(auto-regressive) 특성 때문에, 매 디코딩 스텝마다 큰 수의 key/value 토큰(페이지드 KV cache에 저장됨)에 대해 적은 양의 query 토큰을 계산해야 해요. decode 컨텍스트 병렬의 핵심은 KV cache를 GPU들 사이에 어떻게 샤딩하느냐예요.

kv-head가 H개인 모델에서, 컨텍스트에 T개 토큰이 있는 요청은 KV cache에 H * T개의 key/value 텐서를 저장해야 해요.

  1. GPU 하나가 모두 담을 수 있고 성능도 충분하다면, 병렬화가 필요 없어요.
  2. GPU 하나가 모두 담을 수 없거나 KV cache에 더 많은 요청을 담고 싶다면, 먼저 H 차원을 따라 KV cache를 샤딩해요. 이게 단순 tensor parallel 샤딩이에요. 명령줄에 -tp <num_gpus>를 추가하기만 하면 돼요.
  3. H는 (모델 아키텍처에 의해 결정되는) 제한된 값이라, tensor parallel 크기를 계속 늘리면 각 GPU의 KV cache가 tp_size / H배로 중복돼요. 물론 중복은 효율에 좋지 않아요. 그래서 decode 컨텍스트 병렬을 추가해 T 차원을 따라 KV cache를 더 샤딩하는 거예요. 명령줄에 -dcp <size>를 추가하기만 하면 돼요. size가 GPU 수를 늘리는 게 아니라 KV cache 중복만 줄인다는 점에 주의하세요. dcp 크기는 [1, tp_size/H] 범위에 있어야 해요. dcp 크기가 클수록 KV cache 중복은 줄지만 통신 오버헤드는 늘어나요.

이론적으로는 dcp 크기를 tp_size / H보다 늘려 KV cache를 더 샤딩하고 decode 단계를 가속화하는 것도 가능해요. 하지만 decoding에서 query 토큰 수가 제한돼 있어서, 남는 dcp_size - tp_size / H개의 GPU를 비-어텐션 레이어에 어떻게 쓸지 불분명해요. 단순함을 위해 dcp 크기는 tp_size / H로 상한을 두었어요. decode 단계를 더 가속화하려면 tp_size부터 늘리고, 그 다음 dcp 크기를 늘리는 걸 고려하세요.

디코딩 중 KV cache는 커질 수 있으므로 샤딩 전략을 신중히 구현해야 해요. vLLM은 T 차원을 따라 interleaving 전략으로 KV cache를 샤딩해서, 미래 토큰용 KV cache가 T 차원을 따라 자연히 샤딩되게 해요. 이 방식은 Moonshot의 Chao Hong이 제안했고 이 논문에 자세히 설명돼 있어요.

사례 연구를 볼게요.

  • DeepSeek-R1: MLA가 켜져 있으면 kv-head가 1개예요. -tp 8로 단일 노드 배포하면 KV cache가 8배 중복돼요. -dcp 8을 추가해 중복을 줄일 수 있어요.
  • Kimi-K2: DeepSeek-R1과 비슷하지만 파라미터가 더 많아요. -tp 16으로 배포하면 KV cache가 16배 중복돼요. -dcp 16을 추가하면 통신 오버헤드가 커지는 대신 KV cache 중복을 완전히 없앨 수 있어요. -dcp 8을 추가하면 KV cache 중복을 2배로 줄일 수 있어요. 여전히 2배 중복이지만 DCP 통신이 한 노드 안에서만 일어나므로 통신 오버헤드가 더 작아요.
  • Qwen3-235B-A22B: kv-head가 4개예요. -tp 8로 배포하면 KV cache 중복이 2배예요. -dcp 2를 추가하면 KV cache 중복을 없앨 수 있어요.

요약하면, decode 컨텍스트 병렬에서는 -tp 크기를 만족스러운 성능이 나올 때까지 늘리고, 그 다음 -dcp를 추가해 KV cache 중복을 줄이세요.

decode 컨텍스트 병렬은 vLLM에서 MLA 모델과 GQA 모델 모두 지원해요. 일부 어텐션 백엔드는 decode 컨텍스트 병렬과 MTP(multi-token prediction)의 조합도 지원해 decode 단계를 더 가속화해요.

기술 논의 (Technical Discussions)

주요 논의는 vLLM Slack의 #sig-context-parallel 채널에서 진행돼요.

더 알아보기 (Learn more)