콘텐츠로 이동

KV 캐시

목차

KV 캐시란?

Autoregressive decode에서 각 attention layer는 이전 토큰이 만든 key(K)와 value(V)를 반복해 사용합니다. KV 캐시는 이 텐서를 저장하여 매 decode step마다 전체 prefix를 다시 계산하지 않고 새 토큰의 projection만 수행하게 합니다.

캐시가 없으면 매 단계마다 길어진 prefix 전체를 처리하므로 생성 비용이 지나치게 커집니다. 캐시를 사용하면 저장 상태는 보이는 context에 따라 늘어나지만 일반 decode step의 토큰 projection은 새 토큰 하나만 처리합니다. Attention은 여전히 관련 prefix 또는 window를 읽으므로 전체 attention 비용 자체가 문자 그대로 상수라는 뜻은 아닙니다.

일반적인 grouped-query attention layer에서 K와 V의 형상은 [batch, num_kv_heads, sequence, head_dim]입니다. Query head 수는 KV head 수보다 클 수 있으므로 캐시 용량 계산에 query head 수를 사용하면 안 됩니다.

메모리 증가량

모든 layer의 차원이 같다면 시퀀스 하나에서 토큰당 캐시 증가량은 대략 다음과 같습니다.

2 × num_layers × num_kv_heads × head_dim × bytes_per_element

계수 2는 K와 V를 모두 포함합니다. Layer별 차원이 다른 아키텍처는 각 layer의 K/V 크기를 합산해야 합니다. Batch size, quantized-cache metadata, 정렬, allocator 여유 공간, paged block 단위, 임시 텐서 때문에 실제 사용량은 이 추정치보다 큽니다.

표준 캐시는 토큰마다 재할당하지 않도록 buffer를 일정 단위로 늘립니다. 따라서 할당 크기는 offset이 나타내는 채워진 길이보다 클 수 있습니다.

세 가지 구조적 캐시 레이아웃

src/lib/mlxcel-core/src/cache.rs는 세 가지 attention-cache 레이아웃을 정의합니다. 이 레이아웃은 사용자가 선택하는 KV 저장 모드(예: FP16, uniform quantization, TurboQuant) 및 서버의 dense/paged 저장 backend와 별개의 개념입니다.

표준 KVCache

KVCache는 보이는 prefix 전체를 유지합니다. 일반적인 full-attention layer가 사용하며, 투기적 디코딩에서 수락되지 않은 토큰을 되돌릴 수 있도록 trim()을 제공합니다. Buffer는 chunk 단위로 증가하고 가능한 곳에서는 in-place slice update를 사용합니다.

RotatingKVCache

RotatingKVCache는 sliding-window attention을 구현합니다. 논리 window가 가득 차면 가장 오래된 항목을 덮어쓰고 최근 window를 선형화해 attention에 제공합니다. 메모리는 설정된 window와 구현상 할당 여유 범위 안에서 제한되지만, 매 순간 정확히 window_size × bytes_per_token과 같다고 보장하지는 않습니다.

대표 사용 모델은 Gemma ¾, ExaOne 4와 ExaOne MoE, GPT-OSS, Ministral 3, Mellum, Muse Glimmer, RecurrentGemma, Step 3.5, AFMoE입니다. 일부는 rotating layer와 full-attention layer를 함께 사용합니다.

ChunkedKVCache

ChunkedKVCache는 전역 토큰 offset과 제한된 visible chunk의 시작 위치를 함께 추적합니다. Llama 4의 interleaved/chunked attention layer가 사용하며 같은 모델의 full-attention layer는 표준 캐시를 계속 사용합니다.

Jamba의 Mamba layer는 ChunkedKVCache를 사용하지 않습니다. Jamba는 모델 소유 Mamba 순환 상태를 유지하고 attention layer에만 표준 KVCache를 사용합니다.

아키텍처별 선택

각 모델의 make_caches() 구현이 attention layer별 레이아웃을 선택합니다. 하이브리드 모델은 layer마다 다른 레이아웃을 반환할 수 있습니다.

