A method for optimizing large-scale language model key-value cache based on recent query attention information

Abstract

This invention discloses a method for optimizing the key-value cache (KV Cache) of a large language model based on recent query attention information. The method involves calculating attention for subwords input into the large language model; constructing a record table of attention information for the most recent r queries using the average attention score as a threshold; performing an OR operation on the attention information of each recorded query and deleting key-value cache entries deemed unimportant; and determining whether the total number of generated subwords has reached a preset limit. By recording the attention information of queries in recent inference steps, this invention assesses the importance of information in the KV Cache, deleting less important KV Cache entries while retaining more important ones. Ultimately, this method optimizes over 70% of the KV Cache while maintaining minimal impact on model performance.

Description

Technical Field

[0001] This invention relates to the field of large language models, and more specifically to a method for optimizing the KV Cache of large language models based on recent query attention information. Background Technology

[0002] Large language models perform well in various natural language processing (NLP) tasks, including question answering, text summarization, and multi-turn dialogue. However, current large language models are characterized by a large number of parameters and the fact that the memory usage of the key-value cache during inference increases linearly with the sequence length. This leads to the model consuming a large amount of GPU memory and facing the problem of memory overflow when processing long texts during the inference stage.

[0003] Currently, large language models face the risk of GPU memory overflow during inference due to limited computing resources. GPU memory overhead primarily stems from two sources: the overhead of the model's parameter count itself, and the overhead caused by the linear increase in GPU memory usage by the KV Cache as the inference sequence length grows. Therefore, a solution is needed to optimize the KV Cache and reduce its associated GPU memory overhead. Summary of the Invention

[0004] To address the memory overflow problem faced by large language models during inference, this invention focuses on optimizing the memory overhead caused by the key-value cache. By utilizing the attention information of recent queries during inference, the key-value cache is optimized to reduce its memory overhead. Finally, a method for optimizing the memory usage of key-value cache during the inference process of large language models is proposed.

[0005] First, this invention proposes a method for optimizing a large-scale language model KV Cache based on recent query attention information, comprising the following steps:

[0006] 1) Segment the text to be input into the selected large language model to obtain the segmented sub-words;

[0007] 2) Input the split subwords into a large language model and perform attention calculation to obtain the attention information of each query during the inference process of the large language model; then construct a record table of attention information of the most recent r queries using the average attention score of each key in the query as the threshold.

[0008] 3) Enter the loop reasoning stage of the large language model to reason the next sub-word. After the reasoning is completed, add the attention information of the current query to the record table. If the number of queries in the table exceeds the preset number r, only the attention information of the most recent r queries will be retained. If the number of queries is less than r, proceed directly to step 5).

[0009] 4) Perform an OR operation on the attention information of each recorded query to determine whether the key-value cache in the query is important, and delete the key-value cache that is deemed unimportant;

[0010] 5) Determine whether the total number of generated subwords has reached the preset inference quantity. If not, return to step 3) to generate the next subword. If the quantity has been reached, decode the generated subword to obtain the generated text.

[0011] This invention also provides a system based on the above-mentioned large-scale language model KV Cache optimization method, comprising:

[0012] A word segmenter is used to split the text to be input into a selected large language model into sub-words and to decode the sub-words generated by the large language model to obtain the generated text.

[0013] The attention calculation module is used to perform attention calculations on subwords input into a large language model.

[0014] The record table construction module is used to construct a record table of attention information for the most recent r queries, using the average attention score as a threshold.

[0015] Attention information or operation module, used to perform OR operation on the attention information of each query in the record, and delete key-value cache that is deemed unimportant;

[0016] The sub-word count determination module is used to determine whether the total number of generated sub-words has reached the preset number.

[0017] The present invention also provides an electronic device, including a processor and a memory, wherein the memory stores machine-executable instructions that can be executed by the processor, and the processor executes the machine-executable instructions to implement the above-described large language model KV Cache optimization method.

[0018] The present invention also provides a machine-readable storage medium storing machine-executable instructions, which, when called and executed by a processor, are used to implement the above-described large language model KVCache optimization method.

[0019] Compared with the prior art, the present invention has the following beneficial effects:

