Key value cache pruning method and system
By calculating the importance evaluation value of the key-value cache vector, key-value pairs with large activation values are identified and retained. This solves the problem of slowed inference speed and increased memory cost in long context processing of large language models, and achieves reduced memory requirements and improved performance.
Patent Information
- Application Number
- CN202510794368.8
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-06-13
- Publication Date
- 2025-11-07
AI Technical Summary
Large language models experience slower inference speeds and increased memory costs when dealing with long contexts, and existing technologies struggle to effectively address the storage requirements and performance issues of key-value caching.
By calculating the importance evaluation value of the key-value cache vector, key-value pairs with large activation values are identified and retained, and sparse pruning is performed to reduce unnecessary cache usage, supporting multi-head attention and grouped query attention architectures.
It significantly reduces memory requirements while maintaining model performance, improves inference speed and efficiency, and is suitable for deployment in resource-constrained environments.
Smart Images
Figure CN120910092A_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The present application relates to the technical field of artificial intelligence, and in particular to a key-value cache pruning method and system. BACKGROUND
[0002] For generative large models, it is crucial to extend the context length. However, it is a major challenge for large models to handle very long context windows. On the one hand, this leads to slower inference speed, and on the other hand, it leads to increased memory cost. During the process of generating content by the model, the computed attention keys and values are cached in memory. However, the key-value cache size varies linearly with the context length, thereby bringing challenges to long context inference. SUMMARY
[0003] The present application provides a key-value cache pruning method and system to solve the defects of slow inference speed and increased memory cost of large models in the prior art. The method of the present application not only reduces the storage requirements of key-value cache, but also maintains superior model performance.
[0004] The present application provides a key-value cache pruning method, comprising: obtaining a query vector of a large language model, a set of key cache vectors to be pruned, and a set of value cache vectors to be pruned; calculating a key cache vector importance evaluation value according to the query vector and the set of key cache vectors to be pruned; the key cache vector importance evaluation value is used to identify key cache vectors with large activation values in the set of key cache vectors to be pruned; calculating a value cache vector importance evaluation value according to the set of value cache vectors to be pruned; the value cache vector importance evaluation value is used to identify value cache vectors with large activation values in the set of value cache vectors to be pruned; determining a key-value cache importance measurement result according to the key cache vector importance evaluation value and the value cache vector importance evaluation value; performing key-value cache pruning on the set of key cache vectors to be pruned and the set of value cache vectors to be pruned based on the key-value cache importance measurement result, to retain key cache vectors with large activation values and value cache vectors with large activation values.
[0005] According to the key-value cache pruning method provided by the present application, the query vector of the large language model, the set of key cache vectors to be pruned, and the set of value cache vectors to be pruned are obtained, comprising: in the dialogue application of the large language model, determining the query vector according to the query sentence input by the user; based on the query vector, determining the set of key cache vectors and the set of value cache vectors generated in the pre-filling stage in the inference process of the large language model; according to the set of key cache vectors and the set of value cache vectors in the pre-filling stage, respectively taking the set of key cache vectors and the set of value cache vectors generated in the decoding stage as the set of key cache vectors to be pruned and the set of value cache vectors to be pruned.
[0006] According to the key-value cache pruning method provided by the application, the key cache vector importance evaluation value is calculated according to the query vector and the key cache vector set to be pruned, and the key cache vector importance evaluation value is calculated according to the query vector and the transposed key cache vector set to be pruned.
[0007] According to the key-value cache pruning method provided by the application, the attention weight of the attention head is calculated according to the query vector and the key cache vector set to be pruned, and the query vector and the key cache vector set to be pruned are dot multiplied, and then divided by the square root of the key vector dimension of the key cache vector set to be pruned to obtain the scaled dot product; the scaled dot product is converted into the attention weight of the attention head by the Softmax function.
[0008] According to the key-value cache pruning method provided by the application, the value cache vector importance evaluation value is calculated according to the value cache vector set to be pruned, and the maximum absolute value of the value cache vector in the value cache vector set to be pruned in all feature dimensions is calculated to obtain the value cache vector importance evaluation value.
[0009] According to the key-value cache pruning method provided by the application, the method supports multi-head attention and grouped query attention architecture.
[0010] The application also provides a key-value cache pruning system, comprising: an acquisition module for acquiring a query vector of a large language model, a key cache vector set to be pruned and a value cache vector set to be pruned; a first calculation module for calculating a key cache vector importance evaluation value according to the query vector and the key cache vector set to be pruned; the key cache vector importance evaluation value is used to identify key cache vectors with large activation values in the key cache vector set to be pruned; a second calculation module for calculating a value cache vector importance evaluation value according to the value cache vector set to be pruned; the value cache vector importance evaluation value is used to identify value cache vectors with large activation values in the value cache vector set to be pruned; a metric result determination module for determining a key-value cache importance metric result according to the key cache vector importance evaluation value and the value cache vector importance evaluation value; a pruning module for performing key-value cache pruning on the key cache vector set to be pruned and the value cache vector set to be pruned based on the key-value cache importance metric result, so as to retain key cache vectors with large activation values and value cache vectors with large activation values.
[0011] The application further provides an electronic device, including a memory, a processor, and a computer program stored in the memory and executable on the processor, and the processor implements the key-value cache pruning method according to any one of the above when executing the computer program.
[0012] The application further provides a non-transitory computer-readable storage medium, which stores a computer program, and the computer program is executed by a processor to implement the key-value cache pruning method according to any one of the above.
[0013] The application further provides a computer program product, which includes a computer program, and the computer program is executed by a processor to implement the key-value cache pruning method according to any one of the above.
[0014] The application provides a key-value cache pruning method and system, which first acquires a query vector of a large language model, a key cache vector set to be pruned, and a value cache vector set to be pruned; calculates a key cache vector importance evaluation value according to the query vector and the key cache vector set to be pruned; calculates a value cache vector importance evaluation value according to the value cache vector set to be pruned; determines a key-value cache importance measurement result according to the key cache vector importance evaluation value and the value cache vector importance evaluation value; and performs key-value cache pruning on the key cache vector set to be pruned and the value cache vector set to be pruned based on the key-value cache importance measurement result, so as to retain key cache vectors with large activation values and value cache vectors with large activation values. The method provided by the application not only reduces the storage requirement of the key-value cache, but also maintains superior model performance. BRIEF DESCRIPTION OF DRAWINGS
[0015] In order to more clearly illustrate the technical solutions of the present application or the prior art, the following will briefly introduce the drawings needed in the embodiments or prior art description. Obviously, the drawings in the following description are some embodiments of the present application, and those skilled in the art can also obtain other drawings according to these drawings without creative labor.
[0016] Figure 1 is a flowchart of a key-value cache pruning method provided by the application.
[0017] Figure 2 is a principle diagram of a key-value cache pruning method provided by the application.
[0018] Figure 3 is a principle diagram of a large model dialogue application provided by the application.
[0019] Figure 4 is a structure diagram of a key-value cache pruning system provided by the application.
[0020] Figure 5It is a structural schematic diagram of an electronic device provided by the application. DETAILED DESCRIPTION
[0021] To make the objects, technical solutions and advantages of the present application clearer, the technical solutions in the present application will be described below in connection with the drawings in the present application. Obviously, the described embodiments are only some of the embodiments of the present application, rather than all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those of ordinary skill in the art without creative labor fall within the scope of protection of the present application.
[0022] In large language models, the attention mechanism is one of the most computationally expensive parts. In each step of generation, the attention representation of only the last token in the input sequence is used to predict the next token. However, the model still calculates the attention representation of all tokens in parallel, which produces a large amount of redundant calculations, and the longer the input length, the more redundant calculations are produced. Key value cache (KV Cache, Key Value Cache) is a method for accelerating large model inference, which caches the key vectors and value vectors in attention, so that only the similarity between the new token and the cached key vectors and value vectors needs to be calculated during inference. Without affecting any calculation accuracy, the idea of space for time is used to improve inference performance. However, as the context length used by the large language model increases, due to the delay of accessing a large number of key value caches and excessive consumption of storage resources, a huge inference cost will be generated.
[0023] Reference is made to Figure 1 , Figure 1 A flowchart of a key value cache pruning method provided by the application is shown.
[0024] Reference is made to Figure 2 , Figure 2 A principle diagram of a key value cache pruning method provided by the application is shown.
[0025] The application provides a key value cache pruning method, comprising: 101: obtaining a query vector of a large language model, a set of key cache vectors to be pruned, and a set of value cache vectors to be pruned; 102: calculating a key cache vector importance evaluation value according to the query vector and the set of key cache vectors to be pruned; the key cache vector importance evaluation value is used to identify key cache vectors with large activation values in the set of key cache vectors to be pruned; 103: calculating a value cache vector importance evaluation value according to the set of value cache vectors to be pruned; the value cache vector importance evaluation value is used to identify value cache vectors with large activation values in the set of value cache vectors to be pruned; 104: determining a key-value cache importance measure result according to the key cache vector importance evaluation value and the value cache vector importance evaluation value; 105: performing key-value cache pruning on the key cache vector set to be pruned and the value cache vector set to be pruned based on the key-value cache importance measure result, to retain key cache vectors with large activation values and value cache vectors with large activation values.
[0026] The present application provides a key-value cache pruning method, which performs sparse inference by retaining only a small number of important key-value caches. The present application discloses that large activation is the key to identifying these important key-value caches, and is compatible with grouped query attention. The pruning method of the present application operates within each attention group to improve efficiency and minimize performance degradation. The method of the present application can exhibit similar accuracy while significantly reducing key-value cache requirements when using only 1 / 10 of the key-value caches.
[0027] Specifically, first, the query vector (Q) of the large language model, the key cache vector set (K) to be pruned, and the value cache vector set (V) to be pruned are obtained. The query vector is extracted from the input sequence of the model, while the key cache vector set and the value cache vector set are generated in the self-attention layer of the model. Next, according to the query vector and the key cache vector set to be pruned, the key cache vector importance evaluation value S1 is calculated. This step involves calculating the similarity between each key vector and the query vector, usually through the dot product operation, and then applying the Softmax function to obtain the weight of each key vector. These weights reflect the importance of the key vector in the model, and the key vector with a larger weight is considered to have a large activation value. Similarly, according to the value cache vector set to be pruned, the value cache vector importance evaluation value S2 is calculated. This step involves identifying the maximum activation value in the value vector, which can be obtained by calculating the maximum absolute value of the value vector in all feature dimensions. The value vector with a high S2 value is considered to have a large activation value. Then, according to the key cache vector importance evaluation value S1 and the value cache vector importance evaluation value S2, the key-value cache importance measure result is determined. This step may involve combining S1 and S2 to form a comprehensive measure for evaluating the importance of each key-value pair in the model. Finally, based on the key-value cache importance measure result, the key cache vector set to be pruned and the value cache vector set to be pruned are pruned. During pruning, those key vectors with high S1 values and value vectors with high S2 values are retained, while those key-value pairs that have little effect on the model output are removed. In this way, not only the memory occupancy of the model is reduced, but also the performance of the model is maintained, because the key information that has a significant impact on the model output is retained. This method can significantly reduce the memory requirements of the large language model while maintaining its performance, making it more suitable for deployment in resource-constrained environments.
[0028] It is necessary to explain that large activations refer to neuron outputs with high absolute values in large language models (LLMs). The importance of large activations in the key-value cache in the pruning process in LLMs mainly lies in their significant impact on the model output. These large activations existing in the key and value caches have a significant impact on the model's calculation results, and removing these large activations will cause a significant decrease in model performance, so it is crucial to maintain large activations in the key-value cache during the pruning process.
[0029] Large activations in the key cache: The key cache is crucial in shaping the distribution of attention logits. In self-attention architectures, the key cache interacts with the query vectors to compute attention logits. When a particular key cache is important, it will have a significant impact on the Softmax computation, resulting in increased attention logits for a particular position. This increased attention indicates that the relevant value vectors, weighted according to these attention logits, will have a stronger impact on the resulting output. Therefore, pruning large activations in the key cache can weaken the model's ability to focus on important aspects of the input data, leading to significant information loss and a decrease in model performance.
[0030] Large activations in the value cache: Similarly, large activations in the value cache directly determine the size of the output components. Value vectors directly contribute to the final output after being weighted by their corresponding attention logits. Therefore, large activations in the value cache will result in greater contributions to the output, especially when they align with high attention logits. These key elements in the value cache contain information that is crucial to the accuracy and richness of the model's output. Ignoring these large activations during the pruning process can lead to the loss of critical output features, thereby reducing the overall performance and accuracy of the model.
[0031] As a preferred embodiment, obtaining the query vector of the large language model, the set of key cache vectors to be pruned, and the set of value cache vectors to be pruned includes: in the dialogue application of the large language model, determining the query vector according to the query sentence input by the user; based on the query vector, determining the set of key cache vectors and the set of value cache vectors generated in the pre-filling stage during the inference process of the large language model; and based on the set of key cache vectors and the set of value cache vectors in the pre-filling stage, respectively taking the set of key cache vectors and the set of value cache vectors generated in the decoding stage as the set of key cache vectors to be pruned and the set of value cache vectors to be pruned.
[0032] Please refer to Figure 3 , Figure 3 The principle diagram of the large model dialogue application provided by the present application is shown.
[0033] In this embodiment, in the dialog application of large language models, the user inputs a query sentence, and the user-input query sentence is converted into a query vector that is convenient for the large language model to understand. The inference process of the large language model includes two stages: a prefill stage and a decoding stage.
[0034] In the prefill stage, the model processes tokens of the input prompt in parallel to generate key-value caches for each token. This stage is computationally intensive because it involves processing the entire prompt at once. The key-value caches computed in the prefill stage for a sequence of length can be represented as: , where is the set of key cache vectors generated in the prefill stage, is the set of value cache vectors generated in the prefill stage, is the key cache vector generated in the first step of the prefill stage, is the key cache vector generated in the second step of the prefill stage, is the key cache vector generated in the prompt-length step of the prefill stage, is the value cache vector generated in the first step of the prefill stage, is the value cache vector generated in the second step of the prefill stage, is the value cache vector generated in the prompt-length step of the prefill stage.
[0035] After the prefill stage, the model generates output tokens one by one in an autoregressive manner. In each step of the decoding stage, the model attends to previously generated tokens (as well as the initial prompt) by using their key-value caches. For generating the j th output token, the query vector q j attends to the key-value caches of all preceding tokens: , where is the set of key cache vectors in the j th step of the decoding stage, is the set of value cache vectors in the j th step of the decoding stage, is the key cache vector generated in the first step of the decoding stage, is the value cache vector generated in the first step of the decoding stage, is the key cache vector generated in the j th-1 step of the decoding stage, is the value cache vector generated in thej -1 step generated value cache vector.
[0036] This process has high memory occupancy requirements due to the continuous growth of the size of the key-value cache.
[0037] In order to reduce the memory requirements of the decoding stage, the sparse key-value cache technology aims to only retain the most important entries in the key-value cache. Instead of using the complete and , the compressed version and is used. These sparse caches contain a subset of the original key-value cache pairs, which are considered to be the most critical entries to maintain performance.
[0038] , wherein, is the sparse key cache vector set of the sparse decoding stage step j , is the sparse value cache vector set of the sparse decoding stage step j .
[0039] As a preferred embodiment, the key cache vector importance evaluation value is calculated according to the query vector and the key cache vector set to be pruned, including: calculating the attention weight of the attention head according to the query vector and the key cache vector set to be pruned; by aggregating the attention weights of all attention heads in the observation window, the key cache vector importance evaluation value is calculated.
[0040] In this embodiment, in order to effectively identify the key (Key) cache vector with large activation, a method of tracing from high attention logits to its corresponding significant key (Key) cache is adopted. Specifically, the area with higher attention logits value indicates proportional impact on the model output, so it is worth keeping the keys (Key) caches related to these logits. Therefore, the key cache vectors related to these high-amplitude activations are preferentially retained to ensure that important context connections are retained during the pruning process. In addition, in order to adapt to the multi-head attention (MHA) and grouped query attention (GQA) architecture, a group voting mechanism is introduced to facilitate the identification of large activations, reducing the number of key caches in the generation stage, thereby reducing the computational burden of LLMs when processing long contexts. The following definitions are proposed in this embodiment: Prompt Length, : total length of user-provided input.
[0041] Observation Window, This is the closest part of the prompt and is crucial for analyzing how different contexts affect attention distribution.
[0042] Prefix Length This is the length of the input section preceding the view window. It is contained within the prompt but does not overlap with the view window.
[0043] The relationship between these lengths is given by the following formula: .
[0044] To obtain the importance evaluation value of the key cache vector, it is necessary to compute the attention logits of each query header in the current attention group within the observation window, and then aggregate these weights to identify the most salient position in the prefix. For a single batch sequence, the computation can be formally represented as: , in, This is the evaluation value of the importance of the key cache vector within the current attention group. , For the total number of attention heads, For the current attention group N A subset of the Softmax normalized attention logits of each attention head. i To observe the index of the position within the window, h For the index of attention heads.
[0045] As a preferred embodiment, the attention weights of the attention head are calculated based on the query vector and the set of key cache vectors to be pruned, including: taking the dot product of the query vector and the transposed set of key cache vectors to be pruned, and then dividing by the square root of the key vector dimension of the set of key cache vectors to be pruned to obtain a scaled dot product; and converting the scaled dot product into the attention weights of the attention head using the Softmax function.
[0046] In this embodiment, within the context of large language models (LLMs) and during the pre-filling phase of the inference process, specific attention groups... Attention mechanisms in groups The first in h The calculation steps for each attention point are as follows: , in, For the first h Attention weight matrix for each attention head. For the Softmax function, For the first h A query vector with attention heads. For the first a set of key cache vectors of the attention group, a set of value cache vectors of the i-th attention group, a set of value cache vectors of the i-th attention group, an output of the i-th attention head, h an index of the attention group, a dimension of the key vector, a matrix transpose. T
[0047] The attention weight is obtained by calculating the dot product of the query vector and the transpose of the set of key cache vectors , dividing by the square root of the dimension of the key vector of the set of key cache vectors to be pruned , and then applying the Softmax function to ensure that the sum of the weights is one. The final output of the attention head is obtained by weighting the set of value cache vectors with the calculated attention weight . Multi-head attention (MHA) can be regarded as a special case of grouped query attention (GQA), where each group contains exactly one query head.
[0048] As a preferred embodiment, according to the set of value cache vectors to be pruned, the value cache vector importance evaluation value is calculated, including: by calculating the maximum absolute value of the value cache vector in all feature dimensions in the set of value cache vectors to be pruned, the value cache vector importance evaluation value is obtained.
[0049] In this embodiment, in order to identify important information in the set of value cache vectors, the most direct and effective method is to focus on large activations, which are usually an indication of key parts in the model output. These large activations in the set of value cache vectors directly affect the size of the output. The importance of the value cache vector based on these large activations can be calculated in the following way: , wherein, is the value cache vector importance evaluation value within the current attention group, is the value cache vector, d is the dimension of the value cache vector, .
[0050] Captures the maximum activation amplitude of each value vector in the prefix, indicating its potential maximum contribution to the output. The operation selects the maximum absolute value of each prefix token in the feature dimension .
[0051] To comprehensively evaluate the importance of different key-value caches, an integrated indicator is proposed The indicator combines the key cache vector importance evaluation value S1 and the value cache vector importance evaluation value S2. By multiplying the two indicators, the overall importance of a specific key-value cache (key-value cache importance measurement result) is quantified.
[0052] .
[0053] As a preferred embodiment, the method supports multi-head attention and grouped query attention architecture.
[0054] Multi-head attention (MHA) enhances the model's expressive power by using multiple attention heads in parallel to capture different features in the input sequence. The basic idea of multi-head attention is to process the input query (Q), key (K) and value (V) vectors through multiple attention heads in parallel, and then concatenate the results to get the final output. Each attention head learns different features of the data in different subspaces, enabling the model to better understand complex inputs. By parallel processing through multiple attention heads, MHA can focus on information in the input data from different angles, capturing richer contextual relationships. The design of MHA allows simultaneous computation of multiple attention heads on hardware acceleration, improving computational efficiency.
[0055] Grouped query attention (GQA) aims to improve computational efficiency and enhance model capabilities by grouping queries. GQA optimizes the attention computation process by dividing queries into multiple groups and independently computing attention for each group. Queries in GQA share a common key-value cache, and this sharing mechanism aims to optimize the processing efficiency and consistency between grouped queries.
[0056] Table 1 Comparison of the impact of the method of the present invention and the prior art method on model performance Table 1 shows the comparative performance of the method of the present invention and the prior art method on various large language models (LLMs) using grouped query attention (GQA) architecture. The evaluation results are taken from the LongBench dataset, covering Llama-3-8B and Mistral-7B models, and the results of two KV cache sizes are listed in Table 1: 128 and 2048. These two sizes represent two different operation scenarios: memory-efficient scenario and performance-preservation scenario, demonstrating the balance between memory utilization and model performance.
[0057] Overall, the method (SlimKV) of the present application not only supports KV cache pruning in the GQA model architecture, but also significantly outperforms existing methods in terms of performance. In various GQA LLMs, SlimKV outperforms other existing methods: for example, in the Mistral-7B model, when the KV cache budget size is 128, SlimKV achieves an average accuracy of 30.58, which is higher than the average accuracy of 26.60 of other methods. When the KV cache budget size is 2048, this advantage is more pronounced, and the method of the present application achieves an average accuracy of 39.97, which is significantly higher than the average accuracy of 31.12 of other existing methods.
[0058] When the KV cache size budget is 2048, the average accuracy of the model pruned by the method of the present application is improved from 39.76 to 39.97. This performance improvement is not achievable by other existing methods. This phenomenon can be attributed to the existence of certain “noise” contexts, which may actually reduce the performance of the model when included. By removing the key-value cache corresponding to these interfering contexts, the performance of the model can be effectively enhanced.
[0059] The method of the present application can reduce redundant calculations, reduce the overall memory usage, and significantly improve the inference speed of large models. For long text generation tasks, it can effectively handle long sequence inference problems and improve generation efficiency and quality.
[0060] The key-value cache pruning system provided by the present application is described below. The key-value cache pruning system described below can be referred to in conjunction with the key-value cache method described above.
[0061] Please refer to Figure 4 , Figure 4 The structure diagram of a key-value cache pruning system provided by the present application.
[0062] The application further provides a key-value cache pruning system, comprising: an acquisition module 401, configured to acquire a query vector of a large language model, a key cache vector set to be pruned, and a value cache vector set to be pruned; a first calculation module 402, configured to calculate key cache vector importance evaluation values according to the query vector and the key cache vector set to be pruned; the key cache vector importance evaluation values are used to identify key cache vectors with large activation values in the key cache vector set to be pruned; a second calculation module 403, configured to calculate value cache vector importance evaluation values according to the value cache vector set to be pruned; the value cache vector importance evaluation values are used to identify value cache vectors with large activation values in the value cache vector set to be pruned; a metric result determination module 404, configured to determine a key-value cache importance metric result according to the key cache vector importance evaluation values and the value cache vector importance evaluation values; and a pruning module 405, configured to perform key-value cache pruning on the key cache vector set to be pruned and the value cache vector set to be pruned based on the key-value cache importance metric result, so as to retain the key cache vectors with large activation values and the value cache vectors with large activation values.
[0063] Figure 5 An example of a schematic diagram of a physical structure of an electronic device is shown in Figure 5 The electronic device can include a processor 501, a communications interface 502, a memory 503, and a communications bus 504, wherein the processor 501, the communications interface 502, and the memory 503 can communicate with each other through the communications bus 504. The processor 501 can invoke a logical instruction in the memory 503 to execute a key-value cache pruning method, which includes: acquiring a query vector of a large language model, a key cache vector set to be pruned, and a value cache vector set to be pruned; calculating key cache vector importance evaluation values according to the query vector and the key cache vector set to be pruned; the key cache vector importance evaluation values are used to identify key cache vectors with large activation values in the key cache vector set to be pruned; calculating value cache vector importance evaluation values according to the value cache vector set to be pruned; the value cache vector importance evaluation values are used to identify value cache vectors with large activation values in the value cache vector set to be pruned; determining a key-value cache importance metric result according to the key cache vector importance evaluation values and the value cache vector importance evaluation values; and performing key-value cache pruning on the key cache vector set to be pruned and the value cache vector set to be pruned based on the key-value cache importance metric result, so as to retain the key cache vectors with large activation values and the value cache vectors with large activation values.
[0064] In addition, the logic instructions in the memory 503 described above can be implemented in the form of a software function unit and sold or used as an independent product, which can be stored in a computer readable storage medium. Based on such understanding, the technical solutions of the present application essentially or the part that contributes to the prior art or part of the technical solutions can be embodied in the form of a software product, and the computer software product is stored in a storage medium, including a plurality of instructions to make a computer device (which can be a personal computer, a server, or a network device, etc.) execute all or part of the steps of the method described in various embodiments of the present application. The foregoing storage medium includes: a U disk, a mobile hard disk, a read-only memory (ROM, Read-Only Memory), a random access memory (RAM, Random Access Memory), a magnetic disk or an optical disk, and various media that can store program codes.
[0065] In another aspect, the present application also provides a computer program product, which comprises a computer program, the computer program can be stored on a non-transitory computer readable storage medium, and the computer program can be executed by a processor, so that the computer can execute the key-value cache pruning method provided by the above-mentioned method, which comprises: obtaining a query vector of a large language model, a set of key cache vectors to be pruned, and a set of value cache vectors to be pruned; calculating a key cache vector importance evaluation value according to the query vector and the set of key cache vectors to be pruned; the key cache vector importance evaluation value is used to identify the key cache vector with a large activation value in the set of key cache vectors to be pruned; calculating a value cache vector importance evaluation value according to the set of value cache vectors to be pruned; the value cache vector importance evaluation value is used to identify the value cache vector with a large activation value in the set of value cache vectors to be pruned; determining a key-value cache importance measurement result according to the key cache vector importance evaluation value and the value cache vector importance evaluation value; performing key-value cache pruning on the set of key cache vectors to be pruned and the set of value cache vectors to be pruned based on the key-value cache importance measurement result, so as to retain the key cache vector with a large activation value and the value cache vector with a large activation value.
[0066] In yet another aspect, the present application also provides a non-transitory computer-readable storage medium having stored thereon a computer program, which, when executed by a processor, implements the key-value cache pruning method provided by any of the above methods, and the method comprises: obtaining a query vector of a large language model, a set of key cache vectors to be pruned, and a set of value cache vectors to be pruned; calculating a key cache vector importance evaluation value according to the query vector and the set of key cache vectors to be pruned; the key cache vector importance evaluation value is used to identify a key cache vector with a large activation value in the set of key cache vectors to be pruned; calculating a value cache vector importance evaluation value according to the set of value cache vectors to be pruned; the value cache vector importance evaluation value is used to identify a value cache vector with a large activation value in the set of value cache vectors to be pruned; determining a key-value cache importance measurement result according to the key cache vector importance evaluation value and the value cache vector importance evaluation value; and performing key-value cache pruning on the set of key cache vectors to be pruned and the set of value cache vectors to be pruned based on the key-value cache importance measurement result, so as to retain the key cache vector with the large activation value and the value cache vector with the large activation value.
[0067] The device embodiments described above are merely illustrative, wherein the units described as separate components can or can not be physically separated, and the components shown as units can or can not be physical units, i.e., can be located in one place or distributed on multiple network units. Part or all of the modules can be selected to achieve the purpose of the embodiment according to actual needs. Those skilled in the art can understand and implement without creative labor.
[0068] From the above description of the embodiments, those skilled in the art can clearly understand that the embodiments can be realized by means of software plus necessary general hardware platforms, and of course can also be realized by hardware. Based on such understanding, the above technical solutions, essentially or in terms of the contribution to the prior art, can be embodied in the form of a software product. The computer software product can be stored in a computer readable storage medium, such as a ROM / RAM, a magnetic disk, an optical disk, etc., and includes a plurality of instructions to make a computer device (which can be a personal computer, a server, or a network device, etc.) execute the methods described in each embodiment or some parts of the embodiments.
[0069] Finally, it should be noted that: the above embodiments are only used to illustrate the technical solutions of the present application, and not to limit them; although the present application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that: it can still modify the technical solutions recorded in the foregoing embodiments, or make equivalent replacement to some technical features; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of the embodiments of the present application.
Claims
1. A method for pruning a key-value cache, the method comprising: The method comprises the steps of: obtaining a query vector of a large language model, a set of key cache vectors to be pruned, and a set of value cache vectors to be pruned; calculating a key cache vector importance evaluation value according to the query vector and the set of key cache vectors to be pruned; the key cache vector importance evaluation value is used to identify key cache vectors with large activation values in the set of key cache vectors to be pruned; calculating a value cache vector importance evaluation value according to the set of value cache vectors to be pruned; the value cache vector importance evaluation value is used to identify value cache vectors with large activation values in the set of value cache vectors to be pruned; determining a key-value cache importance measurement result according to the key cache vector importance evaluation value and the value cache vector importance evaluation value; performing key-value cache pruning on the set of key cache vectors to be pruned and the set of value cache vectors to be pruned based on the key-value cache importance measurement result, so as to retain key cache vectors with large activation values and value cache vectors with large activation values.
2. The key-value cache pruning method of claim 1, wherein, The method comprises the steps of: In a dialogue application of a large language model, determining the query vector according to a query sentence input by a user; based on the query vector, determining a set of key cache vectors and a set of value cache vectors generated in a pre-filling stage in an inference process of the large language model; respectively taking a set of key cache vectors and a set of value cache vectors generated in a decoding stage as the set of key cache vectors to be pruned and the set of value cache vectors to be pruned according to the set of key cache vectors and the set of value cache vectors in the pre-filling stage.
3. The key-value cache pruning method of claim 1, wherein, The method comprises the steps of: calculating attention weights of attention heads according to the query vector and the set of key cache vectors to be pruned; aggregating the attention weights of all attention heads in an observation window to calculate the key cache vector importance evaluation value.
4. The key-value cache pruning method of claim 3, wherein, The method comprises the steps of: performing dot product on the query vector and a transposed set of key cache vectors to be pruned, and then dividing the dot product by the square root of the dimension of the key vector of the set of key cache vectors to be pruned to obtain a scaled dot product; converting the scaled dot product into the attention weights of the attention heads through a Softmax function.
5. The key-value cache pruning method of claim 1, wherein, The method comprises the steps of: obtaining the value cache vector importance evaluation value by calculating the maximum absolute value of the value cache vectors in the set of value cache vectors to be pruned in all feature dimensions.
6. The key-value cache pruning method according to any one of claims 1 to 5, characterized in that, The method supports multi-head attention and grouped query attention architecture.
7. A key-value cache pruning system, comprising: The method comprises the steps of: an obtaining module, configured to obtain a query vector of a large language model, a set of key cache vectors to be pruned, and a set of value cache vectors to be pruned; The first calculation module is configured to calculate a key cache vector importance evaluation value according to the query vector and the key cache vector set to be pruned; the key cache vector importance evaluation value is used to identify a key cache vector with a large activation value in the key cache vector set to be pruned. The second calculation module is configured to calculate a value cache vector importance evaluation value according to the value cache vector set to be pruned; the value cache vector importance evaluation value is used to identify a value cache vector with a large activation value in the value cache vector set to be pruned. The metric result determination module is configured to determine a key-value cache importance metric result according to the key cache vector importance evaluation value and the value cache vector importance evaluation value. The pruning module is configured to perform key-value cache pruning on the key cache vector set to be pruned and the value cache vector set to be pruned based on the key-value cache importance metric result, so as to retain a key cache vector with a large activation value and a value cache vector with a large activation value.
8. An electronic device comprising a memory, a processor, and a computer program stored on the memory and running on the processor, characterized in that, The processor executes the computer program to implement the key-value cache pruning method according to any one of claims 1 to 6. 9.A non-transitory computer-readable storage medium having stored thereon a computer program, characterized in that, The computer program is executed by the processor to implement the key-value cache pruning method according to any one of claims 1 to 6.
10. A computer program product comprising a computer program, characterized in that, The computer program is executed by the processor to implement the key-value cache pruning method according to any one of claims 1 to 6.