아키텍처 패턴 상태 레이아웃
대부분의 full-attention transformer 표준 KVCache
Gemma ¾ 같은 sliding-window/global hybrid Window layer는 RotatingKVCache, global layer는 KVCache
Llama 4 Interleaved layer는 ChunkedKVCache, full-attention layer는 KVCache
Jamba 모델 소유 Mamba 상태 + 표준 attention KVCache
RecurrentGemma 모델 소유 RGLRU 상태 + rotating attention cache
Mamba, Mamba2, RWKV Attention KV 캐시가 아닌 모델 소유 순환 상태

정확한 조합은 로드된 아키텍처의 구현 속성이며 전역 명령행 선택 항목이 아닙니다.

시퀀스별 상태와 CachePool

CLI와 서버 모두 디코딩 Engine이 소유한 CachePool에 캐시 상태를 둡니다. mlxcel generate는 DirectEngine으로 실행되며, 생성마다 이 풀에 시퀀스 하나를 열고 끝나면 닫습니다. 서버는 요청마다 단조 증가하는 SequenceId를 할당하고 CachePool로 시퀀스별 캐시 세트, prompt 길이, offset, 할당 시점을 추적합니다.

SequenceId는 u64이므로 합리적인 서버 수명 동안 wrap될 가능성은 사실상 없지만 수학적으로 “절대 wrap되지 않는다”는 보장은 아닙니다.

CachePool 자체는 의도적으로 Sync가 아닙니다. 여러 thread가 공유하려면 호출자가 외부 Mutex/RwLock 동기화를 제공해야 합니다. 할당은 max_sequences를 적용하고, release는 끝난 상태를 제거하며, memory_usage_bytes()는 할당된 K/V buffer를 합산합니다. 포화된 요청이 대기할지 busy 응답을 받을지는 CachePool 단독이 아니라 서버 admission queue가 결정합니다.

모델 소유 상태를 쓰는 아키텍처에서는 pool이 scheduler bookkeeping용 경량 placeholder만 담고, 실제 상태는 모델이 SequenceId별로 직접 보관합니다(LanguageModel::sequence_state_layout, prepare_sequence_state, sync_sequence_storage). Scheduler는 모델이 supports_batching()을 보고할 때만 이런 시퀀스 여러 개를 하나의 배치 forward 패스에 넣고, 보고하지 않는 모델은 시퀀스 1개로 제한합니다.

순환·하이브리드 모델

Mamba, Mamba2, RWKV 및 하이브리드 아키텍처는 표준 KV 캐시처럼 전체 context 길이에 따라 커지지 않는 순환 상태를 유지합니다. 이 상태는 모델이 소유하고 시퀀스 id로 구분하며, Mamba, Mamba2, Jamba, RWKV7은 supports_batching()과 BatchedDecodeKind::Native를 보고하므로 서버는 이들의 시퀀스 여러 개를 행 격리를 유지한 채 하나의 배치 forward 패스로 decode합니다. supports_batching()을 보고하지 않는 계열은 순차 단일 슬롯 worker로 실행합니다.

Dense transformer는 지원되는 하드웨어에서 padded prefill을 선택하고 이후 외부 캐시를 실제 prompt 길이로 trim할 수 있습니다. attention 캐시를 모델이 직접 소유하는 계열은 LanguageModel::trim_state로 자기 상태를 trim하는데, Gemma 3가 이 방식으로 dense 모델처럼 padded prefill을 사용합니다. NemotronH 같은 순환 hybrid와 AFMoE, Llama 4, Gemma 4, DeepSeek V4, Bailing linear, Qwen3-Next는 supports_padded_prefill() == false를 반환하므로 정상 생성 경로에서 prefill을 padding하지 않습니다. MLXCEL_FORCE_PADDED_PREFILL은 어떤 하드웨어에서든 tile 정렬을 강제하고 MLXCEL_NO_PADDED_PREFILL은 이를 끄며, 둘 다 CLI와 서버에 적용되고 함께 설정하면 MLXCEL_NO_PADDED_PREFILL이 우선합니다. trim_internal_caches() hook은 어떤 경로가 초과 내부 상태를 버려야 할 때를 위한 방어적 모델별 장치입니다. NemotronH가 정상적으로 padded prefill을 사용한다는 뜻이 아닙니다.

