Hybrid KV Cache Manager

Hybrid KV Cache Manager

경고: 이 문서는 커밋 458e74 기준으로 작성됐습니다. 이 기능은 아직 초기 단계이며 내용은 바뀔 수 있습니다.

하이브리드 모델이란? (What is a hybrid model?)

최근의 많은 "하이브리드" LLM은 하나의 모델 안에서 여러 attention 유형을 결합합니다. 예를 들어:

  • 슬라이딩 윈도우 attention(sw) + 전체 attention(full): gpt-oss, Gemma 2/3, Ministral, cohere 등
  • Mamba + full: Bamba, Jamba, Minimax 등
  • 로컬 청크 attention + full: Llama4

이런 모델을 효율적으로 서빙하려면 KVCacheManager가 다음을 해야 합니다:

  • 레이어 유형마다 다른 슬롯을 할당합니다. 예:
    • 전체 attention 레이어: 모든 토큰에 슬롯을 예약합니다.
    • 슬라이딩 윈도우 레이어: 가장 최근 sliding_window_size 토큰에만 슬롯을 예약합니다.
  • 레이어별 프리픽스 캐시 규칙을 지원합니다. 예:
    • 전체 attention: 캐시 히트 프리픽스는 모든 토큰이 KV 캐시에 남아 있어야 합니다.
    • 슬라이딩 윈도우: 캐시 히트 프리픽스는 마지막 sliding_window_size 토큰만 KV 캐시에 남아 있으면 됩니다.

출처: 문서

본문

정의 (Definitions)

  • kv hidden size: 단일 레이어의 한 토큰 KV 캐시를 저장하는 바이트 수.
  • block: KV 캐시용으로 예약된 메모리가 같은 페이지 크기(아래 정의)로 나뉜 블록.
  • block size: 블록 안의 토큰 수.
  • page size: 블록의 물리 메모리 크기. 다음과 같이 정의됩니다:

$$\text{num_layers} \times \text{block_size} \times \text{kv_hidden_size}$$

num_layers는 모델의 전체 레이어 수를 뜻하지 않습니다. 정확한 숫자는 이 문서의 맥락에 따라 다릅니다.

참고: 이는 코드의 KVCacheSpec.page_size_bytes와 다릅니다. 그 값은 다음과 같이 정의됩니다: $$\text{block_size} \times \text{kv_hidden_size}$$

할당 (Allocation)

높은 수준의 아이디어 (High level idea)

모든 레이어 유형에 단일 메모리 풀을 사용합니다. 메모리 풀은 같은 페이지 크기의 여러 블록으로 나뉩니다. KVCacheManager는 attention 유형에 따라 각 레이어에 다른 수의 블록을 할당합니다.

핵심 과제는 모든 레이어 유형이 같은 페이지 크기를 사용하도록 하는 것입니다. full-attention 전용 모델에서 페이지 크기는 단순합니다:

$$\text{page_size} = \text{block_size} \times \text{num_hidden_layers} \times \text{kv_hidden_size}$$

그러나 하이브리드 모델에서는 num_hidden_layers가 attention 유형에 따라 달라져 보통 페이지 크기가 일치하지 않게 됩니다. 아래 사례들이 어떻게 통합하는지 보여줍니다.

사례 1: 장난감 모델 (Case 1: toy model)

장난감 예시로 시작합니다: full attention 레이어 1개와 sliding window attention 레이어 3개가 있는 모델입니다. 모든 레이어는 같은 kv_hidden_size를 갖습니다.

각 블록이 한 레이어의 block_size개 토큰을 담게 하면:

$$\text{page_size} = \text{kv_hidden_size} \times \text{block_size}$$

KVCacheManager는 각 레이어에 다른 수의 블록을 할당합니다.

이 사례는 장난감 예시일 뿐입니다. 실제 모델은 아래 사례를 참고하세요.

사례 2: 같은 kv_hidden_size와 규칙적인 패턴 (Case 2: same kv_hidden_size and a regular pattern)

모델에 레이어가 더 많을 때, 예를 들어 같은 kv_hidden_size를 가진 sliding window attention 레이어 20개와 full attention 레이어 10개가 있다고 합시다. 레이어마다 한 번씩 할당자 호출(30회)도 가능하지만 비효율적입니다. 해결책으로, 같은 블록 수가 필요한 레이어들의 할당을 그룹화해 호출 수를 줄입니다.

그룹화는 보통 서로 다른 유형 레이어 수 사이에 아름다운 비율이 있기 때문에 가능합니다. 예:

  • Gemma-2: sw 1 : full 1
  • Llama 4: local 3 : full 1

