[vLLM] Disaggregated Prefill/Decode: 분리된 서빙

들어가며

LLM 추론은 두 가지 근본적으로 다른 단계로 구성된다. Prefill은 전체 프롬프트를 한 번에 처리하여 KV 캐시를 생성하는 연산 집약적(compute-bound) 작업이고, Decode는 한 번에 한 토큰씩 생성하는 메모리 집약적(memory-bound) 작업이다. 이 두 단계의 하드웨어 요구사항이 다르므로, 같은 GPU에서 실행하면 둘 다 최적이 아니다. Disaggregated Serving은 Prefill 전용 노드와 Decode 전용 노드를 분리하고, Prefill이 생성한 KV 캐시를 네트워크를 통해 Decode 노드로 전송하는 아키텍처이다.

• 논문: DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving
• 공식 문서: https://docs.vllm.ai

공식 문서

vLLM 공식 문서: Disaggregated Prefill

핵심 구조/코드 분석

KVConnectorBase_V1: 추상 인터페이스

vllm/distributed/kv_transfer/kv_connector/v1/base.py에 정의된 KV 커넥터의 기본 인터페이스이다. 스케줄러 측과 워커 측으로 역할이 나뉜다:

class KVConnectorBase_V1(ABC):
    """
    Scheduler-side:
        get_num_new_matched_tokens() - 원격 KV 캐시에 있는 새 토큰 수
        update_state_after_alloc() - 임시 버퍼 할당 후 상태 갱신
        request_finished() - 요청 완료 시 KV 캐시 해제 여부 결정

    Worker-side:
        start_load_kv() - KV 로드 시작 (비동기 가능)
        wait_for_layer_load() - i번째 레이어 로드 완료 대기
        save_kv_layer() - i번째 레이어 KV 저장 시작
        wait_for_save() - 모든 저장 완료 대기
        get_finished() - 비동기 전송/수신 완료된 요청 ID 반환
    """

이 설계의 핵심은 레이어별 파이프라이닝이다. 전체 KV 캐시 전송을 기다리지 않고, 레이어 0의 KV가 도착하면 즉시 해당 레이어의 디코딩을 시작할 수 있다.

다양한 커넥터 구현

vllm/distributed/kv_transfer/kv_connector/v1/ 디렉토리에는 다양한 전송 백엔드가 구현되어 있다:

커넥터	전송 방식	특징
nixl_connector.py	NIXL	NVIDIA 고성능 전송 라이브러리
lmcache_connector.py	LMCache	외부 KV 캐시 스토어
lmcache_mp_connector.py	LMCache MP	멀티프로세스 버전
p2p/	Point-to-Point	GPU 직접 전송
mooncake/	Mooncake	Moonshot AI의 전송 엔진
offloading_connector.py	CPU Offloading	CPU 메모리로 오프로드
flexkv_connector.py	FlexKV	유연한 KV 관리
multi_connector.py	Multi	여러 커넥터 조합

KVConnectorRole: 역할 분리

class KVConnectorRole(enum.Enum):
    # 스케줄러에서 실행되는 역할
    SCHEDULER = "scheduler"
    # 워커에서 실행되는 역할
    WORKER = "worker"

하나의 KV 커넥터 클래스가 스케줄러와 워커 양쪽에서 인스턴스화된다. 스케줄러 측은 메타데이터 관리(어떤 토큰이 캐시되어 있는지)를, 워커 측은 실제 텐서 전송을 담당한다.

레이어별 KV 전송 프로토콜

워커 측의 핵심 프로토콜은 다음과 같다:

# Prefill 노드 (Sender)
for layer_idx in range(num_layers):
    # 레이어의 forward pass 수행 후 KV 저장
    save_kv_layer(layer_idx, kv_tensors)
wait_for_save()  # 모든 레이어 전송 완료 대기

# Decode 노드 (Receiver)
start_load_kv(request)  # 비동기 로드 시작
for layer_idx in range(num_layers):
    wait_for_layer_load(layer_idx)  # 이 레이어 로드 완료 대기
    # 즉시 이 레이어의 forward pass 실행 가능

이 파이프라인 구조 덕분에 KV 전송과 연산이 겹칠 수 있다. 레이어 0의 KV를 전송하는 동안 레이어 1의 prefill이 동시에 진행된다.

