An adaptive kv cache compression method and system for large language models

Abstract

This invention relates to the field of artificial intelligence and large language model inference optimization technology, and discloses an adaptive key-value cache compression method and system for large language models. The method includes: constructing a lexical importance measurement mechanism; analyzing the attention head distribution structure in large language model inference and constructing multiple pruning strategies; designing an adaptive key-value cache compression hybrid strategy set based on the lexical importance measurement mechanism and the attention head distribution structure, and based on the pruning strategies; constructing a static adaptive key-value cache compression method to automatically allocate key-value cache compression strategies during the pre-filling stage of large language model inference; and adaptively compressing the key-value cache based on the allocated key-value cache compression strategies during the decoding stage. This invention can achieve efficient compression of the key-value cache without relying on explicit attention score calculation, and is compatible with system-level inference optimization frameworks, maintaining generation performance while significantly reducing memory consumption and enhancing long-context inference capabilities.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence and large language model inference optimization technology, specifically to an adaptive key-value caching compression method and system for large language models. Background Technology

[0002] In recent years, with the emergence and development of the Transformer architecture, the performance of large language models has continuously improved. Currently, the computational complexity and memory overhead of these models have become the main bottlenecks restricting their deployment and performance improvement, placing higher demands on the inference capabilities of large language models. Key-value caching compression, as one of the important inference optimization methods for large models, has extremely high research and practical engineering application value. The essence of key-value caching is that during large model inference, by retaining previously calculated key-value vector pairs, there is no need to perform expensive repeated calculations on these vectors, thereby reducing the time complexity required to generate each token to linear, greatly improving inference efficiency.

[0003] However, as the model size increases and the generated length extends, the key-value cache size increases linearly with the length of the input text sequence, leading to a surge in memory pressure and limiting the model's performance under longer input sequences. Research shows that when the input sequence length exceeds the pre-training text length, the text generation capability of LLMs drops sharply. The need to extend the input context length of LLMs contradicts the high model deployment overhead and severe performance loss caused by long context windows; how to balance key-value cache size and model performance remains an unsolved problem. In addition, existing key-value cache compression methods have many shortcomings: current key-value cache pruning methods based on elimination strategies mostly rely on attention scores to determine word importance, but are incompatible with mainstream optimization methods such as Flash Attention, hindering the application of this type of method in mainstream LLMs inference acceleration frameworks; previous methods mostly adopted globally uniform static compression strategies, without considering the different attention head distribution structures between attention heads, which easily leads to high model performance loss when eliminating words; previous methods only extended the input sequence context length, failing to maintain the model's memory capacity, resulting in the loss of some memory in the early stages of the dialogue.

[0004] Therefore, it is urgent to solve the technical problems of poor compatibility, high model performance loss, and low long context processing capability in existing key-value caching compression methods. Summary of the Invention

[0005] This invention provides an adaptive key-value caching compression method and system for large language models, addressing the technical problems of poor compatibility, high performance loss in large language models, and low long context processing capability of key-value caching compression methods. The specific technical solution is as follows:

[0006] Firstly, an adaptive key-value caching compression method for large language models is provided, the method comprising:

[0007] S1: Based on the numerical distribution characteristics of key vectors and attention scores in large language models, a lexical importance measurement mechanism with the L2 norm of key vectors as the core is constructed.

[0008] S2: Analyze the attention head distribution structure in the reasoning of the large language model to obtain the attention head distribution pattern;

[0009] S3: Based on the lexical importance measurement mechanism and the attention head distribution pattern, construct a variety of pruning strategies, and design an adaptive key-value caching compression hybrid strategy set based on the pruning strategies;

[0010] S4: Construct a static adaptive key-value cache compression method based on an adaptive key-value cache compression hybrid strategy set. This method is used to automatically allocate key-value cache compression strategies during the pre-filling stage of large language model inference. During the decoding stage of large language model inference, the key-value cache is adaptively compressed based on the allocated key-value cache compression strategies.

[0011] As a further improvement to the above technical solution:

[0012] Based on the numerical distribution features of key vectors and attention scores in large language models, a lexical importance measurement mechanism with the L2 norm of key vectors as its core is constructed, including:

[0013] Calculate the L2 norm of the key vector of the i-th word in the input word sequence, satisfying the following expression:

[0014] ;

[0015] in, Let i be the key vector of the i-th word. Let j be the value of the key vector of the i-th word;

[0016] Then, the L2 norm values ​​of different layers are normalized to obtain standardized indices. It satisfies the following expression:

[0017] ;

[0018] in, This is the mean of the norms of all key vectors in the current layer. The standard deviation of the norm of all key vectors in the current layer;

[0019] Attention loss due to key-value cache pruning Defined as the sum of attention scores associated with the discarded key-value pairs, it is used to establish the mapping relationship between compression strength and the performance of large language models, satisfying the following expression:

[0020] ;

[0021] in, Let be the attention score of the i-th token in the h-th attention head of the l-th layer of the large language model, and let D represent the positions of the k discarded key-value pairs, satisfying the following condition. ;

[0022] An adaptive threshold adjustment function is defined based on the attention loss function, satisfying the following expression:

[0023] ;

[0024] in, For reference to attention loss, This is to account for the attention loss caused by discarding key-value vector pairs with low L2 norm;

[0025] The cross-layer weighted attention loss index ALI is calculated, satisfying the following expression:

[0026] ;

[0027] Where H is the total number of attention heads in the large language model, L is the number of layers in the large language model, and N is the total number of tokens in the input sequence.

[0028] As a further improvement to the above technical solution:

[0029] Various pruning strategies include:

[0030] : Only special terms are retained in the key-value cache; the corresponding attention head distribution pattern is an attention head that pays more attention to special terms;

[0031] : Only punctuation words are retained in the key-value cache; the corresponding attention head distribution pattern is an attention head that pays more attention to punctuation words;

[0032] Only retain terms in the key-value cache The size of the nearby context window is lexical elements, among which is a hyperparameter; the corresponding attention head distribution pattern is an attention head that pays more attention to words in the recent context window;

[0033] A lexical importance measurement mechanism based on the L2 norm of key vectors adaptively retains key-value pairs corresponding to lexical terms with high importance weights in the key-value cache; the corresponding attention head distribution pattern is an attention head that focuses on the sparse distribution of semantics.

[0034] To retain all terms in the key-value cache, no pruning is performed.

[0035] As a further improvement to the above technical solution:

[0036] The adaptive key-value caching compression hybrid strategy set based on pruning strategy includes:

[0037] Selecting a pruning strategy:

[0038] For the pruning strategy P, define the corresponding key-value cache compression, satisfying the following expression:

[0039] ;

[0040] in, and It is a key-value cache compressed using the pruning strategy P; and It is the key-value cache before compression by the pruning strategy P;

[0041] With cache utilization and attention recovery as joint optimization objectives, an adaptive strategy selection formula is constructed, satisfying the following expression:

[0042] ;

[0043] in, The optimal strategy is one that minimizes memory cost and recovery rate. This indicates the video memory usage of strategy P for the key-value cache. Indicates attention deficit. This is the interlayer adaptive balance factor;

[0044] Define a set of hybrid strategies based on a greedy approach, starting with L2 strategies and special lexical strategies. As the key-value cache budget increases, add hybrid strategies to the set of greedy hybrid strategies to obtain an adaptive key-value cache compression hybrid strategy set. It satisfies the following expression:

[0045] ;

[0046] in, The union operation represents the set of policies. express and union, express , and union, express , , and The union of .

[0047] As a further improvement to the above technical solution:

[0048] Automatically assign key-value caching compression strategies during the pre-filling stage of large language model inference, including:

[0049] Automatically analyze the distribution of attention scores during reasoning in large language models and identify the distribution patterns corresponding to each attention head;