우리 예는 sw 2 : full 1로 볼 수 있습니다. 모델에 sw 2개와 full 1개가 있는 것처럼 블록을 할당하고, 그 결과를 10번 반복해 30개 레이어의 block_ids를 생성할 수 있습니다. 페이지 크기는 다음과 같습니다:

$$10 \times \text{kv_hidden_size} \times \text{block_size}$$

block_size=16, sliding window size=32, 요청 길이=112라고 가정하면, 위 예제 모델에 대해 11개 블록(full용 0-6, sw 그룹 1용 7-8, sw 그룹 2용 9-10)을 할당해야 합니다.

여기서 "/"는 블록이 필요 없음을 뜻합니다(슬라이딩 윈도우 레이어는 초기 토큰용 슬롯이 필요 없습니다).

공식 정의를 보면, 레이어들은 여러 KV Cache Group으로 나뉘어 다음과 같은 특성을 갖습니다:

  • 그룹 내 동일한 attention 유형: 각 그룹은 같은 attention 유형의 레이어만 포함하므로 주어진 요청에 대해 같은 수의 블록이 필요합니다. 덕분에 같은 그룹의 레이어는 메모리 낭비 없이 같은 block id를 공유할 수 있습니다.
  • 그룹 간 동일한 페이지 크기: 메모리 풀은 페이지 크기가 하나뿐이기 때문입니다.

예제 모델은 3개의 KV cache group으로 나뉩니다:

  • Group 0: full attention 레이어 10개 (full.0 - full.9)
  • Group 1: sliding window attention 레이어 10개 (sw.0 - sw.9)
  • Group 2: sliding window attention 레이어 10개 (sw.10 - sw.19)

당연히 규칙 1을 만족합니다. 규칙 2의 경우 세 그룹 모두 페이지 크기가 다음과 같습니다:

$$10 \times \text{kv_hidden_size} \times \text{block_size}$$

사례 3: 같은 kv_hidden_size, 규칙적 패턴 없음 (Case 3: same kv_hidden_size and no regular pattern)

불행히도 모든 모델이 그런 아름다운 비율을 갖는 것은 아니며, 사례 2의 접근은 너무 많은 작은 그룹을 만들 수 있습니다. 예를 들어 Gemma-3-27b는 sliding window attention 레이어 52개와 full attention 레이어 10개가 있습니다. 사례 2의 제약에서는 26개의 sliding window 그룹과 5개의 full attention 그룹이 되고, 각 그룹은 레이어 2개를 담습니다. 할당은 여전히 비효율적입니다. kv cache group 수를 줄이기 위해 모든 attention 유형 중 가장 작은 레이어 수로 레이어를 그룹화합니다. 예를 들어 Gemma-3-27b에서는 min(52, 10)=그룹당 10개 레이어입니다. 그룹화 결과는:

  • Group 0: full attention 레이어 10개 (full.0 - full.9)
  • Group 1: sliding window attention 레이어 10개 (sw.0 - sw.9)
  • Group 2: sliding window attention 레이어 10개 (sw.10 - sw.19)
  • Group 3: sliding window attention 레이어 10개 (sw.20 - sw.29)
  • Group 4: sliding window attention 레이어 10개 (sw.30 - sw.39)
  • Group 5: sliding window attention 레이어 10개 (sw.40 - sw.49)
  • Group 6: sliding window attention 레이어 2개 (sw.50 - sw.51)와 패딩 레이어 8개

새 모델이 나왔을 때 이 휴리스틱이 나쁜 결과를 내면(예: full 20 + sw 30, 그룹 크기는 20이 아니라 10이어야 함) 알고리즘을 갱신하겠습니다.

이 사례는 Gemma-3 시리즈 모델, 그리고 full attention 레이어를 하나 추가하는 eagle 추측 디코딩이 있는 사례 2 모델에서 발생합니다. 해결책은 메모리 낭비가 있고 완벽하지 않습니다. 패딩 오버헤드가 감당할 수 없게 되는 사례가 있으면 알려주세요. 알고리즘을 개선하겠습니다.

사례 4: 서로 다른 kv_hidden_size (주로 하이브리드 mamba 모델) (Case 4: different kv_hidden_size)

일부 아키텍처(예: Bamba, Jamba, Minimax)는 표준 attention 레이어와 Mamba 레이어를 교차 배치하는데, 각 Mamba 레이어의 토큰당 상태 크기는 attention 레이어의 kv_hidden_size보다 훨씬 클 수 있습니다. 모든 그룹에 단일 페이지 크기만 지원하므로 이 서로 다른 hidden size를 조정해야 합니다.

현재 알고리즘은:

  • $$\text{block_size} \times \text{kv_hidden_size}_{att} \ge \text{state_size}$$ 가 될 때까지 attention 레이어의 block_size를 늘립니다.
  • Mamba 상태를 레이어당 $$\text{block_size} \times \text{kv_hidden_size}_{att}$$로 패딩합니다.
  • 사례 3의 그룹화 전략을 적용합니다.

