분리형 프리필링

분리형 프리필링 (Disaggregated Prefilling, 실험적)

이 페이지는 vLLM의 분리형 프리필링(disaggregated prefilling) 기능을 소개합니다. 프리필과 디코드 단계를 서로 다른 vLLM 인스턴스에 나눠 두어, 첫 토큰 시간과 꼬리 지연을 각각 제어할 수 있게 해주는 기능입니다.

출처: 문서

참고: 이 기능은 실험적이며 변경될 수 있습니다.

본문

왜 분리형 프리필링인가? (Why disaggregated prefilling?)

두 가지 주요 이유가 있습니다.

  • TTFT와 ITL을 각각 독립적으로 튜닝. 분리형 프리필링은 LLM 추론의 프리필 단계와 디코드 단계를 서로 다른 vLLM 인스턴스에 둡니다. 이렇게 하면 tp나 pp 같은 서로 다른 병렬 전략을 할당해, ITL에 영향을 주지 않고 TTFT를 튜닝하거나 그 반대도 가능합니다.
  • 꼬리 ITL 제어. 분리형 프리필링이 없으면 vLLM은 한 요청의 디코딩 중간에 프리필 작업을 끼워 넣을 수 있어 꼬리 지연이 높아집니다. 분리형 프리필링은 이 문제를 해결하고 꼬리 ITL을 제어하게 해줍니다. 적절한 청크 크기의 chunked prefill도 같은 목표를 달성할 수 있지만, 실제로 올바른 청크 크기를 찾기는 어렵습니다. 그래서 분리형 프리필링이 꼬리 ITL을 제어하는 훨씬 더 확실한 방법입니다.

참고: 분리형 프리필링은 처리량(throughput)을 향상시키지 않습니다.

사용 예제 (Usage example)

현재 9가지 타입의 커넥터를 지원합니다.

--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'

--kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both"},{"kv_connector":"ExampleConnector","kv_role":"kv_both","kv_connector_extra_config":{"shared_storage_path":"local_storage"}}]}}'
  • OffloadingConnector: KV 데이터를 CPU 메모리로 오프로딩하고, CPU 블록 크기(토큰 단위)와 할당할 총 CPU 메모리 바이트를 지정합니다.
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'

다중 계층 오프로딩(예: CPU + 파일시스템 계층)과 전체 구성 참조는 KV Offloading Usage Guide를 참고하세요.

--kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'

디코드 시 프리필 토큰 ID 재사용 (Reusing prefill token ids on decode)

참고: 이 내용은 /v1/chat/completions 엔드포인트에서 KV 커넥터를 위 Usage 예제처럼 구성해 서빙하는 분리형 프리필/디코드에 적용됩니다. 실험적이며 변경될 수 있습니다.

분리형 서빙에서 프리필과 디코드 단계는 모두 messages에서 채팅 프롬프트를 렌더링하고 토크나이즈합니다. 프리필 단계가 이미 토큰 ID를 만들었으므로, 디코드 단계는 그 토큰 ID를 재사용해 자체 템플릿 작업과 토크나이즈를 건너뛸 수 있습니다. 출력은 여전히 일반 채팅 완료와 동일합니다. 텍스트로 디토크나이즈되며, 도구·추론 파싱, 스트리밍, 구조화된 출력 제약이 모두 그대로 적용됩니다.

토큰 ID는 전송을 조정하기 위해 디코드 요청에 이미 붙어 있는 dict인 kv_transfer_params를 통해 디코드 단계로 전달됩니다.

  • 프리필 요청을 return_token_ids를 켜고 보내고, 응답에서 prompt_token_ids를 읽습니다.
  • 디코드 요청의 kv_transfer_params["prompt_token_ids"]를 그 ID들로 설정합니다. messages는 여전히 필요하지만, ID가 있으면 그 내용은 토크나이즈되지 않습니다.
prefill = client.chat.completions.create(
    model=model,
    messages=messages,
    extra_body={"return_token_ids": True, "kv_transfer_params": {"do_remote_decode": True}},
)
ids = prefill.prompt_token_ids

decode = client.chat.completions.create(
    model=model,
    messages=messages,
    stream=True,
    extra_body={"kv_transfer_params": {"do_remote_prefill": True, "prompt_token_ids": ids}},
)

개발 (Development)

분리형 프리필링은 2개의 vLLM 인스턴스를 실행해 구현합니다. 하나는 프리필용(프리필 인스턴스이라 부름), 하나는 디코드용(디코드 인스턴스라 부름)이며, 커넥터가 프리필 KV 캐시와 결과를 프리필 인스턴스에서 디코드 인스턴스로 전송합니다.

모든 분리형 프리필링 구현은 vllm/distributed/kv_transfer 아래에 있습니다.

분리형 프리필링의 핵심 추상화:

  • Connector: kv consumer가 kv producer로부터 요청 배치의 KV 캐시를 가져올 수 있게 해주는 커넥터.
  • LookupBuffer: 두 API를 제공합니다. insert는 KV 캐시를 버퍼에 넣고, drop_select는 주어진 조건과 일치하는 KV 캐시를 반환하면서 버퍼에서 제거합니다. 의미는 SQL의 insert/drop_select와 유사합니다.
  • Pipe: 텐서 전송을 위한 단방향 FIFO 파이프. send_tensor와 recv_tensor를 지원합니다.

참고: insert는 비차단(non-blocking)이지만 drop_select는 차단(blocking) 연산입니다.

다음 그림은 위 3개 추상화가 어떻게 구성되는지 보여줍니다.

분리형 프리필링의 워크플로는 다음과 같습니다.

여기서 buffer는 LookupBuffer의 insert API에, drop_select는 LookupBuffer의 drop_select API에 대응합니다.

이제 vLLM의 모든 프로세스에는 해당 커넥터가 있습니다. 구체적으로는:

  • Scheduler 커넥터: 스케줄러 프로세스와 같은 프로세스에 있는 커넥터. KV 캐시 전송 연산을 스케줄링합니다.
  • Worker 커넥터: 워커 프로세스에 있는 커넥터. KV 캐시 전송 연산을 실행합니다.

다음 그림은 위 2개 커넥터의 구성 방식을 보여줍니다.

아래 그림은 worker 커넥터가 어텐션 모듈과 함께 동작해 레이어 단위로 KV 캐시를 저장·로드하는 방식을 보여줍니다.

서드파티 기여 (Third-party contributions)

분리형 프리필링은 인프라와 밀접한 관련이 있어, vLLM은 프로덕션 수준의 분리형 프리필링을 위해 서드파티 커넥터에 의존합니다(vLLM 팀은 서드파티 커넥터의 새 PR을 적극적으로 검토하고 병합합니다).

세 가지 구현 방식을 권장합니다.

  • 완전 맞춤 커넥터: 자신만의 Connector를 구현하고 서드파티 라이브러리를 호출해 KV 캐시를 주고받는 방식(커스텀 프리필링을 위해 vLLM의 모델 입력을 수정하는 등 더 많은 것도 가능). 가장 많은 제어권을 주지만, 향후 vLLM 버전과 호환이 안 될 위험이 있습니다.
  • 데이터베이스형 커넥터: SQL처럼 insert와 drop_select API를 지원하는 자신만의 LookupBuffer를 구현.
  • 분산 P2P 커넥터: torch.distributed처럼 send_tensor와 recv_tensor API를 지원하는 자신만의 Pipe를 구현.

더 알아보기 (Learn more)