추론 모델을 위한 가치 인지 확률적 KV 캐시 축출

초록

추론 모델은 긴 사고 체인을 통해 정확도를 향상시키지만, 긴 출력은 메모리 및 연산 병목 현상을 초래한다. KV 캐시 제거(eviction) 방법은 캐시에서 중요하지 않은 키-값 쌍을 제거함으로써 이러한 비용을 줄이지만, 전체 KV 캐시를 유지하는 선택 기반 희소 어텐션 대안보다 종종 더 낮은 정확도를 보인다. 본 연구는 KV 캐시 제거 정확도에 중요한 핵심 요소를 식별한다. 첫째, 소수 값 상태의 크기가 비정상적으로 큰 경우가 있으며, 이를 제거하면 모델이 반복적 추론 루프에 빠지는 치명적 실패가 발생한다. 둘째, 제거 과정에 확률성을 도입하면 캐시 다양성을 증가시켜 정확도를 향상시킨다. 이러한 발견을 바탕으로, 본 연구는 큰 크기의 값 상태를 보호하고 다양한 제거 결정을 촉진하는 학습 없는 방법인 값 인식 확률적 KV 캐시 제거(VaSE)를 제안한다. 여섯 가지 추론 과제에서, 동일한 희소성 수준에서 최첨단 선택 방법보다 4배 KV 캐시 압축을 사용한 Qwen3 모델이 VaSE를 통해 더 높은 평균 정확도를 달성하였으며, 가장 강력한 제거 방법보다 4% 이상 뛰어난 성능을 보였다. 전반적으로 VaSE는 효율성과 정확도 간의 격차를 해소하며, FlashAttention2를 지원하고 추론 모델에 대한 고정 메모리 사용량을 가능하게 한다.

English

Reasoning models improve accuracy through extended chains of thought, but their long outputs create a memory and compute bottleneck. KV cache eviction methods reduce this cost by evicting unimportant key-value pairs from the cache, yet they often yield worse accuracy than selection-based sparse attention alternatives, which keep the full KV cache. We identify key factors crucial to KV cache eviction accuracy. First, a small fraction of value states have abnormally large magnitudes, and evicting them causes catastrophic failure where models enter repetitive reasoning loops. Second, introducing stochasticity during eviction improves accuracy by increasing cache diversity. Based on these findings, we propose Value-aware Stochastic KV Cache Eviction (VaSE), a training-free recipe that protects large-magnitude value states and promotes diverse eviction decisions. Across six reasoning tasks, Qwen3 models using VaSE with 4x KV cache compression yield higher average accuracies than SOTA selection method at the same sparsity, while outperforming the strongest eviction method by more than 4%. Overall, VaSE bridges the gap between efficiency and accuracy, supporting FlashAttention2 and enabling a static memory footprint for reasoning models.