[0050] Based on the distribution pattern corresponding to the identified attention head, a corresponding adaptive key-value caching compression hybrid strategy set is generated. Each strategy in the adaptive key-value caching compression hybrid strategy set corresponds to a combination of one or more attention head distribution patterns.

[0051] Build an initial key-value cache for each strategy.

[0052] As a further improvement to the above technical solution:

[0053] During the decoding phase of large language model inference, adaptive compression of the key-value cache is performed based on the allocated key-value cache compression strategy, including:

[0054] The final compression strategy is selected from the set of adaptive key-value cache compression hybrid strategies generated during the pre-filling stage. ;

[0055] The selected final compression strategy is used to compress the key-value cache. The compressed key-value cache is then used to generate the lexical units for the next time step, satisfying the following expression:

[0056] ;

[0057] in, and It is after the final compression strategy Compressed key-value cache; Indicates adoption right and Compress; and The final compression strategy Key-value cache before compression.

[0058] As a further improvement to the above technical solution:

[0059] The final compression strategy Select using the following method:

[0060] In scenarios where multiple strategies coexist, a weighted fusion mechanism is used to select the final compression strategy. It satisfies the following expression:

[0061] ;

[0062] in, This represents the i-th compression strategy. This represents the dynamic fusion weight of the i-th compression strategy, which is adaptively adjusted based on the ALI value and key-value cache budget of each layer.

[0063] As a further improvement to the above technical solution:

[0064] The method also includes evaluating the compression effect of the key-value cache, and the evaluation metrics include:

[0065] The perplexity factor PPL is calculated according to the following expression:

[0066] ;

[0067] in, Let be the i-th word in the input; ANLL is the average natural log-likelihood. The higher the PPL, the lower the performance of the large language model in the generation task.

[0068] The F1 score is calculated according to the following expression:

[0069] ;

[0070] Where Precision is the accuracy, Recall is the recall, and F1 is the harmonic mean of precision and recall. A higher F1 score indicates a higher degree of matching between the results generated by the large language model and the real answer.

[0071] The code execution pass rate, Pass@1, is calculated using the following method:

[0072] For each test case in the code generation task, the large language model generates code output only once. If the output passes all the preset unit tests, it is recorded as passed. Pass@1 is the proportion of the test cases that pass.

[0073] As a further improvement to the above technical solution:

[0074] The evaluation of the key-value cache compression effect uses test tasks including long text generation, keyword retrieval, code generation, and mathematical reasoning.

[0075] In a second aspect, a computer system is provided, comprising: a processor, a memory, and a computer program stored in the memory and executable on the processor, wherein the processor, when executing the computer program, implements the adaptive KV cache compression method for large language models as described above.

[0076] Compared with the prior art, the advantages of the present invention are as follows:

[0077] (1) The present invention provides an adaptive KV cache compression method for large language models. By constructing a lexical importance substitution metric mechanism based on the L2 norm of key vectors and attention scores in existing models, it solves the compatibility problem between existing lexical importance evaluation methods based on attention weights and system-level optimization. Unlike existing compression methods based solely on attention scores or fixed thresholds, the present invention achieves cross-layer adaptive compression control by normalizing the L2 norm distribution of key vectors and dynamically adjusting the threshold. This not only avoids the computational dependency of explicitly accessing the attention matrix, but also automatically adjusts the compression intensity according to the characteristics of different layers and attention heads, thus balancing compression efficiency and model performance. Furthermore, this strategy no longer directly uses the size of the L2 norm as a static selection condition, but instead maps the norm information to cross-layer adjustable weight signals. By introducing intra-layer normalization and threshold adjustment mechanisms, the pruning behavior can be dynamically adjusted according to the semantic features of different attention heads, realizing an algorithm upgrade from "static norm pruning" to "adaptive norm weighted compression". By systematically analyzing the horizontal and vertical dimensional distribution structure of attention heads in a large language model, we can distinguish the semantic abstraction differences between different layers vertically and characterize the functional heterogeneity of different attention heads within the same layer horizontally, thereby enabling us to design targeted pruning and retention strategies.