HMA(Hybrid Memory Allocator) 지원

class SupportsHMA(ABC):
    """Hybrid Memory Allocator를 지원하는 커넥터."""
    @abstractmethod
    def request_finished_all_groups(
        self, request, block_ids,
    ) -> tuple[bool, dict[str, Any] | None]:
        """요청 완료 시 모든 KV 캐시 그룹에 대해 호출된다."""

HMA는 GPU 메모리와 CPU 메모리를 함께 사용하는 하이브리드 할당자이다. KV 커넥터가 HMA를 지원하면, 요청 완료 시 블록을 GPU에서 CPU로 오프로드하거나 원격 노드로 전송하는 것을 커넥터가 제어한다.

왜 이 설계인가

1. Prefill/Decode 분리의 근본 이유: Prefill은 높은 연산 활용률(compute utilization)이 필요하고, Decode는 높은 메모리 대역폭이 필요하다. 동일 GPU에서 둘을 시분할하면 어느 쪽도 최적이 아니다. 분리하면 각 단계에 최적화된 배치 크기와 하드웨어를 사용할 수 있다.

2. 레이어별 파이프라이닝: 32레이어 모델의 KV 캐시 전체를 전송하면 수백 ms가 소요될 수 있다. 레이어별로 전송하면 첫 레이어 도착 즉시 디코딩을 시작할 수 있어 지연을 크게 줄인다.

3. 다중 백엔드: NVIDIA GPU 간에는 NVLink/NIXL이 최적이고, 크로스 노드에서는 RDMA 기반 전송이, CPU 오프로드에서는 PCIe가 사용된다. 하드웨어 토폴로지에 따라 최적의 전송 방식이 다르므로, 플러그인 아키텍처로 다양한 백엔드를 지원한다.

4. 비동기 완료 통보: get_finished()는 비동기 전송이 완료된 요청 ID를 반환한다. 스케줄러는 이를 통해 KV 전송이 완료된 요청만 디코딩 큐에 넣을 수 있다.

논문 핵심 내용

DistServe: Disaggregating Prefill and Decoding for Goodput-optimized LLM Serving (2401.09670) 논문은 Prefill과 Decode를 물리적으로 분리하여 서빙 처리량을 극적으로 개선하는 아키텍처를 제안했다.

핵심 아이디어: 기존 시스템은 Prefill과 Decode를 같은 GPU에서 배치 처리하면서, TTFT(Time To First Token)와 TPOT(Time Per Output Token) 사이에 상충 관계가 발생했다. DistServe는 두 단계를 별도 GPU에 할당하고, 각 단계에 최적화된 리소스 할당과 병렬화 전략을 적용한다. 네트워크 인식 배치로 KV 캐시 전송 오버헤드를 최소화했다.

애플리케이션별 Goodput 향상

애플리케이션	모델	vs vLLM	vs DeepSpeed-MII	TTFT SLO	TPOT SLO
Chatbot (ShareGPT)	OPT-13B	2.0-4.6x	-	0.25s	0.1s
Code Completion	OPT-66B	5.7x	1.6x	0.125s	0.2s
Summarization	OPT-66B	4.3x	1.8x	15s	0.15s

주요 수치

메트릭	수치
최대 요청 처리량 향상	7.4x
최대 SLO 강화	12.6x
SLO 충족률	> 90%
KV 캐시 전송 지연 (OPT-175B)	전체 지연의 0.1% 미만
95% 요청의 전송 지연	30ms 미만

KV 캐시 전송 오버헤드가 OPT-175B 같은 초대규모 모델에서도 전체 지연의 0.1% 미만이라는 결과가 인상적이다. 인트라 노드 NVLink(600 GB/s)을 활용하면 전송 비용이 거의 무시할 수 있는 수준이 된다. 크로스 노드(25 Gbps) 환경에서도 95% 이상의 요청이 30ms 미만의 전송 지연만 경험했다.

마무리

Disaggregated Serving은 LLM 추론의 두 단계가 가진 본질적 차이를 활용한 아키텍처 최적화이다. vLLM은 KVConnectorBase_V1이라는 추상 인터페이스 위에 NIXL, LMCache, Mooncake 등 다양한 전송 백엔드를 플러그인 방식으로 지원하여, 하드웨어 환경에 맞는 최적의 KV 전송을 가능하게 한다.