[vllm] vLLM의 대규모 모델 추론을 위한 MRV2 기반 Prefill Context Parallelism(PCP) 도입

PR 링크: vllm-project/vllm#46570 상태: Merged | 변경: +1204 / -196

들어가며

최신 대규모 언어 모델(LLM)은 긴 컨텍스트를 처리하기 위해 많은 메모리와 연산 자원을 요구합니다. 특히 Multi-Head Latent Attention(MLA)을 사용하는 모델의 경우, Prefill 단계에서의 병목 현상이 심각합니다. 이번 vLLM 업데이트에서는 MRV2(Model Runner V2) 아키텍처를 기반으로 PCP(Prefill Context Parallelism)를 도입하여, Prefill 연산을 여러 GPU에 분산 처리함으로써 TTFT(Time To First Token)를 획기적으로 단축했습니다.

주요 변경 사항

1. PCPManager 도입 (vllm/v1/worker/gpu/pcp_manager.py)

가장 핵심적인 변화는 PCPManager의 추가입니다. 이 클래스는 전역적으로 스케줄링된 InputBatch를 각 PCP 랭크가 처리할 수 있는 로컬 가상 배치로 분할합니다.

Before:

# 기존에는 단일 랭크에서 전체 Prefill을 처리하거나 단순 분할만 수행
input_batch = scheduler.get_batch()

After:

# PCPManager가 DualChunkSwap 전략을 사용하여 배치를 로컬 랭크에 맞게 분할
self.pcp_manager = PCPManager(pcp_size=self.pcp_size)
local_batch = self.pcp_manager.partition(global_batch)

이 방식은 DualChunkSwap 알고리즘을 사용하여, 각 랭크가 전체 Prefill의 일부를 담당하되, 인과적 주의(Causal Attention) 연산의 균형을 맞추기 위해 '앞쪽 청크'와 '뒤쪽 청크'를 쌍으로 할당합니다.

2. PCP와 DCP의 직교화 (vllm/distributed/parallel_state.py)

기존에는 PCP와 DCP(Data/KV Parallelism)가 혼재되어 KV 캐시 소유권 계산이 복잡했습니다. 이번 변경으로 두 축을 완전히 분리했습니다.

• TP(Tensor Parallelism): 모델 헤드 분할
• PCP: Prefill 연산 분할
• DCP: KV 캐시 소유권 분할

이제 KV 샤드 수는 PCP × DCP가 아닌 오직 DCP에 의해서만 결정됩니다. 이는 vllm/v1/simple_kv_offload/test_scheduler.py의 테스트 코드 변경에서 명확히 드러납니다.

Before:

cp_world_size = dcp_world_size * pcp_world_size
virtual_block_size = BLOCK_SIZE * cp_world_size

After:

# DCP 소유권만 고려하여 블록 크기 계산
virtual_block_size = BLOCK_SIZE * dcp_world_size

왜 이게 좋은가

이 최적화는 특히 긴 컨텍스트(100k 토큰 이상)에서 강력한 성능을 발휘합니다. 4xB300 GPU 환경에서 GLM-5.2-NVFP4 모델을 테스트한 결과, TP1+PCP4+EP4 구성에서 100k 토큰의 TTFT가 2.859초로 측정되었습니다. 이는 기존의 단순 TP 구성 대비 2배 이상의 성능 향상입니다.

교훈:

1. 연산과 메모리의 분리: Prefill 연산(PCP)과 KV 캐시 저장(DCP)을 분리함으로써, 메모리 용량 제한 없이 연산 성능을 확장할 수 있습니다.
2. DualChunkSwap의 효율성: 단순히 시퀀스를 N등분하는 것보다, 앞뒤 청크를 섞어 할당함으로써 GPU 간의 연산 부하 불균형(Load Imbalance)을 최소화할 수 있습니다.

리뷰 피드백 반영

리뷰 과정에서 GirasoleY는 pcp_manager.py의 replace() 호출 시 필드 상속 문제에 대해 지적했습니다. 이는 dataclasses.replace가 필드를 암시적으로 상속받는 특성 때문이며, 명시적인 필드 오버라이딩을 통해 잠재적인 버그를 방지하도록 개선되었습니다.

References

• vLLM Documentation
• RFC #46358: Unify Context Parallelism in vLLM

참고 자료

• https://docs.vllm.ai/en/latest/
• https://github.com/vllm-project/vllm/issues/46358

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.