[0078] (2) The method of this invention fully considers the differences in the role of different attention heads on different lexical types in the inference process, significantly reduces the damage to key semantic information during lexical elimination, and ensures the inference stability and generation quality of the model under compression conditions. The proposed static adaptive key-value cache compression method realizes adaptive management of KV cache in the inference stage by automatically allocating static compression strategies in the pre-filling stage and dynamically executing strategy mapping in the decoding stage. It realizes automatic selection of the optimal compression scheme according to different task types (such as dialogue, summarization, code generation, etc.), which significantly improves cache utilization and inference speed. Attached Figure Description

[0079] Figure 1 This is a flowchart of an adaptive key-value caching compression method for large language models according to an embodiment of the present invention. Detailed Implementation

[0080] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0081] like Figure 1 As shown, this implementation provides an adaptive key-value caching compression method for large language models. Taking the mainstream large language model DeepSeek-Llama3-8B-Distill (Decoder-only architecture, 32-layer model structure, 8 attention heads per layer, vocabulary size of 128,000) as the target, it addresses the scenario of "multi-turn dialogue long text reasoning" (input sequence length up to 8192 tokens, covering 5-8 turns of dialogue, including instruction tokens [INST], start tokens, and various punctuation marks). The method includes the following steps:

[0082] S1: Based on the numerical distribution characteristics of key vectors and attention scores in large language models, a lexical importance measurement mechanism with the L2 norm of key vectors as its core is constructed, including:

[0083] Calculate the L2 norm of the key vector of the i-th word in the input word sequence, satisfying the following expression:

[0084] ;

[0085] Then, the L2 norm values ​​of different layers are normalized to obtain standardized indices that satisfy the following expression:

[0086] Attention loss due to key-value cache pruning Defined as the sum of attention scores associated with the discarded key-value pairs, it is used to establish the mapping relationship between compression strength and the performance of large language models, satisfying the following expression:

[0087] ;

[0088] in, Let be the attention score of the i-th token in the h-th attention head of the l-th layer of the large language model, and let D represent the positions of the k discarded key-value pairs, satisfying the following condition. ;

[0089] An adaptive threshold adjustment function is defined based on the attention loss function, satisfying the following expression:

[0090] ;

[0091] in, For reference to attention loss, To mitigate the attention loss caused by discarding key-value pairs with low L2 norm, R is used to dynamically adjust the compression ratio of different layers and attention heads. When R is large, the pruning intensity is reduced to ensure the performance of large language models.

[0092] The cross-layer weighted attention loss index ALI is calculated, satisfying the following expression:

[0093] ;

[0094] Where H is the total number of attention heads in the large language model, L is the number of layers in the large language model, and N is the total number of tokens in the input sequence. By calculating ALI across model layers and attention heads during inference of the large language model as a cross-layer adaptive adjustment signal, this invention can achieve self-adjustment of compression intensity based on norm distribution information without accessing the attention matrix, thus ensuring the universality and scalability of the method.

[0095] In this embodiment, the strategy of lexical importance measurement mechanism based on K-vector L2 norm is described. Specifically, the following steps are taken: First, the L2 norm of the K-vector of each word is normalized within the layer and thresholded to obtain the importance weight. Then, based on the compressive strength parameters of each attention head... Determine the retention ratio Ultimately, only those conditions are retained. > The key-value vector pairs corresponding to the lexical units. This mode corresponds to an attention head that focuses more on the importance changes of low-norm regions, and can adaptively adjust the compression intensity between different layers to ensure that key semantic information is still preserved under high compression ratio conditions.

[0096] The above steps calculate the L2 norm of the key vector of each word in the input sequence, and then normalize and correct the inter-layer weights to form a word importance weight matrix, providing a unified and calculable index for subsequent key-value caching and compression strategies. This step does not rely on the explicit calculation of attention scores and can be obtained directly from the intermediate features of the model, avoiding the dependence on access to the attention layer results, thus maintaining compatibility under system-level optimizations such as Flash Attention and PagedAttention.