참고: 이로 인해 attention 레이어의 block_size가 400이 넘을 수 있는데 너무 큽니다. 또 다른 패딩 전략은 다음이 될 때까지 block_size를 늘리는 것입니다: $$\text{block_size} \times \text{kv_hidden_size}{att} \times \text{num_attn_layers} \ge \text{state_size}{mamba}$$ 이 패딩 전략은 아직 진행 중입니다.

사례 5: KV 공유 (Case 5: KV sharing)

KV 공유는 레이어가 다른 레이어의 KV 캐시를 사용하는 것을 뜻합니다(예: gemma-3n). 이런 모델에서 KVCacheManager는 KV 공유 레이어를 모두 무시하고 KV 캐시가 필요한 레이어에만 할당하며, 모델 러너에 몇 가지 패치를 적용해 할당 결과를 KV 공유 레이어에 반영합니다.

프리픽스 캐싱 (Prefix caching)

이 절에서는 단순화를 위해 block_size=1을 가정합니다.

높은 수준의 아이디어 (High level idea)

블록 풀은 tuple(block_hash, group_id) -> block과 유사한 dict를 사용해 전체 블록을 캐시합니다. 즉 서로 다른 그룹의 같은 토큰들이 독립적으로 캐시되고 제거됩니다.

새 요청이 들어오면 각 그룹의 캐시 히트 프리픽스를 확인하고, 이 그룹들의 교집합을 요청의 캐시된 프리픽스로 반환합니다. 한 그룹의 캐시 히트 확인 및 교집합 수행 알고리즘은 아래를 참고하세요.

사례 0: full attention 전용 모델 (Case 0: full attention only models)

full attention 레이어에서 요청의 모든 토큰에 블록이 할당됩니다. 기본 설계에 대한 자세한 내용은 Prefix Caching을 참고하세요.

요청의 가장 긴 캐시 히트 프리픽스를 찾으려면 왼쪽(첫 블록)에서 오른쪽(마지막 블록)으로 열거하며 블록이 캐시됐는지 확인하고, 캐시 미스가 나면 종료합니다. 예를 들어 아래 예에서 (파란 블록이 캐시됨) 첫 7개 토큰(0-6)을 캐시 히트 프리픽스로 반환합니다.

사례 1: sliding window attention 전용 모델 (Case 1: sliding window attention only models)

sliding window attention 레이어의 메모리 할당 순진한 구현은 sliding_window_size개 블록을 할당하고 round-robin 방식으로 채우는 것입니다. 하지만 이 순진한 구현은 프리픽스 캐싱과 호환되지 않아 이 설계를 채택하지 않았습니다. vLLM에서는 서로 다른 토큰에 서로 다른 블록을 할당하고 슬라이딩 윈도우 밖의 블록을 해제합니다.

새 요청의 캐시 히트 프리픽스는 마지막 sliding_window_size - 1개 토큰만 캐시되면 됩니다. sliding_window_size=4, block_size=1이고 요청이 15-token 프롬프트라고 합시다(파란 블록이 캐시됨).

가능한 캐시 히트 프리픽스가 3개 있습니다:

  • 캐시 히트 길이 5, [2, 3, 4]로 prefill 계산 → [5, 6, …, 14]
  • 캐시 히트 길이 6, [3, 4, 5]로 prefill 계산 → [6, 7, …, 14]
  • 캐시 히트 길이 14, [11, 12, 13]으로 prefill 계산 → [14] (가장 효율적)

캐시 히트를 오른쪽에서 왼쪽으로 확인하고 일치하는 것을 찾으면 조기 종료할 수 있습니다. 이는 왼쪽에서 오른쪽으로 확인하고 실패 시 조기 종료하는 full attention과 반대입니다. (full attention과 비교해) 한 가지 단점은 일치가 없을 때 토큰 전체 목록을 반복하게 되는 것인데, 이는 흔한 경우입니다. 이는 무시할 수 없는 오버헤드를 유발할 수 있지만 아래에서 논의하듯 full + swa에서는 괜찮습니다.

사례 2: sliding window attention + full attention 모델 (Case 2: sliding window attention + full attention models)

첫 번째 문제는 캐시 히트 프리픽스를 어떻게 찾는가입니다. 전역 및 sliding window attention 레이어의 캐시 히트를 교집합으로 "intersect"해야 합니다:

  • full attention에 대한 가장 긴 캐시 히트를 얻습니다(왼쪽에서 오른쪽으로 스캔).
  • 그 길이 안에 있는 sliding window attention의 가장 긴 캐시 히트를 얻습니다. full attention의 캐시 히트 길이에서 시작해 오른쪽에서 왼쪽으로 캐시 히트를 확인해 구현합니다.

