Context Parallel Deployment

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

Context parallel은 주로 긴 컨텍스트(long context) 요청을 서빙하는 문제를 해결해요. prefill과 decode는 특성이 매우 다르고 SLO(서비스 수준 목표)도 다르기 때문에, 두 단계를 별도로 구현해야 해요. 주요 고려 사항은 긴 컨텍스트 prefill에서는 TTFT(time to first token)를 prefill 연산 시간을 쿼리 토큰에 걸쳐 분산(amortize)시켜 제어해야 하고, 긴 컨텍스트 decode에서는 KV 캐시 공간을 늘려 배치 크기(그리고 처리량)를 키워야 한다는 점이에요.

출처: 문서

본문

Prefill Context Parallel

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

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

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

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

Decode Context Parallel

디코딩의 자기회귀(auto-regressive) 특성 때문에 모든 디코딩 단계는 페이징된 KV 캐시에 저장된 많은 수의 key/value 토큰에 대해 소량의 쿼리 토큰을 계산해야 해요. decode context parallel의 핵심은 KV 캐시를 GPU 간에 어떻게 샤딩(sharding)하느냐예요.

engine_client.vllm_config.cache_config.effective_attention_block_size는 DCP를 포함한 초기화된 full-attention 블록 크기를 토큰 수로 보고해요. DCP=4에서 16 토큰 물리 블록은 64 토큰을 나타내요. 같은 값은 gRPC Control.GetServerInfo.effective_attention_block_size로도 사용할 수 있어요. 물리 블록 크기 필드는 기존 의미를 유지해요. 공통 full-attention 크기가 없거나 스케줄러가 노출하지 않거나 엔진이 이전 버전이면 새 값은 사용 불가능해요(Python에서 None, gRPC에서 부재). 부분 캐시 이벤트는 자신의 실제 크기를 전달해요.

H개의 kv-head를 가진 모델에서 T개 토큰의 컨텍스트를 가진 요청은 KV 캐시에 H * T개의 key/value 텐서를 저장해야 해요.

  1. 한 GPU가 모두 감당할 수 있고 성능이 충분하다면 병렬화가 필요 없어요.
  2. 한 GPU가 모두 감당할 수 없거나 KV 캐시에 더 많은 요청을 담고 싶다면, 먼저 KV 캐시를 H 차원을 따라 샤딩해요. 이것이 일반적인 tensor parallel 샤딩이며, 커맨드라인에 -tp <num_gpus>를 추가하는 것만으로 충분해요.
  3. H는 모델 아키텍처에 의해 결정되는 유한한 값이므로, tensor parallel 크기를 계속 늘리면 각 GPU의 KV 캐시가 tp_size / H 배로 중복돼요. 물론 중복은 효율에 좋지 않아요. 그래서 decode context parallel을 추가해 KV 캐시를 T 차원을 따라 더 샤딩해야 해요. 이것은 커맨드라인에 -dcp <size>를 추가하는 것만으로 충분해요. size는 실행할 GPU 수를 늘리는 것이 아니라 KV 캐시 중복을 줄일 뿐이에요. dcp 크기는 [1, tp_size/H] 범위여야 해요. dcp 크기가 클수록 KV 캐시 중복은 줄지만 통신 오버헤드가 늘어나요.

이론적으로 dcp 크기를 tp_size / H를 넘겨 KV 캐시를 더 샤딩하고 디코딩 단계를 가속화하는 것도 가능해요. 하지만 디코딩에서는 쿼리 토큰 수가 제한적이므로, 남은 dcp_size - tp_size / H개의 GPU를 non-attention 레이어에 대해 어떻게 할지 불분명해요. 단순성을 위해 dcp 크기는 tp_size / H로 상한을 정해요. 디코딩 단계를 더 가속화하려면 먼저 tp_size를 늘리고, 그다음 dcp 크기를 늘리는 것을 고려할 수 있어요.

KV 캐시는 디코딩 중에 커질 수 있으므로 샤딩 전략을 신중히 구현해야 해요. vLLM은 T 차원을 따라 KV 캐시를 샤딩하는 인터리빙(interleaving) 전략을 사용해 미래 토큰의 KV 캐시도 T 차원을 따라 자연스럽게 샤딩되게 해요. 이는 Moonshot의 Chao Hong이 제안했고 이 논문에 자세히 설명돼 있어요.

사례 연구

  • DeepSeek-R1: MLA가 활성화되면 kv-head가 1개예요. 일반적인 단일 노드 배포인 -tp 8은 8배 KV 캐시 중복을 발생시켜요. -dcp 8을 추가해 KV 캐시 중복을 줄일 수 있어요.
  • Kimi-K2: DeepSeek-R1과 유사한 아키텍처지만 파라미터가 더 많아요. -tp 16으로 배포하면 KV 캐시 중복이 16배예요. -dcp 16을 추가하면 더 많은 통신 오버헤드의 대가로 KV 캐시 중복을 완전히 없앨 수 있어요. -dcp 8을 추가하면 KV 캐시 중복을 2배로 줄일 수도 있는데, DCP 통신이 한 노드 안에서만 발생하므로 통신 오버헤드는 더 작아요.
  • Qwen3-235B-A22B: kv-head가 4개예요. -tp 8로 배포하면 KV 캐시 중복이 2배예요. -dcp 2를 추가하면 중복을 없앨 수 있어요.

요약하면 decode context parallel에서는 만족스러운 성능이 나올 때까지 -tp 크기를 늘린 다음, -dcp를 추가해 KV 캐시 중복을 줄이세요.

Decode context parallel은 MLA와 GQA 모델 모두에서 vLLM에서 지원돼요. 일부 attention 백엔드는 decode context parallel과 MTP(다중 토큰 예측)의 결합을 지원해 디코딩 단계를 더 가속화할 수 있어요.

기술 논의

주요 논의는 vLLM Slack의 #sig-context-parallel 채널에서 이루어져요.

더 알아보기 (Learn more)