[0097] S2: Analyze the attention head distribution structure in large language model inference from two dimensions: vertical (model layer) and horizontal (different attention heads within the same layer); the attention head distribution pattern is obtained, including:

[0098] : Only special terms are retained in the key-value cache; the corresponding attention head distribution pattern is an attention head that pays more attention to special terms, such as the starting term. <s>And instruction words [INST], etc. These words usually contain key information such as dialogue topics and task instructions, and are indispensable in multi-turn dialogues and instruction-following tasks;

[0099] Only punctuation mark words are retained in the key-value cache; the corresponding attention head distribution pattern is to focus more on punctuation mark words, such as ""., ":", etc. Punctuation marks are key markers for text semantic segmentation, and retaining their key-value vector pairs can help the model maintain the integrity of sentence structure.

[0100] Only retain terms in the key-value cache The size of the nearby context window is lexical elements, among which is a hyperparameter; the corresponding attention head distribution pattern is an attention head that pays more attention to words in the recent context window, which aligns with the intuition that "recent information is more relevant";

[0101] A lexical importance measurement mechanism based on the L2 norm of key vectors adaptively retains key-value pairs corresponding to lexical terms with high importance weights in the key-value cache; the corresponding attention head distribution pattern is an attention head that focuses on the sparse distribution of semantics.

[0102] To retain all terms in the key-value cache, no pruning is performed;

[0103] This step analyzes the distribution structure of attention heads in a large language model from two dimensions: the vertical dimension is the model layer, and the horizontal dimension is the different attention heads in the same model layer. The attention heads are divided into 5 modes, corresponding to attention to special words, punctuation words, nearby words, semantically sparsely distributed words, and all words, respectively. This fully considers the heterogeneous impact of different attention heads on inference performance and solves the forgetting problem during model inference.

[0104] S3: Based on the lexical importance measurement mechanism and the attention head distribution pattern, construct multiple pruning strategies, and design an adaptive key-value caching and compression hybrid strategy set based on the pruning strategies to achieve dynamic allocation of strategies among different levels and different attention heads, including:

[0105] After defining five different pruning strategies, it is necessary to select an appropriate pruning strategy:

[0106] For pruning strategy P, define the corresponding key-value cache compression, perform optimal single-strategy selection, and calculate the compressed key vector and value vector, satisfying the following expression:

[0107] ;

[0108] in, and It is a key-value cache compressed using the pruning strategy P; and It is the key-value cache before compression by the pruning strategy P;

[0109] With cache utilization and attention recovery as joint optimization objectives, an adaptive strategy selection formula is constructed, satisfying the following expression:

[0110] ;

[0111] in, The optimal strategy is one that minimizes memory cost and recovery rate. This indicates the video memory usage of strategy P for the key-value cache. Indicates attention deficit. The inter-layer adaptive balancing factor is the attention loss index weighted across layers. Dynamically determined to balance compression efficiency and large language model performance; when When the value is large, the algorithm tends to prioritize compression efficiency; when the value is small, the algorithm tends to prioritize performance preservation. Through this adaptive optimization mechanism, the present invention can dynamically select the optimal strategy under different tasks and different layer depths without manually setting a fixed threshold or ratio.

[0112] Determining the optimal strategy Then, a set of hybrid strategies is selected based on the key-value cache budget. A set of hybrid strategies based on a greedy method is defined, starting with L2 strategy and special lexical strategy. As the key-value cache budget increases, hybrid strategies are added to the set of hybrid strategies based on the greedy method, resulting in an adaptive key-value cache compression hybrid strategy set. It satisfies the following expression:

[0113] ;

[0114] in, This represents the union operation of policy sets, i.e., multiple policies taking effect simultaneously; express and union, express , and union, express , , and The union of .

[0115] In scenarios where multiple strategies coexist, a weighted fusion mechanism is used to select a hybrid strategy, satisfying the following expression:

[0116] ;

[0117] in, This represents the dynamic fusion weights of each strategy, which are adaptively adjusted based on the ALI value and key-value cache budget at each layer. .

