[vLLM] Automatic Prefix Caching: 접두사 캐싱
들어가며
LLM 서빙에서 동일한 시스템 프롬프트를 가진 요청이 반복적으로 들어오는 것은 매우 흔한 패턴이다. 예를 들어 RAG 파이프라인에서는 동일한 지시문 + 다른 문서 조합이, 챗봇 서비스에서는 동일한 시스템 프롬프트 + 다른 사용자 메시지가 반복된다. **Automatic Prefix Caching(APC)**은 이 공통 접두사의 KV 캐시를 자동으로 재활용하여, 중복된 prefill 연산을 완전히 건너뛴다.
- 논문: ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition
- 공식 문서: https://docs.vllm.ai
공식 문서
vLLM 공식 문서: Automatic Prefix Caching
핵심 구조/코드 분석
KVCacheManager: 중앙 관리자
vllm/v1/core/kv_cache_manager.py의 KVCacheManager가 전체 KV 캐시의 할당, 캐싱, 해제를 관리한다:
class KVCacheManager:
def __init__(
self,
kv_cache_config: KVCacheConfig,
max_model_len: int,
hash_block_size: int,
enable_caching: bool = True,
use_eagle: bool = False,
log_stats: bool = False,
...
) -> None:
self.enable_caching = enable_caching
self.coordinator = get_kv_cache_coordinator(
kv_cache_config=kv_cache_config,
max_model_len=self.max_model_len,
enable_caching=self.enable_caching,
hash_block_size=hash_block_size,
...
)
self.block_pool = self.coordinator.block_pool
enable_caching 플래그로 APC를 활성화/비활성화할 수 있다. 실제 캐시 탐색과 관리는 coordinator에 위임된다.
KVCacheBlocks: 블록 추상화
KV 캐시는 고정 크기 블록으로 관리된다:
@dataclass
class KVCacheBlocks:
blocks: tuple[Sequence[KVCacheBlock], ...]
"""blocks[i][j]는 i번째 kv_cache_group의 j번째 블록을 가리킨다."""
def get_block_ids(self) -> tuple[list[int], ...]:
return tuple([blk.block_id for blk in group] for group in self.blocks)
def get_unhashed_block_ids(self) -> list[int]:
return [block.block_id for block in self.blocks[0]
if block.block_hash is None]
각 블록은 block_hash를 가질 수 있다. 해시가 있는 블록은 이미 접두사 캐시에 등록된 것이고, None이면 아직 해싱되지 않은 새 블록이다.
get_computed_blocks: 캐시 히트 탐색
APC의 핵심 로직은 get_computed_blocks 메서드이다:
def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]:
# 접두사 캐싱이 비활성화되었거나 요청이 캐시 읽기를 건너뛰면 빈 결과 반환
if not self.enable_caching or request.skip_reading_prefix_cache:
return self.empty_kv_cache_blocks, 0
# 모든 토큰이 캐시에 있을 때도 마지막 토큰은 재연산 필요 (로짓 생성)
max_cache_hit_length = request.num_tokens - 1
computed_blocks, num_new_computed_tokens = (
self.coordinator.find_longest_cache_hit(
request.block_hashes, max_cache_hit_length
)
)
if self.log_stats:
self.prefix_cache_stats.record(
num_tokens=request.num_tokens,
num_hits=num_new_computed_tokens,
preempted=request.num_preemptions > 0,
)
return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens
request.block_hashes는 요청의 토큰 시퀀스를 블록 크기로 나누어 해시한 값이다. find_longest_cache_hit이 이 해시 리스트와 기존 캐시를 매칭하여 가장 긴 공통 접두사를 찾는다.
주목할 점은 max_cache_hit_length = request.num_tokens - 1이다. 모든 토큰이 캐시에 있더라도 마지막 토큰은 반드시 재연산한다. 이는 로짓을 생성하기 위해 최소 1토큰의 forward pass가 필요하기 때문이다.
allocate_slots: 블록 할당
새 토큰에 대한 KV 캐시 블록을 할당하는 로직이다:
def allocate_slots(
self,
request: Request,
num_new_tokens: int,
num_new_computed_tokens: int = 0,
new_computed_blocks: KVCacheBlocks | None = None,
num_lookahead_tokens: int = 0,
num_external_computed_tokens: int = 0,
delay_cache_blocks: bool = False,
...
) -> KVCacheBlocks | None:
num_external_computed_tokens는 Disaggregated Serving에서 외부 노드로부터 전송받은 토큰 수를 의미한다. APC는 분리된 서빙 환경에서도 동작하도록 설계되어 있다.
can_fit_full_sequence: 어드미션 게이트
def can_fit_full_sequence(
self, request, num_new_computed_tokens=0,
new_computed_blocks=None, ...
) -> bool:
num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
request_id=request.request_id,
num_tokens=full_num_tokens,
new_computed_blocks=new_computed_block_list,
...
)
return num_blocks_to_allocate <= self.block_pool.get_num_free_blocks()
Chunked Prefill에서 첫 청크만 확인하면 전체 시퀀스를 수용할 수 없는 경우가 발생할 수 있다. 이 메서드는 전체 시퀀스에 필요한 블록 수를 미리 계산하여, 과다 수용(over-admission)을 방지한다.
왜 이 설계인가
1. 해시 기반 매칭: 토큰 시퀀스 전체를 비교하는 대신 블록 해시를 비교하면 O(1)에 캐시 히트를 판단할 수 있다. 블록 크기가 16토큰이면 1024토큰 프롬프트는 64번의 해시 비교로 충분하다.
2. 블록 단위 관리: 토큰 단위가 아닌 블록 단위로 캐시를 관리하면 메타데이터 오버헤드가 크게 줄어든다. 또한 GPU 메모리 할당도 블록 단위로 이루어져 단편화를 방지한다.
3. skip_reading_prefix_cache: 프롬프트 로그 확률(prompt logprobs)이 필요하거나 풀링 모델인 경우, 캐시된 결과를 사용하면 로짓 정보를 얻을 수 없다. 이런 경우 캐시 읽기를 건너뛰는 것이 올바른 동작이다.
4. 통계 수집: PrefixCacheStats로 히트율, 토큰 수, 선점 여부를 추적한다. 이 정보는 Prometheus 메트릭으로 노출되어 프로덕션 모니터링에 활용된다.
논문 핵심 내용
SGLang: Efficient Execution of Structured Language Model Programs (2312.07104) 논문은 RadixAttention을 포함한 LLM 프로그램 실행 최적화 시스템을 제안했다.
핵심 아이디어: RadixAttention은 radix tree(기수 트리) 자료구조로 KV 캐시를 관리하여, 공통 접두사를 가진 요청들의 KV 캐시를 자동으로 재사용한다. 캐시 인식 스케줄링(cache-aware scheduling)과 결합하여 최적 히트율에 가까운 성능을 달성한다.
캐시 히트율
| 환경 | 모델 | 캐시 히트율 |
|---|---|---|
| 프로덕션 (Chatbot Arena) | LLaVA-Next-34B | 52.4% |
| 프로덕션 (Chatbot Arena) | Vicuna-33B | 74.1% |
| 벤치마크 범위 | 다양한 태스크 | 50% ~ 99% |
| 캐시 인식 스케줄링 | - | 최적의 96% |
성능 향상
| 메트릭 | 향상 |
|---|---|
| 최대 처리량 | 기존 SOTA 대비 6.4x |
| JSON 디코딩 (FSM 압축) | 1.6x |
| 멀티모달 모델 | 최대 6x |
| LLaVA-v1.5-7B 이미지 처리 | 0.18 → 1.15 images/sec |
| 첫 토큰 지연 (Vicuna-33B) | 평균 1.7x 감소 |
| 전체 지연 | 최대 3.7x 감소 |
RadixAttention 자료구조 관리 오버헤드는 전체 실행 시간의 0.3% 미만(74.3초 중 0.2초)이다. 캐시 재사용 기회가 없는 워크로드에서도 거의 오버헤드가 없다는 뜻이다. 프로덕션 Chatbot Arena 환경에서 Vicuna-33B의 캐시 히트율이 74.1%라는 건, 실제 서비스에서 요청의 3/4이 중복 연산 없이 처리될 수 있다는 의미다. vLLM의 APC는 이 RadixAttention의 핵심 아이디어를 블록 해시 기반으로 구현한 것이다.
마무리
APC는 "동일한 접두사를 반복 계산하지 않는다"는 단순한 아이디어이지만, 해시 기반 블록 매칭, 어드미션 게이트, 분리 서빙 호환 등 정교한 구현을 통해 실질적인 TTFT(Time To First Token) 감소를 달성한다. 특히 시스템 프롬프트가 긴 RAG 파이프라인에서 효과가 극대화된다.
관련 포스트
vLLM 의 다른글
- 이전글 [vLLM] Structured Output: JSON/regex/문법 제약 생성
- 현재글 : [vLLM] Automatic Prefix Caching: 접두사 캐싱
- 다음글 [vLLM] Disaggregated Prefill/Decode: 분리된 서빙
댓글