[0020] Inspired by theoretical analysis and experimental observations, this invention leverages the fact that similar queries tend to have similar attention to key values. Furthermore, the highly sparse attention matrix of large language models ensures that a large amount of KVCache can be optimized away during the inference phase.

[0021] This invention proposes that the importance of key-value pairs can be determined by recording recent query attention information, and the retention and deletion of key-value pairs can be decided based on this information. During the inference stage, only the attention information of recent queries can be recorded to determine the important key-value pairs that need to be retained, while discarding the relatively less important key-value pairs, without the need for additional similarity calculation operations. This significantly reduces the GPU memory space occupied by the KV Cache of large language models during the inference stage.

[0022] The optimization method of this invention requires no additional training operations, and can maintain the inference performance of the original model while optimizing more than 70% of the KV cache.

[0023] The method proposed in this invention is applicable not only to large language models using the Transformer architecture, but also to generative models such as visual language models using the Transformer architecture. Attached Figure Description

[0024] Figure 1 , Figure 2 , Figure 3 This is a sparsity analysis diagram of LLaMA2-7B in an embodiment of the present invention;

[0025] Figure 4 This is an example of attention maps for the 0th attention head of two different layers after the same text passes through LLaMA2-7B in this embodiment of the invention.

[0026] Figure 5 This is a visualization of the cosine similarity of the query vector of a text of length 1024 on LaMA2-7B in an embodiment of the present invention.

[0027] Figure 6 This is a flowchart of the algorithm of the present invention;

[0028] Figure 7 This is a flowchart illustrating the OR operation within the record table and the deletion of key-value pairs in this invention. Detailed Implementation

[0029] The present invention will be further described and illustrated below with reference to specific embodiments. The embodiments described are merely examples of the content of this disclosure and do not limit the scope of the invention. The technical features of each embodiment in the present invention can be combined accordingly, provided that there is no mutual conflict.

[0030] To address the memory overflow issue faced by large language models during inference, this invention first analyzes the attention sparsity that exists in large language models such as LLaMA, Mistral, and Qwen during attention computation. Specifically, during attention computation in inference, only a small portion of the key-value pairs are truly important, while others do not significantly contribute to the attention calculation. In most cases, 20% of the key-value pairs receive a total attention score exceeding 0.9, while the remaining 80% receive a total attention score below 0.1. This implies that large language models do not need to utilize all the key-value cache during inference, but only need to store the portion that continuously plays a crucial role throughout the process. Then, through theoretical analysis and experiments, it is demonstrated that similar queries exhibit similar attention patterns to each key during attention computation. Specifically, for two similar queries (cosine similarity exceeding 0.8), a key important to one query is also important to the other. Therefore, during the inference process, the keys that each previous query considered important can be recorded, the most similar query can be found using the current query, and attention calculation can be performed using the keys that it considered important and their corresponding values, thereby avoiding calling the entire cache for attention calculation.

[0031] The above method has a significant drawback: it requires additional storage of attention information for each query to preserve the keys it deems important, and the introduction of similarity calculation operations adds extra computational complexity. Experiments have shown that the similarity between neighboring queries in the above model is very high; for example, the current query has a high cosine similarity to the previous r queries. Therefore, during the inference phase, attention information from recent queries can be recorded only to determine the important key-value pairs to retain, while discarding relatively less important key-value pairs, without requiring additional similarity calculation operations. This significantly reduces the GPU memory occupied by the KV cache in large language models during the inference phase.

[0032] This invention proposes a method for optimizing KV Cache based on attention information of recent queries in the inference step. This method records the attention information of queries in recent inference steps, judges the importance of information in KV Cache, deletes relatively unimportant KV Cache and retains important cached parts, and ultimately optimizes more than 70% of KV Cache while keeping the model performance unaffected.

[0033] Current mainstream large-scale language models mainly adopt the Transformer structure, in which attention computation plays a crucial role. During the t-th inference step, q... t Represents the query vector, k i (1≤i≤t) represents the key vector at step i, where q t and k i Both are d-dimensional vectors, q t The attention score assigned to the i-th key vector can be represented as:

[0034]