[0118] S4: Construct a static adaptive key-value cache compression method based on a set of adaptive key-value cache compression hybrid strategies. This method is used to automatically allocate key-value cache compression strategies during the pre-filling stage of large language model inference. During the decoding stage of large language model inference, the key-value cache is adaptively compressed based on the allocated key-value cache compression strategies, thereby achieving coordinated optimization of reducing memory usage and expanding the context window during the inference stage.

[0119] The static adaptive key-value cache compression method performs the following steps during the pre-filling stage of large model inference:

[0120] Step 1: In the pre-filling stage, the model profiling tool (AttentionProfiler based on vLLM) is used to statistically analyze the types of attentional terms of each attention head in real time and automatically match 5 pruning strategies;

[0121] Step 2: Based on the distribution pattern corresponding to the identified attention head, select the corresponding strategy set from the hybrid strategy set according to the memory budget of the current task, and generate the corresponding adaptive key-value cache compression hybrid strategy set. Each strategy in the adaptive key-value cache compression hybrid strategy set corresponds to a combination of one or more attention head distribution patterns.

[0122] Step 3: Perform initial compression on the key-value cache generated in the pre-filling stage according to the allocation strategy to build the initial compressed key-value cache.

[0123] The adaptive key-value caching compression method performs the following steps during the decoding phase of large model inference:

[0124] Step 1: Dynamic strategy invocation, selecting the final compression strategy from the set of adaptive key-value cache compression hybrid strategies generated in the pre-filling stage. When generating each token via autoregression, the system monitors memory usage in real time (via NVIDIA-SMI). If memory usage exceeds the budget threshold, a more compact strategy is selected from the strategy set; if memory usage is below the threshold, a more lenient strategy is switched to.

[0125] Step 2: Real-time compression of key-value cache, using the selected final compression strategy. The key-value cache is compressed, and the compressed key-value cache is used to generate the lexical units for calculating the next time step, satisfying the following expression:

[0126] ;

[0127] in, and It is after the final compression strategy Compressed key-value cache; Indicates adoption right and Compress; and The final compression strategy Key-value cache before compression.

[0128] Compressed and Input the attention layer of the model and calculate the generation probability of the next word.

[0129] Step 3: Context window expansion. When the length of the input sequence exceeds the model's default context window, ... (Recent lexical retention strategy) and (Adaptive metric strategy based on L2 norm) Eliminate early non-critical key-value pairs and expand the context window.

[0130] S5: Evaluation of key-value cache compression performance, including the following metrics:

[0131] The perplexity factor PPL is calculated according to the following expression:

[0132] ;

[0133] Wherein, ANLL is the average natural log-likelihood. The higher the PPL, the lower the performance of the large language model in the generation task.

[0134] The F1 score is calculated according to the following expression:

[0135] ;

[0136] Where Precision is the accuracy, Recall is the recall, and F1 is the harmonic mean of precision and recall. A higher F1 score indicates a higher degree of matching between the results generated by the large language model and the real answer.

[0137] The code execution pass rate, Pass@1, is calculated using the following method:

[0138] For each test case in the code generation task, the large language model generates code output only once. If the output passes all the preset unit tests, it is recorded as passed. Pass@1 is the proportion of the test cases that pass.

[0139] Finally, the large language model using the static adaptive key-value caching compression method is evaluated using long text generation test tasks, keyword retrieval task, code generation task, and mathematical reasoning task.

[0140] The long text generation test task includes the LongBench task, the ZeroSCROLLS task, and the InfiniteBench task.

[0141] The keyword retrieval task includes a needle in a haystack task and a password retrieval task.

[0142] The code generation and mathematical reasoning tasks include the HumanEval task, the GSM8k task, and the NaturalQuestion task.

[0143] The compatibility test results of the method in this embodiment on different inference frameworks and hardware are as follows:

[0144] (1) Framework compatibility: It can run normally in vLLM, TensorRT-LLM and Triton Inference Server frameworks, and is compatible with commonly used inference acceleration frameworks such as Flash Attention 2 and Fused Attention;

