NACL : A General and Effective KV Cache Eviction Framework for LLM at Inference Time (2024)

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 5\times with over 95% performance maintenance. Code available at https://github.com/PaddlePaddle/Research/tree/master/NLP/ACL2024-NACL.

## introduction large language models (llms) with longer context window @xcite @xcite @xcite have emerged recently for better conducting long conversations, summarizing long documents, or debugging code at the repository level @xcite . however, their deployment is costly and infeasible on fixed memory hardware, mainly due to the surprisingly large memory consumption of kv cache mechanism. for instance, a 7 billionparameter model with an input batch size of 4 and a sequence length of 32k results in 64gb of kv cache, 4.7× larger than the model weights. to mitigate the pressure on the scarce gpu memory from using kv cache, a number of studies @xcite @xcite have explored sparsity among transformer attention blocks to evict unnecessary tokens from the kv cache. for instance, h2o @xcite utilized the local statistics of accumulated attention scores to retain a balance of recent and heavy hitter tokens during generation. window attention based methods @xcite proposed to keep the initial tokens which is proven vital for generation fluency. this line of work reduced the memory footprint of kv cache for efficient inference with negligible loss in generation quality. in addition, the above methods do not require costly retraining which is more suitable for current open-sourced llms @xcite , compared to those that need specific attention mechanism adaptation @xcite @xcite . however, we argue that the performance reported in the above methods is over-optimistic, as the evaluation metric and tested benchmark is not sufficient. llms may fail in real-life long-context modeling tasks @xcite , though they can achieve low language modeling perplexity which is untrustworthy used as the golden metric in current studies @xcite . furthermore, the local statistics of accumulated attention score is observed to be biased (see fig. 2 ), especially in long context input, meaning that it should be carefully used as the only strategy for measuring the importance of tokens. to fill the gap, we propose nacl, a general and effective kv cache eviction framework to unleash the power of llms for long-context modeling with limited memory budgets. nacl specifically formulates the eviction task in encoding phase which is different from the commonly used one-tokenin one-token out eviction procedure in generation phase. in encoding phase, the eviction can be effectively implemented to apply only once on the whole input by progressively evicting kv caches layer by layer. the one-eviction formulation benefits current eviction policies in a more efficient and optimal way, as multiple costly eviction operations can be combined, then the global statistics of attention scores can be utilized. based on the above formulation, we present proxy-tokens eviction which exploits the global statistics of attention scores gathered from proxy tokens for eviction. in practise, the proxy tokens can be selected from the question input, commonly located at the end of a long text. intuitively, these proxy tokens are more capable of retaining the task-specific tokens in kv cache. as a result, proxy-tokens eviction alleviates the attention bias problem (see sec. 4) occurred in methods using local statistics @xcite or task-irrelevant proxy tokens @xcite . however, proxy-tokens eviction also relies heavily on the statistic of attention scores which may be untrustworthy in long-context input. thus, we incorporate random eviction, a random eviction policy, into proxy-tokens eviction. random eviction randomly samples tokens to evict from the probability distribution in proxy-tokens eviction with different seed on attention heads and layers. this diversified randomness enhances the model's robustness to maintain potentially important tokens in long text generation. we conducted extensive experiments on a single nvidia a100 (80gb) gpu on representative open-sourced llms: llama2-base, llama2-chat @xcite , and evaluated them on both short-and long-text modeling tasks from lmeval-harness @xcite and longbench @xcite . the experiments show that nacl performs kv cache eviction efficiently with negligible degradation on model quality (i.e., saving the inference memory usage of kv cache by up to 5× with over 95% maintenance). specifically, nacl achieve 80% and 75% performance improvement on short-and long-text modeling tasks, respectively, with 50% kv cache reduction, compared to current eviction methods. ## related work efficient inference with limited kv cache budgets emerged for reducing the prominent inference bottleneck caused by kv cache, particularly for long content input. a series of methods @xcite @xcite explored the sparsity among transformer's attention block, then evicted unnecessary tokens from kv cache for efficient inference. for instance, h2o @xcite retained a balance of recent and heavy hitter tokens with the highest accumulated attention scores throughout the sequence. scissorhands @xcite sequentially predicted the potentially pivotal tokens with the attention score above average within a history window. some method @xcite further applied costly eviction policy selection for better performance. however, the above methods relied heavily on the attention score with local statistics which may be sub-optimal in long-context tasks @xcite . meanwhile, some efforts have been made to utilize a learnable mechanism to determine necessary tokens during inference @xcite , or converting the traditional multi-head at-tention(mha) @xcite to multi-query attention (mqa) @xcite or group-query attention (gqa) @xcite . however, these methods involve additional training, while nacl focuses on the inference phase without resource-intensive training. efficient transformers @xcite have been extensively explored @xcite @xcite @xcite @xcite to address the self-attention operation which scales quadratically with the sequence length. for instance, sparse transformer (child et al., 2019) uses a dilated sliding window the reduces the attention complexity. longformer @xcite and bigbird (zaheer et al., 2020) reduced the complexity of selfattention by combining random, window and global attention. recurrence transformers @xcite maintain a memory bank of past kv cache to process the long text in segments. however, the above methods either trade off model quality or require re-training of models, but often failed in achieving memory saving and wall-clock speedup at inference time @xcite . length extrapolation enabled language models to generalize beyond the context window they were trained on. a recent line of research @xcite @xcite focuses on adapting relative positional embedding @xcite widely used in current foundation models @xcite for context window extension. attention sink @xcite and lm-infinite @xcite further exploited the initial tokens to recover the performance of window attention for infinite-length inputs. however, the ability of these methods tested using metric like perplexity is over-optimistic for long context tasks @xcite . ## problem formulation this section defines a two-phased approach for efficient kv cache management during llm inference, tailored for scenarios with limited kv cache budgets. eviction policy we defined the eviction policy f score : s i t ← s i t-1 , subject to |s i t | = |s i t-1 | ≤ c where the scoring function f score assigns low scores to unnecessary tokens for eviction, such that the pre-define kv cache budget c is maintained. s i t denote the indices set of retained tokens in kv cache at t-th time step and i-th transformer layer. encoding phase eviction the model processes the input prompts, x i @xmath0, . . . , x i p ] ∈ r p×d , to compute the initial key cache k i 0 = x i prompt w i @xmath1, where p denotes the encoding prompt length, w i k , w i @xmath2, respectively. the attention scores a i @xmath3 ## observation we present two experimental findings by rethinking previous eviction methods that inspire the design of nacl. ## nacl in this section, we present a hybrid kv cache eviction policy in nacl, including the proxy-tokens eviction in sec. 5.1 and random eviction in sec. 5.2. ## conclusion in this paper, we focus on the accuracy, robustness, and reliability of evaluation for kv cache eviction algorithms deployed in llms for processing long texts. we introduce nacl, a novel approach that combines proxy- tokens eviction and ran-dom eviction for kv cache eviction strategies, significantly reducing memory usage during model inference without the need for training. we model the eviction problem as a combinatorial optimization issue, where proxy-tokens eviction provides eviction references based on importance, and random eviction enhances information richness and robustness through headwise and layerwise composite sampling. through extensive evaluation, we demonstrate that nacl can significantly improve cache eviction strategies, reduce inference memory costs, and minimize the impact on the llm's ability to handle complex tasks.