Layered eviction method for KV Cache
By employing a KV Cache hierarchical eviction method and a differentiated sparsity strategy based on the hierarchical characteristics of the Transformer model, the problem of low utilization efficiency of LLM on edge devices is solved, achieving a synergistic improvement in inference speed and accuracy, and is suitable for various task scenarios.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- HANGZHOU INTERNATIONAL INNOVATION INSTITUTE OF BEIHANG UNIVERSITY
- Filing Date
- 2025-11-27
- Publication Date
- 2026-04-10
AI Technical Summary
Existing technologies struggle to efficiently utilize video memory resources on edge devices to achieve coordinated optimization of LLM inference speed and accuracy, especially in adapting to different task scenarios.
A hierarchical eviction method using KV Cache is adopted. Based on the hierarchical characteristics of the Transformer model, differentiated sparsity strategies are applied to the initial layer, intermediate layer and tail layer respectively. The initial layer retains all KV Cache, the intermediate layer dynamically selects tokens whose attention scores reach the threshold, and the tail layer performs cumulative attention compensation. The eviction strategy is flexibly adapted.
It achieves efficient utilization of video memory, improves the synergistic balance between inference speed and accuracy, adapts to various task scenarios, and resolves the contradiction between memory optimization and inference accuracy in existing technologies.
Smart Images
Figure CN121835772A_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of artificial intelligence large model inference optimization technology, and in particular to a KV Cache hierarchical eviction method. Background Technology
[0002] With the rapid development of artificial intelligence technology, large-scale language models (LLMs) such as GPT and Deepseek have been widely used in many fields such as natural language generation, dialogue interaction, and document processing. However, the continuous expansion of model parameter size and the increasing length of inference context have led to severe memory bottlenecks in edge deployment and efficient inference. Limited memory resources cannot support massive amounts of KV-Cache data, becoming a core obstacle restricting the deployment of LLMs on edge devices and improving inference speed.
[0003] To overcome this bottleneck, KV-Cache optimization technology has become a hot research topic in the industry. It mainly improves the efficiency of video memory utilization through three core approaches: compression, quantization, and sparse loading. Compression technology takes advantage of the natural sparsity of KV-Cache to reduce the cache size by evicting KV tokens with low attention scores. Quantization technology converts KV-Cache data into a low-precision format to reduce memory usage from the storage perspective. Sparse loading technology uses the sparsity characteristics to perform dynamic pruning, loading only a portion of the key KV-Cache data to reduce invalid computation and accelerate the inference process.
[0004] Despite some progress in existing technologies, significant shortcomings remain, making it difficult to achieve a synergistic balance between memory optimization and inference accuracy: Key-Value Cache (KV) hierarchical compression technologies: PyramidKV employs a fixed-budget, dynamic inter-layer allocation strategy, designing the KV-Cache size in a pyramid shape. However, its attention pattern is fixed, making it unable to flexibly adapt to different task scenarios. DynamicKV, while optimizing budget allocation logic based on task awareness, relies excessively on task type identification, and its progressive cache update mechanism is complex. Both are constrained by fixed budgets, often forcing the discard of tokens critical to inference accuracy in practical applications, resulting in performance degradation due to prioritizing memory over precision.
[0005] KV Cache dynamic loading technology: Quest dynamically loads a fixed number of KV-Caches for attention calculation by defining a similarity metric. In essence, it does not evict redundant caches. Although it can ensure accuracy, the memory consumption remains high, resulting in a dilemma of prioritizing accuracy over memory.
[0006] Small model assisting large model technology: IAM utilizes the high similarity of attention scores of large models of different magnitudes in the same series to assist the inference of large models through small models. However, its core focus is on deployment optimization in cluster server scenarios, and it has poor adaptability to resource constraints of end devices, making it difficult to promote and apply.
[0007] Top-p kernel sampling strategy: The Top-p kernel sampling method is introduced into sparse attention to reduce accuracy loss. However, this approach adopts a uniform strategy for all layers, does not consider the differences in attention distribution between different layers of LLM, is not friendly enough to the initial layer where attention is scattered, and does not design an effective information compensation mechanism for discarded tokens. Therefore, there is still a risk of accuracy loss that cannot be ignored.
[0008] Furthermore, LLM's attention distribution exhibits a significant hierarchical characteristic: the initial layer's attention is dispersed, requiring sufficient key-value cache to ensure the stability of inference initialization; the middle layers need to retain a large amount of key-value cache in certain tasks to maintain task relevance; and the tail layer's attention scores are highly concentrated, with a relatively limited impact on the model's overall inference performance. Existing technologies have failed to fully explore and utilize this key characteristic, making it difficult to reconcile the contradiction between memory optimization and inference accuracy. Summary of the Invention
[0009] The purpose of this invention is to propose a KV Cache hierarchical eviction method that fully utilizes the attention distribution characteristics of different layers of LLM and achieves a synergistic unity of efficient memory utilization, improved inference speed and guaranteed inference accuracy through a differentiated sparsity strategy. It is adaptable to deployment in multiple scenarios and solves the core pain points of existing technologies.
[0010] To achieve the above objectives, this invention proposes a hierarchical eviction method for KV Cache, the specific steps of which are as follows: Step S1: Divide the total number of layers of the Transformer self-attention mechanism neural network model into initial layer, intermediate layer and tail layer according to the set parameters; Step S2: Retain all the KV caches of the Transformer model for the initial layer KV cache strategy; Step S3: Use kernel sampling to dynamically select tokens for the KV Cache strategy of the intermediate layer, so that the total attention score of the selected tokens reaches the set threshold. Step S4: For the KV Cache strategy of the tail layer, calculate the cumulative attention for each token remaining after removing the attention convergence token and the sliding window token. For each autoregressive decoding, add the attention of each token to the corresponding cumulative attention, retain the k tokens with the highest cumulative attention scores, and compensate the discarded tokens for their attention scores.
[0011] Preferably, in step S1, the initial layer, intermediate layer, and tail layer satisfy the following: ; in, This represents the total number of layers in the Transformer model. This is the initial number of layers. The number of intermediate layers. This refers to the number of layers at the tail end.
[0012] Preferably, in step S3, the specific steps are as follows: Step S31: Calculate the normalized probability weights and input them into the pruning unit. The formula is as follows: ; in, These are the normalized probability weights. For query vector, For key vectors, The dimension of the key vector. This is a transpose operation; Step S32: The pruning unit, in descending order of probability weight, updates the index of the corresponding token. i Add to the index set IDX until the cumulative probability is greater than or equal to a set threshold. p ; Step S33: Calculate the attention output of the intermediate layer, using the following formula: ; in, The attention output value for the intermediate layer. The diagonal matrix formed by the set of indices IDX selected for the pruned units. It is a value vector.
[0013] Preferably, in step S4, the specific steps are as follows: Step S41: During the decoding process, calculate the attention score using the following formula: ; in, Attention score; Step S42: After removing the tokens that retain the attention sinks and the tokens within the sliding window, add the attention score of each remaining token to the cumulative attention. The formula is as follows: ; in, To accumulate attention sets, For the first i The cumulative attention of each token; Step S43: Based on cumulative attention Calculate the k tokens with the highest retention scores, add the indices of the remaining tokens to the eviction set E, and evict the corresponding KV Cache during the next round of autoregressive decoding; Step S44: Calculate the mean attention score m of the evicted tokens as the attention score compensation for the tokens corresponding to each index in the evicted set E; Step S45: Calculate the attention output of the tail layer, using the following formula: ; in, This is the attention output for the tail layer.
[0014] Preferably, in step S45, for tokens that have been evicted, the attention output is calculated using the mean attention score m instead of the attention score.
[0015] Therefore, this invention proposes a hierarchical eviction method for KV Cache, which has the following advantages: (1) This invention fits the characteristics of LLM layered attention, with full caching in the initial layer, no budget constraints in the middle layer, and mean compensation in the tail layer. This avoids the loss of key tokens and greatly reduces the memory usage, thus solving the dilemma of balancing memory and precision or precision and memory.
[0016] (2) This invention is compatible with different scales of LLM and end-to-end / server-side deployment scenarios through parameterized layered configuration, and can adapt to various tasks such as text generation and long context processing, with high flexibility.
[0017] (3) The dynamic loading of the middle layer Top-p in this invention reduces invalid calculations, and the KV-Cache of the tail layer reduces data transmission overhead. While ensuring accuracy, it accelerates autoregressive decoding and improves inference response speed.
[0018] The technical solution of the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. Attached Figure Description
[0019] Figure 1 This is a flowchart of a KV Cache hierarchical eviction method according to the present invention. Detailed Implementation
[0020] To make the technical solutions, advantages, and objectives of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below. The described embodiments are only some, not all, of the embodiments of the present invention. All other embodiments obtained by those skilled in the art based on the described embodiments of the present invention without creative effort are within the protection scope of the present invention.
[0021] Unless otherwise defined, the technical or scientific terms used in this invention shall have the ordinary meaning as understood by one of ordinary skill in the art to which this invention pertains.
[0022] Example like Figure 1 As shown, this invention provides a hierarchical eviction method for KV Cache, the steps of which are as follows: Step S1: Divide the total number of layers of the Transformer decoder neural network model with self-attention mechanism into initial layers, intermediate layers, and tail layers according to the set parameters, satisfying: ; in, This represents the total number of layers in the Transformer model. This is the initial number of layers. The number of intermediate layers. This refers to the number of layers at the tail end.
[0023] Step S2: For the initial layer's KV Cache strategy, retain all the KV Caches of the Transformer model; Step S3: For the KV Cache strategy of the intermediate layer, the kernel sampling method Top_p is used to dynamically select tokens so that the total attention score of the selected tokens reaches the set threshold. The specific steps are as follows: Step S31: Calculate the normalized probability weights and input them into the pruning unit. The formula is as follows: ; in, These are the normalized probability weights. For query vector, For key vectors, The dimension of the key vector. This is a transpose operation; Step S32: The pruning unit, in descending order of probability weight, updates the index of the corresponding token. i Add to the index set IDX until the cumulative probability is greater than or equal to a set threshold. p ; Step S33: Calculate the attention output of the intermediate layer, using the following formula: ; in, The attention output value for the intermediate layer. The diagonal matrix formed by the set of indices IDX selected for the pruned units. It is a value vector.
[0024] Step S4: For the KV Cache strategy at the tail layer, calculate the cumulative attention for each token remaining after removing the attention convergence token and the sliding window token. Each time autoregressive decoding is performed, the attention of each token is added to the corresponding cumulative attention. The k tokens with the highest cumulative attention scores are retained, and attention score compensation is applied to discarded tokens. The specific steps are as follows: Step S41: During the decoding process, calculate the attention score using the following formula: ; in, Attention score; Step S42: After removing the tokens that retain the attention sinks and the tokens within the sliding window, add the attention score of each remaining token to the cumulative attention. The formula is as follows: ; in, To accumulate attention sets, For the first i The cumulative attention of each token; Step S43: Based on cumulative attention Calculate the k tokens with the highest retention scores, add the indices of the remaining tokens to the eviction set E, and evict the corresponding KV Cache during the next round of autoregressive decoding; Step S44: Calculate the mean attention score m of the evicted tokens as the attention score compensation for the tokens corresponding to each index in the evicted set E; Step S45: Calculate the attention output of the tail layer, using the following formula: ; in, This is the attention output for the tail layer.
[0025] For tokens that have been expelled, the attention output is calculated using the mean attention score m instead of the attention score.
[0026] The invention will be further illustrated below through specific implementation examples.
[0027] Taking Llama 3.1-70B as an example, the initial layer I=8, the intermediate layer M=32, and the tail layer L=40 are set based on this model. The model is then directly loaded using the transformer library, and the relevant code for attention calculation and KV cache management is replaced before deployment.
[0028] First, the user inputs text, which is then transmitted to the backend for word segmentation and other operations, transforming it into a user input sequence embedding. The model obtains the input sequence embedding and enters the pre-filling stage, where it calculates the query vector Q, key vector K, and value vector V corresponding to the sequence and performs basic multi-head attention (MHA) parallel computation to generate the first token.
[0029] Subsequently, the autoregressive serial decoding stage begins. For each new token generated by autoregression, the model calculates its corresponding Q, K, and V and updates the KV Cache (adding the new K and V vectors to the KV Cache). Then, depending on the layer in which the Transformer is located, different attention calculation operations are performed as follows: For the initial layer I=8, basic multi-head attention (MHA) computation is used, and all KV caches are retained; the output of the current layer is processed by the feedforward neural network before entering the next layer.
[0030] For the intermediate layer M=32, in each autoregressive attention calculation, the normalized attention score of each K token and the current Q token is first calculated, as follows: ; in, These are the normalized probability weights. For query vector, For key vectors, For vector dimensions, This is a transpose operation; The score corresponding to each K token is denoted as: S i , i The index corresponds to the K token, and then the process proceeds to the pruning unit.
[0031] In the pruning unit, the indices of the corresponding K tokens are ordered from high to low probability (i.e., normalized attention scores). i Add to the index set IDX until the cumulative probability is greater than or equal to the threshold: ; in, , =0.9.
[0032] Finally, the attention output of the intermediate layer is calculated using the following formula: ; ; in, The attention output value for the intermediate layer. The diagonal matrix formed by the set of indices IDX selected for the pruned units.
[0033] intermediate layer output After undergoing calculations such as those by a feedforward neural network, it proceeds to the next layer.
[0034] For the tail layer L=40, in each autoregressive attention calculation, the attention score between each K token and the current Q token is first calculated, as follows: ; in, Attention score; After removing the tokens that retain the attention sinks and sliding window, for each position i The scores of the K tokens are accumulated to Then, based on the cumulative attention set C, the top_k of the intermediate tokens are retained, and the cumulative attention score is retained. The top k tokens are selected, and the remaining tokens are evicted. Their indices are added to the evictation set E, and the corresponding KV Cache in E is evicted when the KV Cache is updated for the next autoregressive decoding. The cumulative attention set C is represented as follows: ; in, For the first i The cumulative attention of each token; For each expelled token, its mean attention score is calculated and stored, denoted as m, as attention score compensation for the tokens corresponding to each index in the expulsion set E.
[0035] Finally, the attention output of the tail layer is calculated using the following formula: ; in, For the attention output of the tail layer; For tokens that have been evicted, replace their corresponding scores with the value of m in the tail layer output. After computation by the feedforward neural network, it enters the next layer. After passing through all layers, a new token is generated and added to the end of the sequence embedding to start the next round of autoregressive inference.
[0036] After several rounds of autoregressive inference, the model generates a terminator or reaches the maximum generation length, the inference ends, and the backend inference engine returns the model's output to the user.
[0037] It is worth noting that all contents not described in detail in this invention are existing technologies and are well known to those skilled in the art.
[0038] Therefore, this invention provides a KV Cache hierarchical eviction method, which achieves the coordinated unity of efficient memory utilization, improved inference speed and guaranteed inference accuracy in LLM inference through parameterized hierarchical configuration and differentiated sparsity strategy. It is also adaptable to multiple models and scenarios, and solves the core pain points of existing technologies.
[0039] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and not to limit them. Although the present invention has been described in detail with reference to preferred embodiments, those skilled in the art should understand that modifications or equivalent substitutions can still be made to the technical solutions of the present invention, and these modifications or equivalent substitutions cannot cause the modified technical solutions to deviate from the spirit and scope of the technical solutions of the present invention.
Claims
1. A hierarchical eviction method for KV cache, characterized in that, The specific steps are as follows: Step S1: Divide the total number of layers of the Transformer decoder, a self-attention mechanism neural network model, into initial layers, intermediate layers, and tail layers according to the set parameters; Step S2: Retain all the KV caches of the Transformer model for the initial layer KV cache strategy; Step S3: Use kernel sampling to dynamically select tokens for the KV Cache strategy of the intermediate layer, so that the total attention score of the selected tokens reaches the set threshold. Step S4: For the KV Cache strategy of the tail layer, calculate the cumulative attention for each token remaining after removing the attention convergence token and the sliding window token. For each autoregressive decoding, add the attention of each token to the corresponding cumulative attention, retain the k tokens with the highest cumulative attention scores, and compensate the discarded tokens for their attention scores.
2. The KV Cache hierarchical eviction method according to claim 1, characterized in that, In step S1, the initial layer, intermediate layer, and tail layer satisfy the following: ； in, This represents the total number of layers in the Transformer model. This is the initial number of layers. The number of intermediate layers. This refers to the number of layers at the tail end.
3. The KV Cache hierarchical eviction method according to claim 1, characterized in that, In step S3, the specific steps are as follows: Step S31: Calculate the normalized probability weights and input them into the pruning unit. The formula is as follows: ； in, These are the normalized probability weights. For query vector, For key vectors, The dimension of the key vector. This is a transpose operation; Step S32: The pruning unit, in descending order of probability weight, updates the index of the corresponding token. i Add to the index set IDX until the cumulative probability is greater than or equal to a set threshold. p ; Step S33: Calculate the attention output of the intermediate layer, using the following formula: ； in, The attention output value for the intermediate layer. The diagonal matrix formed by the set of indices IDX selected for the pruned units. It is a value vector.
4. The KV Cache hierarchical eviction method according to claim 1, characterized in that, In step S4, the specific steps are as follows: Step S41: During the decoding process, calculate the attention score using the following formula: ； in, Attention score; Step S42: After removing the tokens that retain the attention sinks and the tokens within the sliding window, add the attention score of each remaining token to the cumulative attention. The formula is as follows: ； in, To accumulate attention sets, For the first i The cumulative attention of each token; Step S43: Based on cumulative attention Calculate the k tokens with the highest retention scores, add the indices of the remaining tokens to the eviction set E, and evict the corresponding KV Cache during the next round of autoregressive decoding; Step S44: Calculate the mean attention score m of the evicted tokens as the attention score compensation for the tokens corresponding to each index in the evicted set E; Step S45: Calculate the attention output of the tail layer, using the following formula: ； in, This is the attention output for the tail layer.
5. The KV Cache hierarchical eviction method according to claim 4, characterized in that, In step S45, for tokens that have been evicted, the attention output is calculated by replacing the attention score with the mean attention score m.