vLLM x Mooncake를 이용한 대규모 에이전틱 워크로드 서빙

에이전틱 워크로드는 대규모 재사용 가능한 접두사(prefix)를 생성합니다

에이전틱 워크로드는 재사용률이 높은 긴 컨텍스트를 생성하므로 prefix caching이 매우 중요합니다. Codex/SWE‑bench Pro 트레이스를 살펴보면 30번째 턴에서 컨텍스트 길이가 약 80K 토큰에 도달하며 180K 토큰을 초과할 수 있는 반면, 각 턴마다 추가되는 새로운 토큰은 수백 개에서 수천 개에 불과합니다. 트레이스당 중앙값 33턴인 610개의 트레이스 전체에서 데이터셋은 다음과 같은 특성을 보입니다:

94.2% 캐시 히트율(cache hit rate)
131:1 입출력 토큰 비율
턴당 약 2,242 토큰의 평균 컨텍스트 증가
트레이스당 12K에서 80K 토큰 사이의 컨텍스트 증가 중앙값
5.2s(중앙값)에서 81.4s(P99) 사이의 턴 간 지연 시간 이 수치들은 입력의 대부분이 재사용 가능한 prefix임을 확인시켜 주며, 이를 캐싱하면 불필요한 prefill 과정을 제거할 수 있습니다.

Mooncake Store를 통한 분산 KV cache 풀로 prefix 공유 문제 해결

Mooncake Store를 통합하면 공유 KV cache를 통해 히트율을 높이고 인스턴스 간 재사용을 가능하게 합니다. Mooncake는 KV cache 전송 및 분산 저장을 위한 오픈 소스 라이브러리입니다. vLLM은 이미 MooncakeConnector를 통해 prefill-decode 분리(disaggregation)를 위해 Mooncake를 사용하고 있습니다. 새로운 통합 방식은 여러 vLLM 인스턴스가 Mooncake 클라이언트를 포함하고 클러스터 전체의 Mooncake Store 마스터에 연결되는 분산 KV cache 풀을 구축합니다. 마스터는 KV-block 메타데이터, 서비스 디스커버리 및 클라이언트 상태를 관리하며, 워커는 RDMA를 통해 GPU HBM과 분산 DRAM/SSD 풀 사이에서 KV 블록을 전송합니다. 스케줄러 측면에서 vLLM은 프롬프트 토큰 블록을 해싱하고, 일치하는 KV cache 블록을 찾기 위해 Mooncake 마스터에 쿼리를 보낸 후, 그 결과를 스케줄링에 활용합니다. 워커 측면에서 각 GPU 워커는 Mooncake 클라이언트를 실행하고, GPU KV cache 메모리를 RDMA 버퍼로 등록하며, 백그라운드 I/O 스레드를 통해 데이터를 이동시킵니다.

설계 하이라이트: SM-free GPUDirect RDMA, 비동기 전송, MultiConnector

구현에는 zero-copy를 위한 GPUDirect RDMA가 사용되며, 전송은 백그라운드 I/O 스레드에서 실행되고, MultiConnector를 통해 PD 분리와 함께 작동합니다. KV 블록 이동은 GPUDirect RDMA를 활용하여 staging buffer와 SM 소비를 피하며, 더 높은 대역폭을 위해 여러 RNIC를 풀링할 수 있습니다. 모든 RDMA 작업은 전용 백그라운드 I/O 스레드에서 실행되므로 전송 경로가 완전히 비동기적으로 처리되어 메인 CPU 경로의 스톨(stall)을 방지합니다. MultiConnector 인터페이스는 서브 커넥터들을 체인 형태로 연결합니다. prefill 인스턴스는 KV 블록을 분산 풀에 저장하고 캐시 히트를 위해 이를 검색할 수 있으며, decode 인스턴스는 prefill에서 볼 수 있도록 블록을 풀에 씁니다.

Codex 트레이스에서의 성능 결과 및 확장성

실제 에이전틱 트레이스에서 캐시 히트율이 1.7%에서 92.2%로 급증하여, 60 GB200 GPU로 확장 시에도 3.8배 높은 처리량, 46배 낮은 TTFT, 8.6배 낮은 지연 시간을 달성하며 >95%의 히트율과 거의 선형적인 처리량 증가를 유지합니다. Codex 에이전틱 트레이스 실험(1P1D, 12 GB200 GPU)에서 분산 KV cache 풀은 처리량을 3.8배 증가시켰고, P50 TTFT를 46배 감소시켰으며, 엔드 투 엔드 지연 시간을 8.6배 단축했습니다. 이러한 이점은 캐시 히트율이 1.7%(시스템 프롬프트만 캐싱됨)에서 92.2%(거의 전체 prefix가 캐싱됨)로 상승한 데서 기인합니다. 확장성 테스트를 위해 Codex 워크로드에서 파생된 합성 데이터셋을 사용하여 노드 간 라운드 로빈 라우팅을 적용했습니다. 시스템은 모든 규모에서 95% 이상의 캐시 히트율을 유지했으며, 12개에서 60개의 GB200 GPU로 확장함에 따라 처리량이 거의 선형적으로 증가하여, 분산 풀이 성능을 저하시킬 수 있는 인스턴스 간 미스(miss)를 방지함을 입증했습니다.

향후 과제

계획된 확장 기능에는 분산 디스크 오프로딩, 하이브리드 모델 지원, 캐시 인식 라우팅(cache-aware routing) 및 데이터 경로 최적화가 포함됩니다. 향후 작업은 더 큰 캐시 용량을 위해 저장 계층을 NVMe SSD 및 분산 파일 시스템으로 확장할 것입니다. 혼합 어텐션 메커니즘을 가진 하이브리드 모델에 대한 지원이 추가되어 레이어당 서로 다른 캐싱 전략을 사용할 수 있게 됩니다. 캐시 인식 라우팅은 요청 라우터와 KV cache 풀을 공동 설계하여 턴을 관련 prefix를 보유한 인스턴스로 직접 전달함으로써, 분산 풀로 넘어가기 전에 로컬 히트를 극대화합니다. 데이터 경로 최적화는 NVIDIA multi-node NVLink와 prefill 및 decode 인스턴스로부터의 DualPath 방식과 유사한 동시 KV 로딩을 탐색하여 총 대역폭을 높일 것입니다.

감사의 말

vLLM Mooncake Store 통합은 vLLM-Ascend의 이전 연구에서 영감을 받았습니다. 초기 구현을 담당한 Ant Group의 Chao Lei와 에이전틱 트레이스 및 분석을 담당한 Inferact의 Zijing Liu에게 감사드립니다. 또한 Approaching.AI의 Jiahao Lu, Zuoyuan Zhang, Zihan Tang, Ke Yang; Huawei의 Pengbo Zhao, Fuqiao Duan, Tianyu Xu; Alibaba Cloud Computing의 Tianchen Ding, Xuchun Shang, Xingrui Yi, Teng Ma; Ant Group의 Yunxiao Ning, Dejiang Zhu, Shoujian Zheng; 그리고 9#AISoft의 Feng Ren에게 기술적 피드백에 대한 감사를 표합니다. 광범위한 vLLM 및 Mooncake 커뮤니티와 긴밀히 협력해 준 Inferact 팀의 지원에 감사드립니다.