[0035] This invention uses the current average attention score as a threshold, and calculates it as follows: The attention score in the attention calculation exceeds the threshold (i.e.) Key-value pairs are considered important, and those that do not exceed a threshold (i.e., ...) Key-value pairs of unimportant key-value pairs are considered unimportant. This invention uses the proportion of unimportant key-value pairs to represent the sparsity of attention; a higher proportion indicates higher sparsity. For example... Figure 1 , 2 Figure 3 illustrates the sparsity of each attention layer and some attention heads during inference in the mainstream open-source large-scale language model LLaMA2-7B. The figure shows that the attention matrix of the large-scale language model exhibits very high sparsity during attention calculation. The sparsity of attention layers with larger index values ​​(deeper layers) even exceeds 90%, and different attention heads within the same layer also show high sparsity.

[0036] Similar queries show similar levels of focus on different keys, such as Figure 4 The image shows attention maps of the same text after passing through LLaMA2-7B at the 0th attention head of two different layers. Dark green indicates high attention level, and the background color indicates low attention level. The horizontal axis represents the position of the key in the sequence, and the vertical axis represents the position of the query in the sequence. Each row in the graph represents the attention of the query at the corresponding position to all previous keys. The vertical bar on the right represents the similarity between the query at the corresponding position and the query at position 255. This invention discretizes the attention scores, displaying important keys in bright green. Each attention map contains two red borders: the bottom border represents the important keys actually attended by the current query, and the other border represents the important keys attended by the query most similar to the current query.

[0037] Query q i and q j (where j>i) The formula for calculating cosine similarity is as follows:

[0038]

[0039] If you query q i and q j For perfect similarity, we have cosine_similarity(q) i ,q j When ) = 1, we have:

[0040] q j =m·q i

[0041] Where m is a positive real number. Then for the first i key values ​​(k1, k2, ..., k... i ), q j The attention weights assigned to each key-value pair are: q i The attention weights assigned to each key-value pair are: At this point:

[0042]

[0043] ...

[0044]

[0045] That is (q) j ·k1,q j ·k2,…,q j ·k i )=m(q i ·k1,q i ·k2,…,q i ·k i This means that for q i Important key values ​​for q j Equally important.

[0046] This experimentally demonstrates that similar queries exhibit similar levels of attention to keys. Therefore, during inference, we can record the keys considered important by each previous query, find the most similar query using the current query, and calculate attention using the keys it considered important. For example... Figure 5 As shown, by visualizing the cosine similarity of the query vector of a text of length 1024 on LLaMA2-7B, it was also found that the similarity of neighboring queries is very high in mainstream large-scale language models. Therefore, during the inference stage, attention information of recent queries can be recorded only to determine the important key-value pairs that need to be retained, while discarding the relatively less important key-value pairs.

[0047] like Figure 6The diagram illustrates the specific implementation process of this algorithm, with the main steps as follows:

[0048] Enter a text;

[0049] First, the text is segmented using a tokenizer to obtain the split tokens, which are then input into a large language model.

[0050] After the attention layer calculates the attention information, this invention uses a table to record the attention information of the most recent r queries. If the number of input tokens is less than r, the attention information of all queries is recorded. The average attention score is used as a threshold. Key-value pairs with an attention score less than the threshold are marked as 0, and the remaining key-value pairs are marked as 1. At the same time, the most recent r key-value pairs are marked as 1 to ensure that recent key-value pairs are not discarded prematurely.

[0051] Enter the loop reasoning phase to reason for the next token. After reasoning, add the attention information of the current query to the table. The r most recent key-value pairs are marked as 1 to prevent recent key-value pairs from being eliminated prematurely, thus affecting reasoning performance. If the number of queries in the table exceeds r, only the attention information of the r most recent queries is retained. If the number of queries is less than r, proceed directly to step 6.

[0052] like Figure 7 As shown, the attention information of each Query record is then ORed. Key-value pairs that are still marked as 0 after the operation (i.e., not considered important by any of the r Queries) are considered unimportant key-value pairs and are directly deleted from the KV Cache. The remaining key-value pairs are still stored in the KV Cache for the next step of inference generation.

[0053] If the preset reasoning length has been reached, proceed to step 7; otherwise, proceed to step 4 to generate the next token.

[0054] The generated token is decoded using a tokenizer to obtain the generated text.

[0055] In this scheme, r is a configurable parameter. A larger r value results in lower KV cache compression with minimal performance loss, while a smaller r value leads to higher KV cache compression but greater performance loss. A recommended range for r is 64–256, with the specific value determined based on the model and desired performance. This invention applies this algorithm to various attention heads in the attention layer of a large language model, ultimately maintaining the original model's performance while reducing KV cache by more than 70%.

[0056] The embodiments described above are merely illustrative of several implementations of the present invention, and while the descriptions are specific and detailed, they should not be construed as limiting the scope of the present invention. Those skilled in the art can make various modifications and improvements without departing from the concept of the present invention, and these modifications and improvements all fall within the scope of protection of the present invention.

Claims

1. A method for optimizing a large-scale language model KV cache based on recent query attention information, characterized in that, Includes the following steps: 1) Segment the text to be input into the selected large language model to obtain the segmented sub-words; 2) Input the split sub-words into a large language model and perform attention calculation to obtain the attention information of each query during the inference process of the large language model; Then, a record table of attention information for the most recent r queries is constructed using the average attention score of each key in the query as a threshold. 3) Enter the loop reasoning stage of the large language model to reason the next sub-word. After the reasoning is completed, add the attention information of the current query to the record table. If the number of queries in the table exceeds the preset number r, only the attention information of the most recent r queries will be retained. If the number of queries is less than r, proceed directly to step 5). 4) Perform an OR operation on the attention information of each recorded query to determine whether the key-value cache in the query is important, and delete the key-value cache that is deemed unimportant; 5) Determine whether the total number of generated subwords has reached the preset inference quantity. If not, return to step 3) to generate the next subword. If the quantity has been reached, decode the generated subword to obtain the generated text.

