본문으로 건너뛰기

[vLLM] Automatic Prefix Caching: 접두사 캐싱

들어가며

LLM 서빙에서 동일한 시스템 프롬프트를 가진 요청이 반복적으로 들어오는 것은 매우 흔한 패턴이다. 예를 들어 RAG 파이프라인에서는 동일한 지시문 + 다른 문서 조합이, 챗봇 서비스에서는 동일한 시스템 프롬프트 + 다른 사용자 메시지가 반복된다. **Automatic Prefix Caching(APC)**은 이 공통 접두사의 KV 캐시를 자동으로 재활용하여, 중복된 prefill 연산을 완전히 건너뛴다.

공식 문서

vLLM 공식 문서: Automatic Prefix Caching

핵심 구조/코드 분석

KVCacheManager: 중앙 관리자

vllm/v1/core/kv_cache_manager.pyKVCacheManager가 전체 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 의 다른글