HeteroCache: 장문 맥락 LLM 추론을 위한 이종 키-값 캐시 압축을 위한 동적 검색 접근 방식

HeteroCache: A Dynamic Retrieval Approach to Heterogeneous KV Cache Compression for Long-Context LLM Inference

키-값(KV) 캐시의 선형적인 메모리 증가는 장문 맥락 작업에서 LLM 추론의 상당한 병목 현상을 야기합니다. 기존의 정적 압축 방법은 종종 전역적으로 중요한 정보를 보존하지 못하는데, 이는 토큰의 중요성이 동적으로 변화하는 '어텐션 드리프트' 현상을 간과하기 때문입니다. 최근의 동적 검색 접근 방식은 이 문제를 해결하려고 시도하지만, 일반적으로 세분화되지 않은 캐싱 전략으로 인해 빈번한 데이터 전송으로 인한 높은 입출력(I/O) 오버헤드를 발생시킵니다. 이러한 한계를 극복하기 위해, 우리는 학습이 필요 없는 동적 압축 프레임워크인 HeteroCache를 제안합니다. 우리의 방법은 두 가지 핵심적인 통찰력에 기반합니다. 첫째, 어텐션 헤드는 다양한 시간적 이질성을 보입니다. 둘째, 동일 레이어 내의 헤드 간에 상당한 공간적 중복성이 존재합니다. 이러한 통찰력을 바탕으로, HeteroCache는 헤드의 안정성과 중복성에 따라 헤드를 분류합니다. 결과적으로, 우리는 빠르게 변화하는 어텐션을 보이는 헤드에 더 큰 캐시 자원을 할당하는 세분화된 가중치 전략을 적용하여, 맥락 변화를 효과적으로 포착하고, 기존의 세분화되지 않은 전략의 비효율성을 해결합니다. 또한, 일부 대표적인 헤드가 어텐션 변화를 모니터링하고, 필요에 따라 CPU에서 맥락을 비동기적으로 검색하는 계층적 저장 메커니즘을 사용하여, I/O 지연을 효과적으로 숨깁니다. 실험 결과, HeteroCache는 여러 장문 맥락 벤치마크에서 최첨단 성능을 달성했으며, 224K 맥락에서 원래 모델보다 최대 3배 빠른 디코딩 속도를 보입니다. 저희의 코드는 오픈 소스로 공개될 예정입니다.

The linear memory growth of the KV cache poses a significant bottleneck for LLM inference in long-context tasks. Existing static compression methods often fail to preserve globally important information, principally because they overlook the attention drift phenomenon where token significance evolves dynamically. Although recent dynamic retrieval approaches attempt to address this issue, they typically suffer from coarse-grained caching strategies and incur high I/O overhead due to frequent data transfers. To overcome these limitations, we propose HeteroCache, a training-free dynamic compression framework. Our method is built on two key insights: attention heads exhibit diverse temporal heterogeneity, and there is significant spatial redundancy among heads within the same layer. Guided by these insights, HeteroCache categorizes heads based on stability and redundancy. Consequently, we apply a fine-grained weighting strategy that allocates larger cache budgets to heads with rapidly shifting attention to capture context changes, thereby addressing the inefficiency of coarse-grained strategies. Furthermore, we employ a hierarchical storage mechanism in which a subset of representative heads monitors attention shift, and trigger an asynchronous, on-demand retrieval of contexts from the CPU, effectively hiding I/O latency. Finally, experiments demonstrate that HeteroCache achieves state-of-the-art performance on multiple long-context benchmarks and accelerates decoding by up to $3\times$ compared to the original model in the 224K context. Our code will be open-source.