2. The large-scale language model KV cache optimization method according to claim 1, characterized in that, The attention calculation in step 2) is specifically as follows: Where t is the number of inference steps in the large language model, and q t For the query vector, k i Let q represent the key vector at step i, 1 ≤ i ≤ t. t and k i Both are d-dimensional vectors, α t,i For query vector q t Attention score assigned to the i-th key vector.

3. The large-scale language model KV cache optimization method according to claim 1, characterized in that, In step 2), when constructing a record table of attention information for the most recent r queries using the average attention score as a threshold, key-value pairs with attention scores less than the threshold are marked as 0, and the remaining key-value pairs are marked as 1. At the same time, the most recent r key-value pairs in each query are marked as 1 to ensure that recent key-value pairs are not discarded prematurely. If the number of input subwords is less than r, the attention information for all queries is recorded.

4. The large-scale language model KV cache optimization method according to claim 1, characterized in that, In step 3), the nearest r key-value pairs will be marked as 1.

5. The method for optimizing a large language model KV cache according to claim 1, characterized in that, The OR operation performed on the attention information of each query record in step 4) specifically involves: decision=m t-r+1 ∨m t-r+2 ∨…∨m t Where, decision is the result of an OR operation used to determine unimportant key-value caches and perform KV Cache optimization; m t-r+1 ,m t-r+2 ,…,m t This provides attention information for the corresponding r queries, where "∨" represents the logical OR operator.

6. The large-scale language model KV cache optimization method according to claim 5, characterized in that, Key-value pairs whose flag remains 0 after the OR operation are considered unimportant and deleted directly. The remaining key-value pairs are still saved and used for the next step of reasoning generation.

7. A system based on the large language model KV Cache optimization method of claim 1, characterized in that, include: A word segmenter is used to split the text to be input into a selected large language model into sub-words and to decode the sub-words generated by the large language model to obtain the generated text. The attention calculation module is used to perform attention calculations on subwords input into a large language model. The record table construction module is used to construct a record table of attention information for the most recent r queries, using the average attention score as a threshold. Attention information or operation module, used to perform OR operation on the attention information of each query in the record, and delete key-value cache that is deemed unimportant; The sub-word count determination module is used to determine whether the total number of generated sub-words has reached the preset number.

8. An electronic device, characterized in that, The system includes a processor and a memory, the memory storing machine-executable instructions that can be executed by the processor, the processor executing the machine-executable instructions to implement the large language model KV Cache optimization method according to any one of claims 1-6.

9. A machine-readable storage medium, characterized in that, The machine-readable storage medium stores machine-executable instructions, which, when called and executed by the processor, are used to implement the large language model KV Cache optimization method according to any one of claims 1-6.

Images