NACL: A General and Effective KV Cache Eviction Framework for LLMs at Inference Time

TL;DR

NaCl addresses the costly memory footprint of KV caches in LLM inference with long contexts by introducing an encoding-phase, one-shot eviction framework that couples Proxy-Tokens Eviction with Random Eviction. By leveraging global attention statistics from a small set of proxy tokens and adding diversified randomness, NaCl reduces KV-cache usage up to 50% while preserving more than 95% of performance across short- and long-text tasks. The method achieves these gains with a simple, training-free design compatible with FlashAttention-2, and demonstrates strong robustness through extensive ablations on LLaMA2-7B variants and LongBench tasks. This work enables more efficient, scalable long-context inference and provides a practical path toward memory-efficient deployment of large language models.

Abstract

Large Language Models (LLMs) have ignited an innovative surge of AI applications, marking a new era of exciting possibilities equipped with extended context windows. However, hosting these models is cost-prohibitive mainly due to the extensive memory consumption of KV Cache involving long-context modeling. Despite several works proposing to evict unnecessary tokens from the KV Cache, most of them rely on the biased local statistics of accumulated attention scores and report performance using unconvincing metric like perplexity on inadequate short-text evaluation. In this paper, we propose NACL, a general framework for long-context KV cache eviction that achieves more optimal and efficient eviction in a single operation during the encoding phase. Due to NACL's efficiency, we combine more accurate attention score statistics in PROXY TOKENS EVICTION with the diversified random eviction strategy of RANDOM EVICTION, aiming to alleviate the issue of attention bias and enhance the robustness in maintaining pivotal tokens for long-context modeling tasks. Notably, our method significantly improves the performance on short- and long-text tasks by 80% and 76% respectively, reducing KV Cache by up to 50% with over 95% performance maintenance. The code is available at https://github.com/PaddlePaddle/Research/tree/master/NLP/ACL2024-NACL.

Figures (11)

• Figure 1: Traditional eviction algorithms perform step-by-step greedy search for tokens for eviction. Our framework searches globally for tokens within a chunk and then performs one single eviction.
• Figure 2: Attention score bias in eviction policy. The darker color in Fig. (a,b,c) shows the retained tokens.
• Figure 3: NaCl consists of a hybrid eviction policy by incorporating Random Eviction into Proxy-Tokens Eviction. Proxy-Tokens Eviction utilizes proxy tokens for more accurate eviction, while Random Eviction performs head-wise sampling from the scoring function of Proxy-Tokens Eviction to enhance the robustness.
• Figure 4: The memory usage of KV Cache with respect to the sequence length in the setting of comparable downstreaming performance between NaCl and other methods.
• Figure 5: The average accuracy is reported with different KV Cache budget (Left), Proxy tokens budget (Middle), and Random budget ratio (Right).