[0145] (2) Hardware compatibility: It can run stably on NVIDIA RTX 3090 24GB and NVIDIA RTX 4090 24GB, and supports 8192 token sequence reasoning on edge devices (NVIDIA Jetson AGX Orin 32GB).

[0146] This embodiment also provides a computer system, including: a processor, a memory, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, it implements the adaptive KV cache compression method for large language models as described above.

[0147] The computer system can be a computing device such as a mobile phone, desktop computer, laptop, handheld computer, or cloud server. The computer system may include, but is not limited to, a processor and memory. For example, the computer system may also include input / output devices, network access devices, and a bus.

[0148] The above description is merely a preferred embodiment of the present invention, and the scope of protection of the present invention is not limited to the above embodiments. For those skilled in the art, improvements and modifications obtained without departing from the inventive concept should also be considered within the scope of protection of the present invention.< / s>

Claims

1. An adaptive key-value caching compression method for large language models, characterized in that, The method includes: S1: Based on the numerical distribution characteristics of key vectors and attention scores in large language models, a lexical importance measurement mechanism with the L2 norm of key vectors as the core is constructed. S2: Analyze the attention head distribution structure in the reasoning of the large language model to obtain the attention head distribution pattern; S3: Based on the lexical importance measurement mechanism and the attention head distribution pattern, construct a variety of pruning strategies, and design an adaptive key-value caching compression hybrid strategy set based on the pruning strategies; S4: Construct a static adaptive key-value cache compression method based on an adaptive key-value cache compression hybrid strategy set, which is used to automatically allocate key-value cache compression strategies during the pre-filling stage of large language model inference; and to adaptively compress the key-value cache based on the allocated key-value cache compression strategies during the decoding stage of large language model inference. Based on the numerical distribution features of key vectors and attention scores in large language models, a lexical importance measurement mechanism with the L2 norm of key vectors as its core is constructed, including: Calculate the L2 norm of the key vector of the i-th word in the input word sequence, satisfying the following expression: ； in, Let i be the key vector of the i-th word. Let j be the value of the key vector of the i-th word; Then, the L2 norm values ​​of different layers are normalized to obtain standardized indices. It satisfies the following expression: ； in, This is the mean of the norms of all key vectors in the current layer. The standard deviation of the norm of all key vectors in the current layer; Attention loss due to key-value cache pruning Defined as the sum of attention scores associated with the discarded key-value pairs, it is used to establish the mapping relationship between compression strength and the performance of large language models, satisfying the following expression: ； in, Let be the attention score of the i-th token in the h-th attention head of the l-th layer of the large language model, and let D represent the positions of the k discarded key-value pairs, satisfying the following condition. ; An adaptive threshold adjustment function is defined based on the attention loss function, satisfying the following expression: ； in, For reference to attention loss, This is to account for the attention loss caused by discarding key-value vector pairs with low L2 norm; The cross-layer weighted attention loss index ALI is calculated, satisfying the following expression: ； Where H is the total number of attention heads in the large language model, L is the number of layers in the large language model, and N is the total number of tokens in the input sequence; The adaptive key-value caching compression hybrid strategy set based on pruning strategy includes: With cache utilization and attention recovery as joint optimization objectives, an adaptive strategy selection formula is constructed, satisfying the following expression: ； in, The optimal strategy is one that minimizes memory cost and recovery rate. This indicates the video memory usage of strategy P for the key-value cache. Indicates attention deficit. This is the interlayer adaptive balance factor.

2. The adaptive KV cache compression method for large language models according to claim 1, characterized in that, Various pruning strategies include: : Only special terms are retained in the key-value cache; the corresponding attention head distribution pattern is an attention head that pays more attention to special terms; : Only punctuation words are retained in the key-value cache; the corresponding attention head distribution pattern is an attention head that pays more attention to punctuation words; Only retain terms in the key-value cache The size of the nearby context window is lexical elements, among which is a hyperparameter; the corresponding attention head distribution pattern is an attention head that pays more attention to words in the recent context window; A lexical importance measurement mechanism based on the L2 norm of key vectors adaptively retains key-value pairs corresponding to lexical terms with high importance weights in the key-value cache; the corresponding attention head distribution pattern is an attention head that focuses on the sparse distribution of semantics. To retain all terms in the key-value cache, no pruning is performed.

