KV 캐시 시스템 (KV Cache System)

KV 캐시 시스템 (KV Cache System)

LLM이 문장을 생성할 때마다 앞서 계산한 키·값을 다시 계산하지 않도록 저장해 두는 게 KV 캐시예요. TensorRT-LLM의 KV 캐시 시스템은 여기서 한 걸음 더 나아가, 요청 간 캐시 재사용을 지원하고 오프로딩과 우선순위 기반 축출 같은 기법으로 재사용률을 높여요. 또 제한된 어텐션 윈도우 크기와 MHA·MQA·GQA 같은 멀티 헤드 어텐션 최적화도 함께 다룹니다.

출처: 공식 문서 - KV Cache System

요청 간 재사용 (Reuse Across Requests)

블록은 검색 트리에서 축출되기 전까지 재사용 가능해요. 축출은 새(빈) 블록이 필요할 때 일어나는데, 핵심 축출 방식은 우선순위 기반 LRU예요. 모든 블록에는 0부터 100까지의 우선순위가 붙고, 100이 가장 중요함을 뜻합니다.

보존 정책 (Retention Policy)

블록 우선순위는 요청의 보존 정책(KvCacheRetentionConfig)을 따라 부여돼요. 우선순위 점수가 낮은 블록이 높은 블록보다 먼저 해제되죠. 보존 정책은 TokenRangeRetentionConfig 객체들의 리스트로, 예를 들어 "토큰 10부터 61까지는 우선순위 X를 부여"처럼 특정 토큰 구간에 우선순위를 지정합니다.

스펙큘레이티브 디코딩 (Speculative Decoding)

모든 스펙큘레이티브 디코딩 모델에서 요청 간 재사용이 지원돼요. 자세한 내용은 스펙큘레이티브 디코딩 문서를 참고하면 됩니다.

제한된 어텐션 윈도우 크기 (Limited Attention Window Size)

TensorRT-LLM은 어텐션 윈도우 크기가 제한된 레이어를 활용해 연산과 메모리 사용을 줄여요. 어텐션 윈도우를 벗어난 블록은 해제되어 라디스(radix) 검색 트리로 들어가 재사용을 기다립니다.

MQA / GQA

TensorRT-LLM은 그룹 쿼리 어텐션(GQA)을 활용해 메모리를 절약해요. KV 캐시는 헤드의 개별 쿼리 그룹 상태만 담을 수 있도록 블록을 만들죠. MHA는 헤드당 하나의 그룹, MQA는 모든 헤드에 단 하나의 그룹이고, GQA는 그 사이의 균형을 잡습니다.

KV 캐시 동작 제어 (Controlling KV Cache Behavior)

KV 캐시 시스템의 많은 기능은 선택 사항이거나 사용자가 정의한 속성으로 동작을 바꿀 수 있어요. KvCacheConfig 클래스를 통해 KV 캐시 기능을 제어합니다.

데이터 타입 (Datatype)

가장 중요한 속성은 캐시에 저장될 데이터 타입을 정하는 dtype이에요. 기본값 auto는 모델 설정에서 데이터 타입을 자동으로 추론합니다.

KV 캐시에 할당할 메모리 크기

free_gpu_memory_fraction 속성은 0보다 크고 1보다 작은 비율로, 비어 있는 GPU 메모리 중 얼마만큼을 KV 캐시에 할당할지 정해요. 기본값은 90% 즉 0.9입니다. max_tokens도 설정되어 있다면 KV 캐시는 max_tokens을 담는 데 필요한 메모리를 계산하고, max_tokens와 free_gpu_memory_fraction 중 더 작은 값을 할당해요.

요청 간 재사용 켜고 끄기

요청 간 블록 재사용은 기본적으로 켜져 있는데, enable_block_reuse를 False로 두면 끌 수 있어요. scheduler_config.enable_prefix_aware_scheduling은 스케줄러 쪽에서 접두사 재사용 추정치를 쓸지 여부만 제어합니다. True(기본값)면 스케줄러가 재사용 가능한 KV 토큰 추정치를 이용해 중복되는 첫 번째 청크 컨텍스트 요청을 미루고, 캐시된 접두사 블록을 재사용할 것으로 예상되는 요청에는 토큰 예산 계산을 줄여줘요.

런타임 KV 블록 재사용은 계속 켜두면서 접두사 인지 스케줄링만 끄려면 이렇게 하면 됩니다.

kv_cache_config:
  enable_block_reuse: true
scheduler_config:
  enable_prefix_aware_scheduling: false

KV 캐시 매니저 선택

TensorRT-LLM은 두 가지 KV 캐시 매니저 구현을 제공해요. use_kv_cache_manager_v2가 둘 사이를 고르며 기본값은 auto로, 모델이 선호하는 방식을 따르다가 선언하지 않은 모델은 V1 C++ 매니저로 폴백합니다. true나 false로 두면 모델 기본값을 덮어씁니다. auto 아래에서 모델 기본 V2가 해당 조합에 대해 V1로 폴백하거나, use_kv_cache_manager_v2: true를 명시하면 오류가 발생하기도 해요.

Mamba 스냅샷 경계

per_conversation 블록 재사용 정책은 주기적 Mamba 스냅샷을 비활성화해요. 하이브리드 Mamba 모델과 함께 쓸 때는 안정적인 경계를 하나 이상 명시적으로(보통 끝 오프셋 0) 구성해야 합니다. 예를 들면 이렇게요.

