본문으로 건너뛰기

[vLLM] KV Cache Coordinator: 하이브리드 KV 캐시 조율

들어가며

최신 LLM 아키텍처 중에는 Full Attention과 Sliding Window Attention을 혼합하여 사용하는 하이브리드 모델이 있다. 이런 모델에서는 어텐션 타입마다 KV 캐시 관리 전략이 달라야 한다. vLLM의 vllm/v1/core/kv_cache_coordinator.py는 이 복잡성을 코디네이터 패턴으로 추상화한다.

공식 문서

vLLM 공식 문서: Hybrid KV Cache Manager

핵심 구조/코드 분석

KVCacheCoordinator 추상 클래스

class KVCacheCoordinator(ABC):
    def __init__(self, kv_cache_config, max_model_len, use_eagle, enable_caching, ...):
        self.block_pool = BlockPool(
            kv_cache_config.num_blocks, enable_caching, hash_block_size,
            enable_kv_cache_events, metrics_collector,
        )
        self.single_type_managers = tuple(
            get_manager_for_kv_cache_spec(
                kv_cache_spec=kv_cache_group.kv_cache_spec,
                block_pool=self.block_pool,
                enable_caching=enable_caching,
                kv_cache_group_id=i,
                ...
            )
            for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups)
        )

하나의 BlockPool을 모든 KV 캐시 그룹이 공유하되, 각 그룹은 전용 SingleTypeKVCacheManager로 관리된다.

세 가지 코디네이터 구현

팩토리 함수가 상황에 따라 적절한 코디네이터를 선택한다:

def get_kv_cache_coordinator(kv_cache_config, ...):
    if not enable_caching:
        return KVCacheCoordinatorNoPrefixCache(...)
    if len(kv_cache_config.kv_cache_groups) == 1:
        return UnitaryKVCacheCoordinator(...)
    return HybridKVCacheCoordinator(...)
  • NoPrefixCache: prefix caching 비활성화 시 사용. 0개 이상의 KV 캐시 그룹을 지원하지만, 캐시 히트 탐색은 항상 빈 결과를 반환한다.
  • Unitary: 단일 KV 캐시 그룹(예: 모든 레이어가 Full Attention). 가장 일반적인 경우다.
  • Hybrid: 복수 KV 캐시 그룹(예: Full + Sliding Window). 각 그룹의 캐시 히트를 조율해야 한다.

UnitaryKVCacheCoordinator의 캐시 히트

class UnitaryKVCacheCoordinator(KVCacheCoordinator):
    def find_longest_cache_hit(self, block_hashes, max_cache_hit_length):
        hit_blocks = self.single_type_managers[0].find_longest_cache_hit(
            block_hashes=block_hashes,
            max_length=max_cache_hit_length,
            kv_cache_group_ids=[0],
            block_pool=self.block_pool,
            kv_cache_spec=self.kv_cache_spec,
            use_eagle=self.use_eagle,
            alignment_tokens=self.block_size,
        )
        return hit_blocks, len(hit_blocks[0]) * self.block_size

단일 그룹이므로 단순하다. 블록 해시를 좌에서 우로 스캔하여 연속된 히트를 찾는다.

HybridKVCacheCoordinator: 반복 수렴 알고리즘

class HybridKVCacheCoordinator(KVCacheCoordinator):
    def find_longest_cache_hit(self, block_hashes, max_cache_hit_length):
        while True:
            curr_hit_length = hit_length
            for spec, group_ids, manager_cls in self.attention_groups:
                is_full_attn = isinstance(spec, FullAttentionSpec)
                if is_full_attn and cached_blocks is not None:
                    # Full attention: 이전 결과를 재사용 (하향 폐쇄 속성)
                    num_blocks = curr_hit_length // spec.block_size
                    curr_hit_length = num_blocks * spec.block_size
                else:
                    hit_blocks = manager_cls.find_longest_cache_hit(...)
                    curr_hit_length = len(hit_blocks[0]) * spec.block_size

            if curr_hit_length >= hit_length:
                break  # 수렴
            hit_length = curr_hit_length
            if is_simple_hybrid:
                break  # 단순 하이브리드는 1회 반복으로 충분

핵심은 반복 고정점 알고리즘이다. 각 어텐션 타입이 현재 후보 길이를 수용하거나 축소한다. 어떤 타입이 길이를 축소하면, 다시 모든 타입을 순회한다. 길이는 단조 감소하고 0 이상이므로 반드시 수렴한다.

LCM 블록 크기

block_sizes = [spec.block_size for spec, _, _ in attention_groups]
self.lcm_block_size = lcm(*block_sizes)

서로 다른 블록 크기를 가진 어텐션 그룹의 캐시 히트 길이는 모든 블록 크기의 LCM(최소공배수)의 배수여야 한다. 부분 블록 캐시 히트를 지원하지 않기 때문이다.

Full Attention 우선 정렬

self.attention_groups = sorted(
    attention_groups,
    key=lambda x: not isinstance(x[0], FullAttentionSpec),
)

Full Attention을 먼저 탐색한다. Full Attention의 효율적인 좌->우 스캔이 더 타이트한 초기 상한을 제공하여, 후속 그룹의 탐색 범위를 줄인다. 또한 Full Attention은 하향 폐쇄 속성(downward-closed property)을 가지므로, 길이가 축소되면 이전 결과에서 블록을 잘라내기만 하면 된다.

왜 이 설계인가

  1. 블록 풀 공유: 모든 어텐션 그룹이 하나의 블록 풀을 공유하면, 그룹 간 메모리 할당의 유연성이 높아진다. 특정 그룹의 캐시가 더 필요하면 다른 그룹의 블록을 재활용할 수 있다.

  2. 단순 하이브리드 최적화: Full Attention 1개 + 다른 타입 1개인 경우(가장 흔함)는 1회 반복으로 충분하다고 판단하여 루프를 일찍 탈출한다. EAGLE 사용 시 복수 반복에서 블록 드롭이 중복 적용되는 문제도 회피한다.

  3. CrossAttentionManager 특수 처리: 인코더-디코더 모델의 교차 어텐션은 인코더 토큰 수에 기반하여 한 번만 정적으로 블록을 할당한다. 디코더의 점진적 블록 할당 패턴과 다르므로, get_num_blocks_to_allocate에서 별도 분기 처리한다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글