본문으로 건너뛰기

[vLLM] KV Transfer Connectors: KV 캐시 전송 프레임워크

들어가며

대규모 LLM 서빙에서 prefill(프롬프트 처리)과 decode(토큰 생성)를 서로 다른 GPU 클러스터에서 실행하는 Disaggregated Serving이 주목받고 있다. 이를 위해서는 prefill 서버의 KV 캐시를 decode 서버로 전송해야 한다. vLLM의 KV Transfer Connector 프레임워크는 이 전송을 추상화하여 NIXL, LMCache, 호스트 오프로딩 등 다양한 백엔드를 플러그인 방식으로 지원한다.

소스 경로: vllm/distributed/kv_transfer/kv_connector/v1/

공식 문서

vLLM 공식 문서: P2P NCCL Connector

핵심 구조/코드 분석

KVConnectorBase_V1 - 기본 추상 클래스

class KVConnectorBase_V1(ABC):
    def __init__(self, vllm_config, role: KVConnectorRole,
                 kv_cache_config=None):
        self._connector_metadata: KVConnectorMetadata | None = None
        self._vllm_config = vllm_config
        self._kv_transfer_config = vllm_config.kv_transfer_config
        self._role = role  # SCHEDULER or WORKER

모든 KV 커넥터의 기본 클래스이다. role로 스케줄러 측과 워커 측을 구분한다. 이 두 역할은 서로 다른 프로세스에서 실행되며, 메타데이터를 통해 통신한다.

역할 분리: Scheduler-side vs Worker-side

스케줄러 측 메서드는 KV 캐시의 논리적 상태를 관리한다:

@abstractmethod
def get_num_new_matched_tokens(
    self, request, num_computed_tokens
) -> tuple[int | None, bool]:
    """원격 KV 캐시에서 로드할 수 있는 토큰 수 반환.
    None이면 아직 확인 중이므로 나중에 다시 쿼리해야 함."""

@abstractmethod
def update_state_after_alloc(
    self, request, blocks, num_external_tokens):
    """블록 할당 후 커넥터 상태 갱신"""

@abstractmethod
def build_connector_meta(self, scheduler_output) -> KVConnectorMetadata:
    """이번 스텝의 전송 메타데이터 빌드"""

워커 측 메서드는 실제 GPU 메모리 간 데이터 전송을 수행한다:

@abstractmethod
def start_load_kv(self, forward_context, **kwargs):
    """비동기 KV 로드 시작 (포워드 패스 전에 호출)"""

@abstractmethod
def wait_for_layer_load(self, layer_name):
    """특정 레이어의 KV 로드 완료 대기"""

@abstractmethod
def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
    """KV 캐시 레이어 저장 시작 (어텐션 레이어에서 호출)"""

@abstractmethod
def wait_for_save(self):
    """모든 저장 완료 대기"""

레이어별 파이프라이닝

# 포워드 컨텍스트 진입 시
start_load_kv(forward_context)

# 각 어텐션 레이어에서
wait_for_layer_load("model.layers.0.self_attn")  # 이 레이어 로드 완료 대기
# ... 어텐션 연산 ...
save_kv_layer("model.layers.0.self_attn", kv_layer, attn_metadata)  # 비동기 저장

# 포워드 컨텍스트 종료 시
wait_for_save()  # 모든 저장 완료

KV 로드와 저장이 레이어 단위로 파이프라이닝된다. 레이어 0의 KV를 로드하는 동안 모델은 이미 로드 완료된 레이어들의 어텐션을 계산할 수 있다.

비동기 완료 추적

def request_finished(self, request, block_ids) -> tuple[bool, dict | None]:
    """요청 완료 시 호출. True 반환하면 블록을 비동기로 해제."""
    return False, None

def get_finished(self, finished_req_ids) -> tuple[set | None, set | None]:
    """비동기 전송이 완료된 요청 ID 반환.
    (보내기 완료 ID, 받기 완료 ID)"""
    return None, None

KV 전송이 비동기인 경우, request_finished()에서 True를 반환하면 블록 해제 책임을 커넥터가 가져간다. 전송이 완료되면 get_finished()를 통해 스케줄러에 알린다.

구현체 목록

vllm/distributed/kv_transfer/kv_connector/v1/
├── base.py                # 기본 추상 클래스
├── nixl_connector.py      # NIXL 원격 전송 (RDMA)
├── lmcache_connector.py   # LMCache 통합
├── offloading_connector.py    # CPU/디스크 오프로딩
├── simple_cpu_offload_connector.py  # 간단한 CPU 오프로딩
├── flexkv_connector.py    # FlexKV
├── multi_connector.py     # 여러 커넥터 조합
└── example_connector.py   # 예제 구현

CUDA Graph 호환성

@classmethod
def requires_piecewise_for_cudagraph(cls, extra_config) -> bool:
    """비동기 레이어별 연산을 사용하는 커넥터는
    PIECEWISE CUDA graph 모드가 필요함을 알림"""
    return False

레이어별 wait_for_layer_load/save_kv_layer는 CUDA 그래프 안에서 캡처될 수 없다. 이런 커넥터는 PIECEWISE 모드를 요구하여, 그래프 조각 사이에 Python 동기화 코드를 실행할 수 있게 한다.

왜 이 설계인가

  1. Disaggregated Serving: Prefill과 Decode를 분리하면 각각의 GPU 특성에 최적화할 수 있다. Prefill은 연산 집약적이고, Decode는 메모리 대역폭 제한적이므로 서로 다른 하드웨어가 적합하다.

  2. 스케줄러/워커 분리: 스케줄러 측은 "어떤 KV를 전송할지" 결정하고, 워커 측은 "어떻게 전송할지" 실행한다. 이 분리로 전송 백엔드(NIXL, LMCache 등)를 교체해도 스케줄러 로직을 변경할 필요가 없다.

  3. 레이어별 파이프라이닝: 모든 레이어의 KV를 한꺼번에 전송하면 포워드 패스가 완전히 블로킹된다. 레이어 단위로 전송하면 통신과 연산을 오버랩하여 전체 레이턴시를 줄인다.

  4. 플러그인 아키텍처: KVConnectorBase_V1을 상속하여 새로운 전송 백엔드를 쉽게 추가할 수 있다. example_connector.py가 구현 가이드 역할을 한다.

참고

댓글

관련 포스트

vLLM 의 다른글