CompressKV: 의미 검색 기반 KV 캐시 압축을 통한 자원 효율적인 장문 컨텍스트 LLM 추론

CompressKV: Semantic-Retrieval-Guided KV-Cache Compression for Resource-Efficient Long-Context LLM Inference

장문 컨텍스트 대규모 언어 모델(LLM)의 추론은 키-값(KV) 캐시의 메모리 사용량과 디코딩 비용에 의해 점점 더 제약되고 있으며, 이는 자원이 제한적인 하드웨어에서의 지속 가능한 배포를 어렵게 만듭니다. 기존 KV 캐시 제거 방법은 일반적으로 GQA 기반 LLM의 모든 헤드에 대한 휴리스틱 토큰 점수를 적용합니다. 이러한 방법은 어텐션 헤드의 다양한 기능을 무시하여 중요한 토큰이 제거되고, 결과적으로 LLM의 성능이 저하됩니다. 이 문제를 해결하기 위해, 본 논문에서는 GQA 기반 LLM을 위한 자원 효율적인 KV 캐시 압축 프레임워크인 CompressKV를 제안합니다. CompressKV는 모든 헤드의 어텐션 점수를 집계하는 대신, 프롬프트의 초기 및 최종 토큰과 의미적으로 중요한 중간 컨텍스트 정보를 모두 포착하는 Semantic Retrieval Heads (SRH)를 식별하고, 이를 사용하여 KV 쌍을 유지해야 할 토큰을 선택합니다. 또한, CompressKV는 레이어별 제거 오류에 대한 오프라인 추정치를 기반으로 각 레이어에 캐시 예산을 할당합니다. LongBench 및 Needle-in-a-Haystack 데이터셋에 대한 실험 결과, CompressKV는 다양한 메모리 예산에서 기존 KV 캐시 제거 방법보다 일관되게 우수한 성능을 보였습니다. 특히, LongBench의 질의 응답 작업에서 전체 캐시 성능의 97% 이상을 유지하면서 KV 캐시의 3%만 사용하고, Needle-in-a-Haystack에서는 정확도 90%를 달성하며 KV 저장 공간의 0.7%만을 사용했습니다. 이러한 결과는 장문 컨텍스트 LLM 추론에 대한 개선된 자원-성능 균형을 보여줍니다. 저희 코드는 다음 링크에서 공개적으로 이용하실 수 있습니다: https://github.com/TUDa-HWAI/CompressKV

Long-context large language model (LLM) inference is increasingly constrained by the memory footprint and decoding cost of key-value (KV) caches, limiting sustainable deployment on resource-constrained hardware. Existing KV cache eviction methods typically apply heuristic token scoring over all heads in GQA-based LLMs. These methods ignore the different functionalities of attention heads, leading to the eviction of critical tokens and thus degrading the performance of LLMs. To address this issue, we propose CompressKV, a resource-efficient KV-cache compression framework for GQA-based LLMs. Instead of aggregating attention scores from all heads, CompressKV identifies Semantic Retrieval Heads (SRHs) that capture both the initial and final tokens of a prompt and semantically important mid-context evidence, and uses them to select tokens whose KV pairs should be retained. Furthermore, CompressKV allocates cache budgets across layers according to offline estimates of layer-wise eviction error. Experiments on LongBench and Needle-in-a-Haystack show that CompressKV consistently outperforms existing KV-cache eviction methods across memory budgets. Notably, it preserves over 97\% of full-cache performance using only 3\% of the KV cache on LongBench question-answering tasks and achieves 90\% accuracy with just 0.7\% KV storage on Needle-in-a-Haystack. These results demonstrate an improved resource--performance trade-off for long-context LLM inference. Our code is publicly available at: https://github.com/TUDa-HWAI/CompressKV