Prefix Caching¶

Overview¶

Prefix caching은 요청들이 공유하는 prefix의 캐시된 key-value(KV) 상태를 재사용하여 prefill 계산량과 첫 토큰까지의 시간을 줄일 수 있습니다.

Sub-block prefix caching¶

Upstream vLLM은 완전한 KV 캐시 블록 경계까지만 캐시된 prefix를 매칭합니다. RBLN처럼 블록이 크면 마지막 블록을 채우지 못한 공유 토큰은 다시 계산해야 합니다. RBLN vLLM은 이 토큰들을 sub-block이라는 더 작은 단위로 재사용하기 위해 sub-block prefix caching을 추가합니다.

KV 캐시 블록 크기가 4,096, sub-block 크기가 128이고, 캐시된 3,000개 토큰의 prefix를 더 긴 프롬프트가 공유하는 경우입니다.

새 요청에서 추가된 토큰은 두 경우 모두 계산합니다.

Configuration¶

sub_block_size는 sub-block 크기를 토큰 단위로 지정합니다. 기본값은 prefill chunk size인 max_num_batched_tokens입니다.

Python에서는 아래 예시처럼 additional_config["sub_block_size"]로 설정합니다. CLI에서는 --rbln-sub-block-size를 사용합니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

from vllm import LLM

llm = LLM(
    model="<MODEL_ID>",
    model_impl="<MODEL_IMPL>",  # "vllm" 또는 "transformers"
    enable_prefix_caching=True,
    max_num_batched_tokens=128,
    block_size=4096,
    additional_config={
        "sub_block_size": 128,
    },
)

Limitation: KV cache allocation and copying¶

Upstream vLLM은 완전한 KV 캐시 블록을 복사 없이 그대로 재사용합니다. RBLN vLLM의 vllm model implementation도 full-block 매칭에서는 같은 방식으로 동작합니다.

반면 sub-block 매칭에서는 두 model implementation 모두 별도 블록을 할당하고 캐시된 KV 데이터를 그 블록으로 복사합니다. 재계산은 생략할 수 있지만, 할당과 복사 오버헤드가 추가됩니다.

Prefix cache-aware routing¶

Prefix cache-aware router는 각 요청을 해당 prefix가 이미 캐시된 엔진으로 보냅니다.

Sub-block caching이 활성화되면 sub-block 단위로 KV 캐시 이벤트를 발행합니다. 라우팅은 다음과 같이 설정합니다.

1. Upstream vLLM의 KVEventsConfig 설정을 사용하여 --kv-events-config로 이벤트 발행을 활성화하고, 라우터가 이벤트를 수신하도록 설정합니다.
2. 라우터의 토큰 처리 블록 크기를 엔진의 sub-block 크기에 맞춥니다. --rbln-sub-block-size를 지정하지 않으면 엔진은 --max-num-batched-tokens 값을 sub-block 크기로 사용합니다.

Sub-block caching이 비활성 상태이면 일반 KV 캐시 이벤트에 맞춰 --block-size를 사용합니다.

Interaction with KV connectors¶

스케줄러는 요청마다 sub-block caching과 KV connector 중 하나만 사용하며, 두 방식을 결합하지 않습니다. 추가로 재사용할 수 있는 캐시된 토큰이 더 많은 쪽을 선택하고, 같으면 sub-block caching을 사용합니다.

KV connector는 KV 캐시 오프로딩과 disaggregated prefill / decode에 사용됩니다.

Upstream prefix caching 동작은 vLLM APC guide와 prefix caching design에서 확인할 수 있습니다.