본문으로 건너뛰기

[vLLM] KV Cache Offloading: GPU에서 CPU로의 KV 캐시 오프로딩

들어가며

긴 시퀀스를 처리할 때 KV 캐시는 GPU 메모리의 상당 부분을 차지한다. vLLM v1은 vllm/v1/kv_offload/ 모듈에서 KV 캐시 블록을 GPU에서 CPU로 오프로딩하여 GPU 메모리를 절약하는 시스템을 구현하고 있다. 스케줄러 측(OffloadingManager)과 워커 측(worker)이 분리된 깔끔한 아키텍처를 가진다.

핵심 구조/코드 분석

OffloadingManager 추상 인터페이스

vllm/v1/kv_offload/abstract.py에 정의된 핵심 인터페이스는 5가지 원시 연산을 제공한다.

class OffloadingManager(ABC):
    @abstractmethod
    def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:
        """첫 번째부터 연속으로 오프로딩된 블록의 최대 길이를 반환"""
        pass

    @abstractmethod
    def prepare_load(self, block_hashes: Iterable[BlockHash]) -> LoadStoreSpec:
        """블록들을 읽기 준비. 완료까지 eviction에서 보호"""
        pass

    @abstractmethod
    def prepare_store(self, block_hashes: Iterable[BlockHash]) -> PrepareStoreOutput | None:
        """블록들을 쓰기 준비. eviction 필요 시 처리"""
        pass

lookup은 스케줄러가 요청의 KV 캐시가 이미 CPU에 있는지 확인할 때 사용한다. None을 반환하면 나중에 재시도하라는 의미다. prepare_load/prepare_store는 비동기 I/O를 위한 2단계 프로토콜이다.

CPUOffloadingManager와 캐시 정책

_CACHE_POLICIES: dict[str, type[CachePolicy]] = {
    "lru": LRUCachePolicy,
    "arc": ARCCachePolicy,
}

class CPUOffloadingManager(OffloadingManager):
    def __init__(self, block_size, num_blocks, cache_policy="lru", enable_events=False):
        self.block_size = block_size
        self._num_blocks = num_blocks
        self._free_list: list[int] = []
        policy_cls = _CACHE_POLICIES.get(cache_policy)
        self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks)

CPU 오프로딩 매니저는 플러그인형 캐시 정책(LRU 또는 ARC)을 지원한다. ARC(Adaptive Replacement Cache)는 최근성과 빈도를 모두 고려하는 고급 캐시 전략으로, 워크로드 패턴에 따라 자동으로 적응한다.

블록 풀 관리

def _allocate_blocks(self, block_hashes: list[BlockHash]) -> list[BlockStatus]:
    num_fresh = min(len(block_hashes), self._num_blocks - self._num_allocated_blocks)
    num_reused = len(block_hashes) - num_fresh
    blocks: list[BlockStatus] = []
    for _ in range(num_fresh):
        blocks.append(BlockStatus(self._num_allocated_blocks))
        self._num_allocated_blocks += 1
    for _ in range(num_reused):
        blocks.append(BlockStatus(self._free_list.pop()))
    return blocks

새 블록은 순차적으로 할당하고, 해제된 블록은 free list에서 재사용한다. 이중 할당 전략으로 fragmentation을 최소화한다.

Store/Load 프로토콜

def prepare_store(self, block_hashes):
    block_hashes_to_store = [bh for bh in block_hashes_list if self._policy.get(bh) is None]
    if num_blocks_to_evict > 0:
        protected = set(block_hashes_list)  # 이번 요청의 블록은 eviction 대상에서 제외
        evicted = self._policy.evict(num_blocks_to_evict, protected)
    blocks = self._allocate_blocks(block_hashes_to_store)
    for block_hash, block in zip(block_hashes_to_store, blocks):
        self._policy.insert(block_hash, block)
    return PrepareStoreOutput(block_hashes_to_store, store_spec, to_evict)

prepare_store의 핵심은 protected set이다. 현재 저장하려는 블록은 eviction 대상에서 제외하여, 방금 저장한 블록이 바로 제거되는 문제를 방지한다.

이벤트 시스템

@dataclass
class OffloadingEvent:
    block_hashes: list[BlockHash]
    block_size: int
    medium: str
    removed: bool  # True면 제거, False면 저장

def take_events(self) -> Iterable[OffloadingEvent]:
    if self.events is not None:
        yield from self.events
        self.events.clear()

오프로딩 이벤트를 추적하여, 외부 시스템(예: KV 이벤트 퍼블리셔)이 블록의 상태 변화를 구독할 수 있다.

왜 이 설계인가

  1. 2단계 prepare/complete 프로토콜: 비동기 I/O를 지원하기 위해 준비(prepare)와 완료(complete)를 분리했다. prepare_load 호출 시 ref_cnt를 증가시켜 eviction을 방지하고, complete_load 시 해제한다. 이로써 GPU-CPU 간 비동기 복사 중에도 안전성이 보장된다.

  2. 스케줄러-워커 분리: OffloadingManager는 스케줄러 프로세스에서 실행되며, 실제 데이터 전송은 워커가 LoadStoreSpec을 받아 수행한다. 이 분리로 스케줄링 결정과 데이터 이동이 독립적으로 진행될 수 있다.

  3. ARC 캐시 정책: LRU는 단순하지만, 한 번 접근한 뒤 재접근하지 않는 scan 패턴에 취약하다. ARC는 이런 상황에서 자동으로 적응하여, 다양한 LLM 워크로드에서 더 나은 히트율을 제공한다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글