결과적인 sliding window attention 레이어의 캐시 히트가 full attention 레이어의 캐시 히트이기도 하다는 것을 보장할 수 있습니다. 이는 각 그룹의 모든 가능한 프리픽스를 찾아 교집합하는 것보다 효율적입니다. 캐시 히트가 없으면 조기 종료할 수 있기 때문입니다.

이 알고리즘은 정확히 두 가지 attention 유형(full attention + X)을 가진 모델에 적용됩니다. 여기서 X는 sliding window, llama 4 local attention, mamba 같은 임의의 효율적 attention 알고리즘이 될 수 있습니다. full attention 레이어가 없는 모델과 2개보다 많은 attention 유형을 가진 모델은 지원하지 않습니다. 이 문서를 쓰는 시점의 대부분 하이브리드 모델에는 충분합니다.

두 번째 질문은 캐시 제거 정책입니다. 현재 모든 kv cache group에 LRU 큐 하나를 사용합니다. 블록은 요청이 끝났거나 블록이 슬라이딩 윈도우 밖으로 나갔을 때 해제되며 LRU 큐에 추가됩니다.

사례 3: mamba 모델 (Case 3: mamba models)

mamba 모델의 프리픽스 캐싱 지원은 진행 중입니다. 구현되면 mamba 레이어 + full attention 레이어 모델을 사례 2의 full attention + X 알고리즘으로 지원할 수 있습니다.

구현 (Implementation)

개요 (Overview)

KVCacheManager는 3개 계층으로 구성됩니다:

  • KVCacheManager: 스케줄러와 kv cache 관리 시스템 사이의 인터페이스.
  • KVCacheCoordinator: 그룹별 SingleTypeKVCacheManager를 조정해 요청의 할당 결과를 생성합니다. 모델 구성에 따라 다음 코디네이터 중 하나가 선택됩니다:
    • KVCacheCoordinatorNoPrefixCache: 프리픽스 캐싱이 비활성화됐을 때 사용.
    • UnitaryKVCacheCoordinator: KV cache group이 하나뿐일 때. 교집합이 필요 없어 프리픽스 캐싱 로직이 단순화됨.
    • HybridKVCacheCoordinator: 정확히 두 KV cache group을 처리(반드시 full-attention 그룹 하나 + 다른 효율적 attention 그룹 하나 포함). 다른 경우는 구현되지 않습니다. 프리픽스 캐싱을 비활성화해 KVCacheCoordinatorNoPrefixCache를 사용할 수 있습니다.
  • SingleTypeKVCacheManager: 각 인스턴스가 하나의 KV cache group에 대한 할당과 프리픽스 캐싱을 관리하며 attention 유형별 로직(full attention, sliding window, Mamba)을 구현합니다.

위 그림의 파란 상자는 full attention 레이어 10개와 sliding window attention 레이어 20개가 있는 경우를 보여줍니다. 따라서:

  • HybridKVCacheCoordinator를 사용합니다.
  • 3개 KVCacheGroup에 대해 FullAttentionManager 1개와 SlidingWindowManager 2개를 사용합니다.

메모리 레이아웃 (Memory Layout)

각각 m개 레이어를 가진 n개 KVCacheGroup이 있는 모델에 대해 m개 버퍼를 할당합니다. 각 버퍼는 각 그룹에서 하나씩 가져온 n개 레이어가 공유합니다.

다음 그림은 full attention 레이어 10개(full.0 - full.9)와 sliding window attention 레이어 20개(sw.0-sw.19)가 있는 모델에 대한 것입니다. "Allocation" 절의 "case 2"를 따르며 3개 그룹으로 나뉩니다:

  • Group 0: full attention 레이어 10개 (full.0 - full.9)
  • Group 1: sliding window attention 레이어 10개 (sw.0 - sw.9)
  • Group 2: sliding window attention 레이어 10개 (sw.10 - sw.19)

요청에 대해 group 0에 block_id 0-6인 블록 11개, group 1에 7-8, group 2에 9-10을 할당합니다.

이 예에서 물리 메모리는 10개 버퍼(KVCacheTensor 0 - KVCacheTensor 9)로 나뉩니다. 각 버퍼는 3개 레이어가 공유하며(예: KVCacheTensor 0은 group 0의 full.0, group 1의 sw.0, group 2의 sw.10이 공유) 크기 block_size * kv_hidden_size 조각으로 나뉩니다. 이 3개 attention 레이어의 KV 캐시는 할당된 block_ids에 따라 버퍼의 서로 다른 조각에 저장됩니다.

참고: 하나의 논리 "block"은 물리 메모리의 10개 버퍼에 있는 10개 조각으로 매핑됩니다.

더 알아보기 (Learn more)