원문 보기
KVzap: Fast, Adaptive, and Faithful KV Cache Pruning
Paper

KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

Maximilian Jeblick
2026.01.20
·Arxiv·by 이호민
#KV Cache Pruning#LLM Inference#Transformer#KVzip#KVzap

핵심 포인트

  • 1기존 KV 캐시 pruning 방법들의 속도 및 적용 한계를 극복하기 위해, 본 논문은 KVzip을 빠르고 적응적으로 근사하는 KVzap을 제안합니다.
  • 2KVzap은 모델의 hidden state로부터 중요도 점수를 예측하도록 학습된 경량 surrogate model(linear layer 또는 MLP)을 사용하여, 임계값 𝜏 이하의 KV 쌍을 제거함으로써 캐시를 압축합니다.
  • 3Qwen3-8B, Llama-3.1-8B-Instruct, Qwen3-32B 모델에서 다양한 long-context 및 reasoning task 전반에 걸쳐 KVzap은 2–4배의 KV 캐시 압축률을 달성하면서도 정확도 손실을 무시할 수 있는 수준으로 유지하며, SOTA 성능을 보여줍니다.

KVzap은 transformer 기반 LLM 추론에서 주요 병목인 KV 캐시의 메모리 사용량과 처리량을 줄이기 위한 빠르고 적응적이며 신뢰성 있는(faithful) pruning 방법론입니다. 기존 KV 캐시 pruning 연구는 속도-정확도 trade-off와 실제 시스템 통합의 어려움으로 인해 널리 채택되지 못했습니다. 특히, 현재까지 가장 좋은 성능을 보이는 KVzip(Kim et al., 2025)은 prefilling 단계에서 입력 프롬프트의 두 배 길이에 해당하는 context를 두 번 처리해야 하는 비효율성과 decoding 단계에서의 사용 불가라는 한계가 있었습니다. KVzap은 이러한 KVzip의 단점을 극복하면서도 유사하거나 더 나은 성능을 달성합니다.

1. KVzip 및 KVzip+ (Core Methodology Detail)

KVzap은 KVzip의 향상된 버전을 근사하는 방식으로 작동합니다.

  • KVzip:
KVzip은 "Repeat the previous context exactly."라는 pretext task를 사용하여 각 KV pair의 중요도를 측정합니다. 주어진 <prompt><prompt>에 대해, 모델이 이 <prompt><prompt>를 반복할 때 원본 <prompt><prompt>의 position ii에 해당하는 토큰에 얼마나 많은 attention을 주는지 측정합니다. 수식은 다음과 같습니다:
si=max⁡j∈<prompt>ajis_i = \max_{j \in \text{<prompt>}} a_{ji}
여기서 sis_i는 position ii의 KV pair 점수이고, ajia_{ji}는 반복되는 <prompt><prompt>의 position jj에서 원본 <prompt><prompt>의 position ii로 향하는 attention weight입니다. 이 점수가 낮으면 해당 KV pair가 정보를 거의 담고 있지 않으므로 제거할 수 있다는 가정에 기반합니다.

  • KVzip+:
KVzap의 개발자들은 Devoto et al. (2025)의 분석을 바탕으로 KVzip의 scoring 방식을 개선했습니다. 이 개선된 KVzip+는 각 토큰의 기여도를 출력 hidden state에 미치는 영향의 크기로 정규화합니다. Transformer head의 decoding step jj에서 hidden-state update는 hout,j=hj+∑i≤jajiWOvi\mathbf{h}_{out,j} = \mathbf{h}_j + \sum_{i \leq j} a_{ji} \mathbf{W}_O \mathbf{v}_i로 주어지며, 여기서 ajiWOvia_{ji} \mathbf{W}_O \mathbf{v}_i는 토큰 ii가 residual stream hj\mathbf{h}_j에 기여하는 바를 나타냅니다. KVzip+ score는 이 정규화 항을 포함하여 다음과 같이 정의됩니다:
si+=max⁡j∈<prompt>aji∥WOvi∥∥hj∥s_i^+ = \max_{j \in \text{<prompt>}} a_{ji} \frac{\|\mathbf{W}_O \mathbf{v}_i\|}{\|\mathbf{h}_j\|}
이 정규화는 단순히 attention weight가 높은 것뿐만 아니라, 해당 value vector가 실제 hidden state에 미치는 영향력을 고려하여 보다 정확한 중요도를 반영합니다.

2. KVzap의 작동 방식

KVzap은 KVzip의 느린 속도와 decoding 단계 사용 불가라는 한계를 해결하기 위해, KVzip+의 score log⁡(si+)\log(s_i^+)를 직접적으로 예측하는 경량 surrogate 모델을 학습시킵니다.

  • Surrogate Model 학습:
각 transformer layer마다 별도의 surrogate 모델이 학습됩니다. 이 모델은 입력 hidden state ht∈RDh\mathbf{h}_t \in \mathbb{R}^{D_h} (여기서 DhD_h는 hidden dimension)를 받아, 각 KV head에 대한 HH개의 예측 점수 ∈RH\in \mathbb{R}^H를 출력합니다. 학습 데이터는 (hidden state, log⁡(s+)\log(s^+)) 쌍으로 구성되며, Nemotron-Pretraining-Dataset-sample에서 1.2M 쌍을 추출하여 사용했습니다. surrogate 모델은 두 가지 종류가 제안됩니다:
  • KVzap-Linear: 단일 linear layer로 구성됩니다.
  • KVzap-MLP: 두 개의 linear layer와 GELU activation으로 구성된 MLP입니다. 중간 hidden layer의 크기는 Dh/8D_h/8입니다.
이 모델들은 각 sequence position tt에서 독립적으로 작동하며, 단일 또는 두 번의 matrix multiplication으로 점수를 계산하기 때문에 계산적으로 매우 효율적입니다.

  • Pruning 정책 (Thresholding):
KVzap은 예측된 score가 고정된 threshold τ\tau보다 낮은 KV pair들을 discard합니다. 이는 KVzip의 고정된 compression budget(예: 50% 유지)과 달리, 입력의 정보 밀도에 따라 압축률이 동적으로 조절되는 adaptive compression을 가능하게 합니다. 즉, 복잡한 입력에는 더 많은 토큰을 유지하고, 불필요한 입력에는 더 적은 토큰을 유지합니다.

  • Sliding Window:
모든 KV pair를 무작정 discard하는 것을 방지하고 local context를 보존하기 위해, KVzap은 가장 최근 w=128w=128개 토큰에 해당하는 KV pair는 점수에 관계없이 항상 유지하는 sliding window 정책을 적용합니다. 이는 StreamingLLM(Xiao et al., 2023)에서 영감을 받았습니다.

  • Algorithm (Pseudocode):
pythondef compress(hidden_states, keys, values, kvzap_model, threshold, window=128): scores = kvzap_model(hidden_states) scores[..., -window:] = float("inf") # 최근 w개 토큰은 무한대 점수를 부여하여 항상 유지 indices = torch.where(scores >= threshold) # threshold 이상인 KV pair의 인덱스 추출 return keys[indices], values[indices] # 해당 KV pair만 반환

3. 실험 및 결과

KVzap은 Qwen3-8B, Llama-3.1-8B-Instruct, Qwen3-32B 모델을 대상으로 RULER(long-context), LongBench(long-context), AIME25(reasoning) 벤치마크에서 평가되었습니다.

  • KVzap 학습 결과:
KVzip+ score와 KVzap 예측값 간의 평균 Squared Pearson correlation(R2R^2)은 KVzap-Linear에서 0.60–0.74, KVzap-MLP에서 0.66–0.77 범위를 보이며, KVzip+ score가 hidden states로부터 근사 가능함을 입증했습니다. KVzap-MLP가 KVzap-Linear보다 일관되게 높은 R2R^2를 보였습니다.

  • 성능:
    • RULER 4k: KVzap은 Qwen3-8B 및 Llama-3.1-8B-Instruct에서 3–4배 압축률까지 거의 완벽한 정확도를 유지하며 SOTA 성능을 달성했습니다. KVzip+가 KVzip과 동등하거나 더 나은 성능을 보여 KVzip+의 유효성을 검증했습니다.
    • LongBench: KVzap 모델들은 2–3배 압축률까지 거의 완벽한 정확도를 유지했습니다. RULER보다 낮은 압축률을 보였는데, 이는 RULER가 합성적이고 반복적인 데이터인 반면 LongBench는 실제 데이터로 정보 밀도가 높기 때문으로 분석됩니다.
    • AIME25 (Reasoning): KVzap-MLP는 2배를 초과하는 압축률에서도 견고한 추론 성능을 유지했습니다.
    • Adaptive Compression: KVzap의 thresholding 방식은 벤치마크마다 다른 최적 압축률을 자동으로 달성하며, 평균 2.7–3.5배의 KV 캐시 압축을 이루면서도 정확도 저하가 미미했습니다.
  • 오버헤드:
KVzap-MLP의 compute cost는 1.1%, KVzap-Linear는 0.02%에 불과하여, 무시할 만한 수준의 오버헤드를 가집니다. 메모리 오버헤드 또한 미미하며, long-context 시나리오에서는 KVzap의 추가 FLOPs가 GPU의 idle cycles을 효율적으로 활용하여 전반적인 효율성을 높입니다.

4. 결론

KVzap은 경량 surrogate 모델을 통해 KVzip+의 score를 예측하고, 이를 바탕으로 KV 캐시를 pruning하는 혁신적인 방법론입니다. 이러한 접근 방식은 KVzip의 단점을 해결하여 빠른 속도와 decoding 단계에서의 사용을 가능하게 합니다. Qwen3 및 Llama 모델에 대한 광범위한 실험을 통해, KVzap은 2-4배의 KV 캐시 압축을 달성하면서도 정확도 손실이 거의 없음을 입증했습니다. 특히, input-adaptive thresholding과 sliding window 정책은 실제 사용 환경에서 높은 유연성과 견고성을 제공합니다. KVzap은 학술적 연구와 실제 추론 엔진 통합 사이의 간극을 메울 수 있는 유력한 생산 환경 배포 후보로 평가됩니다.