IceCache: 장기 시퀀스 LLM을 위한 메모리 효율적인 KV 캐시 관리

IceCache: Memory-efficient KV-cache Management for Long-Sequence LLMs

키-값(KV) 캐시는 중간 어텐션 상태를 저장하고 오토리게시브 생성 과정에서 불필요한 연산을 피함으로써, 대규모 언어 모델(LLM)의 추론 속도를 향상시키는 데 중요한 역할을 합니다. 그러나 KV 캐시의 메모리 사용량은 시퀀스 길이에 따라 선형적으로 증가하며, 이는 자원이 제한된 환경에서 심각한 메모리 병목 현상을 야기할 수 있습니다. 기존 연구에서는 KV 캐시를 CPU로 오프로드하고 GPU에는 일부만 유지하는 방법을 제안했지만, 이러한 방식은 종종 부정확한 토큰 선택에 의존하며, 연쇄적 추론과 같은 장기 생성 작업에서 성능 저하를 초래합니다. 본 논문에서는 의미론적 토큰 클러스터링과 PagedAttention을 통합한 새로운 KV 캐시 관리 전략인 IceCache를 제안합니다. IceCache는 의미적으로 관련된 토큰을 계층적이고 동적으로 업데이트 가능한 데이터 구조로 관리되는 연속적인 메모리 영역으로 구성하여, 보다 효율적인 토큰 선택과 CPU-GPU 전송 시 메모리 대역폭의 향상된 활용을 가능하게 합니다. LongBench 데이터셋에 대한 실험 결과, IceCache는 256개의 토큰 예산으로 전체 KV 캐시 모델이 달성한 원래 정확도의 99%를 유지합니다. 또한, 다른 오프로딩 기반 방법과 비교했을 때, IceCache는 KV 캐시 토큰 예산의 25%만을 사용하면서 경쟁력 있는, 또는 더 나은 지연 시간과 정확도를 달성하여, 장기 시퀀스 시나리오에서 효과적인 성능을 입증합니다. 코드 및 추가 정보는 프로젝트 웹사이트 https://yuzhenmao.github.io/IceCache/ 에서 확인할 수 있습니다.

Key-Value (KV) cache plays a crucial role in accelerating inference in large language models (LLMs) by storing intermediate attention states and avoiding redundant computation during autoregressive generation. However, its memory footprint scales linearly with sequence length, often leading to severe memory bottlenecks on resource-constrained hardware. Prior work has explored offloading KV cache to the CPU while retaining only a subset on the GPU, but these approaches often rely on imprecise token selection and suffer performance degradation in long-generation tasks such as chain-of-thought reasoning. In this paper, we propose a novel KV cache management strategy, IceCache, which integrates semantic token clustering with PagedAttention. By organizing semantically related tokens into contiguous memory regions managed by a hierarchical, dynamically updatable data structure, our method enables more efficient token selection and better utilization of memory bandwidth during CPU-GPU transfers. Experimental results on LongBench show that, with a 256-token budget, IceCache maintains 99% of the original accuracy achieved by the full KV cache model. Moreover, compared to other offloading-based methods, IceCache attains competitive or even superior latency and accuracy while using only 25% of the KV cache token budget, demonstrating its effectiveness in long-sequence scenarios. The code is available on our project website at https://yuzhenmao.github.io/IceCache/.