서버의 decode lookahead는 모델이 투기적으로 붙인 단일 토큰 append를 정확히 되돌릴 수 있을 때(supports_decode_lookahead_rewind())만 모델 소유 계열을 파이프라인으로 처리합니다. Gemma 3는 sliding 레이어에서 decode 쓰기가 덮어쓰는 행을 기록해 두므로, 모든 sliding 레이어가 FP16이면 lookahead 파이프라인으로 decode합니다. 다른 모델 소유 계열은 동기식 decode를 유지합니다.

Attention 분기와 in-place decode 쓰기

attention 실행 방식은 모델이 아니라 캐시가 정합니다. 모델이 KVCache::attend(행 여러 개를 한 번에 처리할 때는 cache::attend_batched)를 호출하면, 캐시가 자기 상태의 저장 방식에 따라 dense SDPA와 paged decode 커널 중 하나를 고릅니다. RotatingKVCache::attend와 ChunkedKVCache::attend도 window 의미를 유지하면서 같은 분기를 적용합니다.

단일 토큰 FP16 decode는 slice update로 캐시를 다시 만들지 않고 새 K/V 행을 캐시 버퍼에 직접(in place) 씁니다. 기본으로 켜져 있으며, MLXCEL_KV_INPLACE_WRITE=0(또는 false, off, no)으로 slice update 방식으로 되돌립니다.

Paged decode 저장소와 커널

서버는 시퀀스별 KV 상태를 하나의 dense allocation 대신 고정 크기 page에 저장할 수 있습니다. 이 방식은 제한된 --kv-cache-budget 아래에서 block 재사용, eviction, preemption, prefix sharing을 가능하게 합니다. Paged storage는 architecture attention layer가 full, rotating, chunked 중 무엇인지와 별개이며 wrapper가 모델의 visible-window semantics를 보존합니다.

Production paged decode 경로는 fused paged attention, sparse page-indirection decode, 충분히 긴 whole-page prefix를 공유하는 group을 위한 opt-in cascade attention을 사용할 수 있습니다. 주요 제어 변수는 MLXCEL_PAGED_DECODE_SPLITS, MLXCEL_SPARSE_PAGED_ATTENTION, MLXCEL_SPARSE_PAGED_MIN_SPARSITY, MLXCEL_CASCADE_ATTENTION, MLXCEL_CASCADE_MIN_MEMBERS, MLXCEL_CASCADE_MIN_SHARED_PAGES입니다. 기본값을 바꾸기 전에 성능 튜닝을 참조하십시오.

Matrix-absorbed MLA decode

DeepSeek 계열 multi-head latent attention은 확장된 head별 K/V 대신 압축 latent를 캐시할 수 있습니다. MLXCEL_MLA_ABSORBED=1이 이 저메모리 경로를 활성화합니다. 흡수는 multi-token 단계에도 적용됩니다. mlxcel_core::mla::use_absorbed_mla가 호출마다 흡수 비용과 윈도우 materialize 비용을 비교하므로, 긴 캐시 뒤에서 이어지는 chunked-prefill 구간이나 투기적 검증 블록은 흡수 경로로 실행되고 콜드 프리필은 materialize 경로를 유지해 기준 구현과 같은 수치를 냅니다. MLXCEL_MLA_SPLIT_KV=1은 긴 latent range를 추가로 분할하며 absorbed MLA가 필요합니다. 둘 다 기본값은 off입니다.

투기적 디코딩의 rewind

투기적 디코딩은 여러 토큰을 제안하고 target 모델로 검증합니다. 검증이 prefix 일부만 수락하면 수락 위치 이후의 캐시 상태를 제거해야 합니다. KVCache::trim(n)은 논리 offset을 되돌리고 투기적 round 루프(프롬프트 조회와 풀 모델 드래프터용 엔진 토큰 전용 루프, MTP·DFlash 루프)는 각 round에 필요한 target·draft 상태를 trim합니다. 다른 캐시 레이아웃과 모델 소유 drafter 경로는 모든 캐시를 단순한 grow-only dense tensor로 가정하지 않고 자체 position/window 불변식을 유지합니다.