3. The adaptive KV cache compression method for large language models according to claim 2, characterized in that, The adaptive key-value caching compression hybrid strategy set based on pruning strategy includes: Selecting a pruning strategy: For the pruning strategy P, define the corresponding key-value cache compression, satisfying the following expression: ； in, and It is a key-value cache compressed using the pruning strategy P; and It is the key-value cache before compression by the pruning strategy P; Define a set of hybrid strategies based on a greedy approach, starting with L2 strategies and special lexical strategies. As the key-value cache budget increases, add hybrid strategies to the set of greedy hybrid strategies to obtain an adaptive key-value cache compression hybrid strategy set. It satisfies the following expression: ； in, The union operation represents the set of policies. express and The union of, express , and union, express , , and The union of .

4. The adaptive KV cache compression method for large language models according to claim 3, characterized in that, Automatically assign key-value caching compression strategies during the pre-filling stage of large language model inference, including: Automatically analyze the distribution of attention scores during reasoning in large language models and identify the distribution patterns corresponding to each attention head; Based on the distribution pattern corresponding to the identified attention head, a corresponding adaptive key-value caching compression hybrid strategy set is generated. Each strategy in the adaptive key-value caching compression hybrid strategy set corresponds to a combination of one or more attention head distribution patterns. Build an initial key-value cache for each strategy.

5. The adaptive KV cache compression method for large language models according to claim 4, characterized in that, During the decoding phase of large language model inference, adaptive compression of the key-value cache is performed based on the allocated key-value cache compression strategy, including: The final compression strategy is selected from the set of adaptive key-value cache compression hybrid strategies generated during the pre-filling stage. ; The selected final compression strategy is used to compress the key-value cache. The compressed key-value cache is then used to generate the lexical units for the next time step, satisfying the following expression: ； in, and It is after the final compression strategy Compressed key-value cache; Indicates adoption right and Compress; and The final compression strategy Key-value cache before compression.

6. The adaptive KV cache compression method for large language models according to claim 5, characterized in that, The final compression strategy Select using the following method: In scenarios where multiple strategies coexist, a weighted fusion mechanism is used to select the final compression strategy. It satisfies the following expression: ； in, This represents the i-th compression strategy. This represents the dynamic fusion weight of the i-th compression strategy, which is adaptively adjusted based on the ALI value and key-value cache budget of each layer.

7. The adaptive KV cache compression method for large language models according to claim 6, characterized in that, The method also includes evaluating the compression effect of the key-value cache, and the evaluation metrics include: The perplexity factor PPL is calculated according to the following expression: ； in, Let be the i-th word in the input; ANLL is the average natural log-likelihood. The higher the PPL, the lower the performance of the large language model in the generation task. The F1 score is calculated according to the following expression: ； Where Precision is the accuracy, Recall is the recall, and F1 is the harmonic mean of precision and recall. A higher F1 score indicates a higher degree of matching between the results generated by the large language model and the real answer. The code execution pass rate, Pass@1, is calculated using the following method: For each test case in the code generation task, the large language model generates code output only once. If the output passes all the preset unit tests, it is recorded as passed. Pass@1 is the proportion of the test cases that pass.

8. The adaptive KV cache compression method for large language models according to claim 7, characterized in that, The evaluation of the key-value cache compression effect uses test tasks including long text generation, keyword retrieval, code generation, and mathematical reasoning.

9. A computer system, characterized in that, It includes a processor, a memory, and a computer program stored in the memory and executable on the processor, wherein the processor, when executing the computer program, implements the adaptive KV cache compression method for large language models as described in any one of claims 1-8.

Images