[vllm] vLLM의 분산 추론 성능 극대화: 양방향 KV 캐시 전송을 통한 Prefill 최적화

PR 링크: vllm-project/vllm#32553 상태: Merged | 변경: +None / -None

들어가며

vLLM의 분산 추론(Disaggregated Serving) 환경에서 Prefill 노드와 Decode 노드는 각각의 역할을 수행합니다. 하지만 멀티턴 대화 시, 이전 턴의 Decode 노드에 저장된 KV 캐시를 Prefill 노드가 활용하지 못하고 매번 재계산(Recompute)하는 비효율이 존재했습니다. 본 PR은 Prefill 노드가 Decode 노드로부터 KV 캐시를 직접 가져오는 '양방향 KV 캐시 전송' 기능을 도입하여, 중복 계산을 제거하고 TTFT(Time To First Token)를 획기적으로 개선합니다.

코드 분석

1. disagg_proxy_multiturn.py (신규 프록시 서버)

멀티턴 대화의 상태를 관리하기 위해 ConversationKVCache 클래스가 도입되었습니다. 각 대화는 conversation_id로 식별되며, Decode 노드에서 생성된 KV 전송 파라미터를 저장했다가 다음 턴의 Prefill 요청 시 주입합니다.

# Before: 상태 관리 없음, 매 턴마다 Prefill 재계산
# After: 프록시가 KV 캐시 정보를 캐싱하여 Prefill에 전달
class ConversationKVCache:
    def get(self, conversation_id: str) -> dict[str, Any] | None:
        entry = self._store.pop(conversation_id, None)
        # ... 캐시 만료 확인 및 반환

2. nixl_connector.py (커넥터 로직 개선)

NixlConnector는 이제 bidirectional_kv_xfer 설정을 통해 양방향 전송을 지원합니다. 특히 remote_pull_threshold를 도입하여, 특정 토큰 수 이상의 경우에만 원격 풀링을 수행하도록 하여 오버헤드를 제어합니다.

# Before: 일방향(P->D) 전송에 최적화
# After: P 노드가 D 노드의 KV 캐시를 pull 할 수 있도록 확장
if self.kv_transfer_config.get_from_extra_config("bidirectional_kv_xfer", False):
    # D 노드로부터 KV 블록을 가져오는 로직 수행

왜 이게 좋은가

성능 개선 수치

• TTFT(Time To First Token) 개선: Prefix Caching이 비활성화된 환경에서 2배 이상의 성능 향상을 보였습니다.
• Prefix Caching 활성화 시: P50 기준 TTFT가 307.86ms에서 219.61ms로 약 28.6% 개선되었습니다.

핵심 교훈

1. 중복 계산 제거: LLM 추론에서 가장 비용이 큰 Prefill 단계를 캐시 재사용으로 대체함으로써 리소스 효율성을 극대화할 수 있습니다.
2. 상태 관리의 중요성: 분산 환경에서 노드 간의 상태(KV 캐시)를 동기화하기 위해 프록시 계층에서 컨텍스트(conversation_id)를 관리하는 패턴이 효과적입니다.
3. 임계치 제어: 모든 요청에 대해 원격 풀링을 수행하는 대신 remote_pull_threshold를 두어, 네트워크 오버헤드와 계산 비용 간의 균형을 맞추는 설계가 필수적입니다.

리뷰 과정에서 NickLucche는 하이브리드 KV(Full + Sliding Window Attention) 지원과 remote_pull_threshold의 기본값 설정을 강조했으며, 이를 통해 더욱 안정적인 프로덕션 환경을 구축할 수 있게 되었습니다.

참고 자료

• https://github.com/vllm-project/vllm/issues/32733
• https://github.com/vllm-project/vllm/pull/40186

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