LookaheadKV: 미래를 예측하여 빠른 정확도 기반 KV 캐시 제거 기술 (생성 과정 없이)

LookaheadKV: Fast and Accurate KV Cache Eviction by Glimpsing into the Future without Generation

트랜스포머 기반 대규모 언어 모델(LLM)은 오토리거시브 추론 과정에서 불필요한 연산을 줄이기 위해 키-값(KV) 캐싱을 사용합니다. 이 메커니즘은 효율성을 크게 향상시키지만, 캐시 크기는 입력 시퀀스 길이에 따라 선형적으로 증가하며, 이는 긴 문맥 처리에서 병목 현상을 초래합니다. 기존의 해결 방안은 중요도 추정 점수를 기반으로 중요도가 낮은 프롬프트 KV를 제거하여 이 문제를 완화합니다. 주목할 만한 최근 연구에서는 '미래를 엿보는' 방식을 통해 제거 품질을 향상시키는데, 이는 드래프트 생성기를 사용하여 대상 모델의 실제 응답을 근사하는 가짜 미래 응답을 생성하고, 이를 통해 캐시된 KV의 중요도를 더욱 정확하게 추정하는 방식입니다. 그러나 이러한 접근 방식은 계산 비용이 많이 드는 드래프트 생성에 의존하며, 이는 상당한 사전 처리 오버헤드를 유발하고 실제 환경에서의 활용성을 제한합니다. 이러한 문제를 해결하기 위해, 우리는 명시적인 드래프트 생성이 필요 없는 경량화된 제거 프레임워크인 LookaheadKV를 제안합니다. LookaheadKV는 트랜스포머 레이어에 실제 중요도 점수를 높은 정확도로 예측하도록 학습된 파라미터 효율적인 모듈을 추가합니다. 우리의 설계는 기존의 저렴한 휴리스틱과 유사한 무시할 만한 런타임 오버헤드를 보장하면서도, 더 비용이 많이 드는 근사화 방법보다 우수한 정확도를 달성합니다. 다양한 모델을 사용하여 수행된 광범위한 실험 결과, LookaheadKV는 다양한 긴 문맥 이해 작업에서 최첨단 기술보다 뛰어난 성능을 보일 뿐만 아니라, 제거 비용을 최대 14.5배까지 줄여 토큰 생성 시간을 크게 단축시킵니다. 저희 코드와 관련 자료는 다음 링크에서 확인할 수 있습니다: https://github.com/SamsungLabs/LookaheadKV.

Transformer-based large language models (LLMs) rely on key-value (KV) caching to avoid redundant computation during autoregressive inference. While this mechanism greatly improves efficiency, the cache size grows linearly with the input sequence length, quickly becoming a bottleneck for long-context tasks. Existing solutions mitigate this problem by evicting prompt KV that are deemed unimportant, guided by estimated importance scores. Notably, a recent line of work proposes to improve eviction quality by "glimpsing into the future", in which a draft generator produces a surrogate future response approximating the target model's true response, and this surrogate is subsequently used to estimate the importance of cached KV more accurately. However, these approaches rely on computationally expensive draft generation, which introduces substantial prefilling overhead and limits their practicality in real-world deployment. To address this challenge, we propose LookaheadKV, a lightweight eviction framework that leverages the strength of surrogate future response without requiring explicit draft generation. LookaheadKV augments transformer layers with parameter-efficient modules trained to predict true importance scores with high accuracy. Our design ensures negligible runtime overhead comparable to existing inexpensive heuristics, while achieving accuracy superior to more costly approximation methods. Extensive experiments on long-context understanding benchmarks, across a wide range of models, demonstrate that our method not only outperforms recent competitive baselines in various long-context understanding tasks, but also reduces the eviction cost by up to 14.5x, leading to significantly faster time-to-first-token. Our code is available at https://github.com/SamsungLabs/LookaheadKV.