kv_cache_config:
  enable_block_reuse: true
  use_kv_cache_manager_v2: true
  avg_seq_len: 2048
  block_reuse_config:
    policy: per_conversation
    max_num_turns: 2
  mamba_state_config:
    periodic_snapshot_interval: 0
    additional_snapshot_offsets_from_start: [128]

pool_ratio는 레이어 그룹 ID 순서대로 그룹당 하나씩, 양수로 정규화된 캐시 계층 할당 가중치를 담아요. avg_seq_len도 pool_ratio도 설정하지 않으면 하이브리드 Mamba 모델은 경고를 띄우고 max_seq_len의 절반으로 폴백하는데, 이는 최적이 아닌 풀 분할을 만들 수 있습니다.

캐시 식별용 멀티모달 UUID 지원

비전-언어 모델처럼 멀티모달 모델을 다룰 때는 어떤 캐시 블록이 어떤 멀티모달 입력(이미지, 영상 등)에 대응하는지 식별해야 해요. 기본적으로 시스템은 각 멀티모달 입력에 콘텐츠 기반 해시로 고유 식별자를 만듭니다. 하지만 세션 간 캐시 관리에는 한계가 있어서, 같은 콘텐츠도 같은 해시를 만들려면 다시 처리해야 하죠. 결정적 캐시 관리를 위해 multi_modal_uuids 파라미터로 나만의 UUID 문자열을 줄 수 있어요. UUID를 주면 KV 캐시 이벤트에는 계산된 콘텐츠 해시 대신 그 UUID가 반환되고, 캐시 키 자체는 정확성을 위해 UUID와 콘텐츠 둘 다로 계산됩니다.

from tensorrt_llm.inputs import TextPrompt

# Provide custom UUIDs for your images
prompt = TextPrompt(
    prompt="Describe these images.",
    multi_modal_data={"image": [image1, image2]},
    multi_modal_uuids={"image": ["image-uuid-001", "image-uuid-002"]}
)

주요 특징을 정리하면 이래요.

  • 캐시 정확성: UUID를 주면 캐시 키는 UUID와 콘텐츠를 함께 BLAKE3(UUID || Content)로 계산해요. 같은 UUID라도 콘텐츠가 다르면 캐시 항목이 달라집니다.
  • 사용자 격리: 같은 콘텐츠라도 UUID가 다르면 캐시 항목이 달라져, 사용자별·세션별 캐시 격리가 가능해요.
  • 안정적인 이벤트 식별자: 원래 UUID 문자열이 get_kv_cache_events()를 통한 KV 캐시 이벤트에 보존되어 반환돼요.
  • 부분 UUID 지원: 일부 항목에는 UUID를, 나머지는 None으로 둬서 콘텐츠 전용 해시로 폴백할 수 있어요.
  • 교차 모달 지원: 이미지와 영상처럼 서로 다른 모달리티가 각자의 UUID를 가질 수 있어요.

UUID 형식은 제한이 없어요. "image-123", "user-session-img-a" 같은 문자열이나 데이터베이스 키를 써도 되고, 원래 UUID 문자열은 KV 캐시 이벤트에 보존되어 반환됩니다.

호스트 메모리로의 오프로딩 (Enable Offloading to Host Memory)

블록이 GPU 메모리에서 축출되기 전에 선택적으로 호스트(CPU) 메모리로 오프로딩할 수 있어요. 블록은 호스트 메모리에서 축출되기 전까지 재사용 가능하고, 오프로딩된 블록을 재사용할 때는 먼저 GPU 메모리로 복사해 옵니다. 오프로딩은 얼마나 많은 호스트 메모리(바이트)를 할당할지 정하는 host_cache_size 속성으로 제어하며 기본값은 0이에요. 오프로딩을 켜면 클라이언트가 블록 우선순위를 조절해 특정 블록이 오프로딩되지 않도록 막을 수 있어요. 특정 임계값보다 우선순위가 낮은 블록은 오프로딩되지 않고 GPU 메모리에서 바로 축출되어 GPU와 호스트 사이 트래픽을 줄입니다. 이 우선순위는 secondary_offload_min_priority로 정하며 기본값은 35로, 35보다 낮은 우선순위 블록은 오프로딩되지 않아요.

부분 재사용 (Partial Reuse)

일부 토큰만 일치할 때 블록을 부분적으로 재사용할 수 있어요. 기본적으로 켜져 있고, enable_partial_reuse를 False로 두면 끕니다. copy_on_partial_reuse 속성은 부분 재사용을 위해 블록을 복사할지 여부를 정해요. 복사를 끄면 부분 일치 블록은 다른 요청이 사용 중이지 않을 때만 재사용할 수 있어요.

어텐션 윈도우 크기

max_attention_window 속성은 모델의 각 레이어별 최대 어텐션 윈도우 크기를 정수 리스트로 지정해요. 리스트 길이가 레이어 수보다 짧으면 필요한 만큼 반복됩니다. 예를 들어 모델이 전부 전체 어텐션 레이어이고 최대 시퀀스 길이가 4096이라면 max_attention_window = [4096]으로 지정하면 돼요. 첫 레이어가 전체 어텐션, 두 번째가 윈도우 256의 제한 어텐션이고 이후 반복된다면 max_attention_window = [4096,256]으로 지정해서, 1·3·5번째 레이어는 전체, 2·4·6번째 레이어는 제한 어텐션이 되게 합니다.

더 알아보기 (Learn more)