Large-scale language model KV Cache optimization method based on recent query attention information

Abstract

The invention discloses a large-scale language model KV Cache optimization method based on recent query attention information. According to the method, attention calculation is carried out on sub-words input into a large-scale language model; constructing a record table of attention information of recent r queries by taking the average attention score as a threshold value; performing OR operation on the recorded attention information of each query, and deleting the key-value cache which is determined to be unimportant; judging whether the number of all generated sub-words reaches a preset number or not; according to the method, the importance degree of information in the KV Cache is judged by recording the attention information of the Query in the recent reasoning step, the relatively less important KV Cache is deleted, the important cache part is reserved, and finally, the performance of the model can be kept not to be obviously influenced under the condition that more than 70% of the KV Cache can be optimized.

Description

Technical Field

[0001] The present invention relates to the field of large language models, and in particular to a large language model KV Cache optimization method based on recent query attention information. Background Art

[0002] Large language models have performed well in various natural language processing (NLP) tasks, including question answering, text summarization, and multi-round dialogues. However, current large language models have the characteristics of large number of parameters and the memory usage of the key-value cache (KV Cache) during inference increases linearly with the length of the sequence. This causes the model to occupy a large amount of memory when processing long texts during the inference phase and faces the problem of memory overflow.

[0003] At present, large language models are subject to the risk of memory overflow during inference due to limited computing resources. The memory overhead mainly comes from two aspects: one is the overhead caused by the model parameters themselves, and the other is the overhead caused by the linear growth of the memory occupied by KV Cache as the length of the inference sequence increases. Therefore, a solution is needed to optimize KV Cache to reduce the memory overhead it brings. Summary of the invention

[0004] In order to solve the problem of video memory overflow faced by large language models during the inference process, the present invention mainly focuses on optimizing the video memory overhead caused by KV Cache. By utilizing the attention information of recent queries (Query) during the inference process, KVCache is optimized to reduce the video memory overhead it brings; finally, a method for optimizing the video memory occupancy of key-value cache during the inference process of large language models is proposed.

[0005] First, the present invention proposes a large language model KV Cache optimization method 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 split subwords;

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

[0008] 3) Enter the loop reasoning stage of the large language model to infer the next subword. After the reasoning is completed, the attention information of the current query is added 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 is retained. If the number of queries is less than r, directly enter step 5);

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

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

[0011] The present invention also provides a system based on the above large language model KV Cache optimization method, comprising:

[0012] A tokenizer, which is used to split the text to be input into the selected large language model into subwords and to decode the subwords generated by the large language model to obtain the generated text;

[0013] The attention calculation module is used to perform attention calculation on subwords input into the large language model;

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

[0015] An attention information OR operation module, used to perform an OR operation on the attention information of each query recorded, and delete the key-value cache that is determined to be unimportant;

[0016] The subword quantity determination module is used to determine whether the number of all generated subwords has reached a 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-mentioned large language model KV Cache optimization method.

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

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

[0020] The present invention is inspired by theoretical analysis and experimental observation that similar queries pay similar attention to key values. At the same time, the attention matrix of a large language model is very sparse, which ensures that a large amount of KVCache can be optimized in the inference stage.

[0021] The present invention proposes to record the attention information of recent queries to realize the importance judgment of key-value pairs, and decide the retention and deletion of key-value pairs based on this information. In the reasoning stage, only the attention information of recent queries can be recorded to determine the important key-value pairs that need to be retained, and relatively less important key-value pairs can be discarded without the need for additional similarity calculation operations, thereby greatly reducing the video memory space occupied by the KV Cache of large language models in the reasoning stage.

[0022] The optimization method of the present invention does not require additional training operations, and can also maintain the reasoning effect of the original model while optimizing more than 70% of the KV Cache.

[0023] The method proposed in the present invention is not only applicable to large language models using the Transformer architecture, but also to generative models such as visual language models using the Transformer architecture. BRIEF DESCRIPTION OF THE DRAWINGS

[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 The attention graphs of the 0th attention heads of two different layers after the same text passes through LLaMA2-7B in an embodiment of the present invention;

[0026] Figure 5 A visualization diagram of the cosine similarity of a query vector of a text with a length of 1024 on LaMA2-7B in an embodiment of the present invention;

[0027] Figure 6 It is the algorithm flow chart of the present invention;

[0028] Figure 7 The present invention records the in-table OR operation and the deletion flow chart of the key-value pair. DETAILED DESCRIPTION

[0029] The present invention is further described and illustrated below in conjunction with specific embodiments. The embodiments are merely exemplary of the present disclosure and do not define the scope of limitation. The technical features of each embodiment of the present invention may be combined accordingly without conflicting with each other.

[0030] In order to solve the problem of video memory overflow faced by large language models during the reasoning process, the present invention first analyzes the attention sparsity of large language models such as LLaMA, Mistral, and Qwen when calculating attention, that is, the key-value pairs that are really important during the attention calculation in the reasoning process only account for a small part, while other key-value pairs do not play a significant role in the attention calculation process. Specifically, in most cases, the sum of attention scores assigned to 20% of the key-value pairs exceeds 0.9, while the sum of attention scores assigned to the remaining 80% of the key-value pairs is less than 0.1, which means that the large language model does not need to use all KV Caches during the reasoning process, but only needs to save the part that continues to play a role in the entire reasoning process. Then, through theoretical analysis and experiments, it is proved that similar queries pay similar attention to each key (key) when performing attention calculation. Specifically, for two similar queries (cosine similarity exceeds 0.8), the key that is more important to one query is also more important to the other query. Therefore, during the reasoning process, we can record the keys that each previous query considers important, use the current query to find the most similar query, and use the keys that it considers important and their corresponding values ​​to perform attention calculations, thereby avoiding calling the entire cache for attention calculations.

[0031] The above method has a significant drawback, that is, it is necessary to store the attention information of each query in order to save the keys that it considers important, and introduces similarity calculation operations, which increases the additional computational complexity. Experiments have found that the similarity of adjacent queries in the above model is very high. For example, the cosine similarity between the current query and the previous r adjacent queries is high. Therefore, in 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, and relatively unimportant key-value pairs can be discarded without additional similarity calculation operations, thereby greatly reducing the video memory space occupied by the KV Cache of large language models in the inference stage.

[0032] The present invention proposes a method for optimizing KV Cache based on the attention information of recent queries in the reasoning step. This method records the attention information of queries in the recent reasoning step, judges the importance of information in the KV Cache, deletes relatively unimportant KV Cache and retains important cache parts. Ultimately, while optimizing more than 70% of the KV Cache, the performance of the model can be kept unaffected.

[0033] The current mainstream large-scale language models mainly use the Transformer structure, in which attention calculation plays an important role. t represents the query vector, k i (1≤i≤t) represents the key vector of the i-th step, where q t and k i are both d-dimensional vectors, q t The attention score assigned to the i-th key vector can be expressed as:

[0034]

[0035] The present invention uses the current average attention score as the threshold, calculated as The attention score in the attention calculation exceeds the threshold (i.e. ) is considered important and does not exceed the threshold (i.e. ) is considered unimportant. The present invention uses the proportion of unimportant key-value pairs to indicate the sparsity of attention. The higher the proportion of unimportant key-value pairs, the higher the sparsity. Figure 1 , 2 As shown in Figure 3, the sparsity of each attention layer and some attention heads of the mainstream open source large language model LLaMA2-7B during inference is shown. From the figure, it can be seen that the sparsity of the attention matrix of the large language model is very high when performing attention calculation. The sparsity of the attention layer with a large index value (deep layer) even exceeds 90%, and different attention heads in the same layer also show high sparsity.

[0036] Similar queries pay similar attention to different keys, such as Figure 4 As shown in the figure, it is the attention map of the 0th attention head of two different layers after the same text passes through LLaMA2-7B. Dark green indicates a high degree of attention, and the background color indicates a low degree of attention. 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 figure represents the attention of the Query at the corresponding position to all previous Keys, and the vertical bar on the right represents the similarity between the Query at the corresponding position and the Query at the current position 255. The present invention discretizes the attention scores, and those important Keys are displayed in bright green. Each attention map contains two red borders, the bottom border represents the important keys that the current Query actually pays attention to, and the other border represents the important keys that the Query most similar to the current Query pays attention to.

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

[0038]

[0039] If the query q i and q j If they are completely similar, then cosine_similarity(q i ,q j )=1, then 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 weight assigned to each key-value pair is q i The attention weight assigned to each key-value pair is At this time there are:

[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 ), which means that for q i Important key values ​​for q j Equally important.

[0046] This also proves experimentally that similar queries pay similar attention to keys. So during the reasoning process, we can record the keys that each previous query considers important, use the current query to find the most similar query, and use the keys that it considers important to perform attention calculation. Figure 5 As shown, by visualizing the cosine similarity of the query vector of a text of length 1024 on LLaMA2-7B, it is also found that in the mainstream large language model, the similarity of neighboring queries is very high. Therefore, in 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, and relatively less important key-value pairs can be discarded.

[0047] like Figure 6As shown in the figure, the specific implementation process of this algorithm is shown. The main process is as follows:

[0048] Enter a text;

[0049] First, use a tokenizer to segment the text into tokens, which are then input into a large language model.

[0050] After the attention calculation of the attention layer, the present invention uses a table to record the attention information of the 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 the threshold, and the 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 are marked as 1 to ensure that the recent key-value pairs will not be lost too early;

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

[0052] like Figure 7 As shown in the figure, an OR operation is performed on the attention information recorded by each Query. After the operation, the key-value pairs whose marks are still 0 (that is, not considered important by any of the r Queries) will be regarded as unimportant key-value pairs and directly deleted from the KV Cache. The remaining key-value pairs are still stored in the KV Cache for the next step of reasoning generation.

[0053] If the preset inference length has been reached, go to step 7, otherwise go to step 4 and loop to generate the next token;

[0054] Use the Tokenizer to decode the generated token to get the generated text;

[0055] In this scheme, r is a configurable parameter. The larger the setting, the lower the compression rate of KV Cache and the slight performance loss. The smaller the setting, the higher the compression rate of KV Cache and the greater the performance loss. It is recommended to set r in the range of 64 to 256, and the specific value needs to be set according to the model and the expected effect. The present invention applies the algorithm to each attention head of the attention layer in a large language model, and ultimately can maintain the performance of the original model while reducing more than 70% of KV Cache.

[0056] The above-mentioned embodiments only express several implementation methods of the present invention, and the description is relatively specific and detailed, but it cannot be understood as limiting the scope of the present invention. For ordinary technicians in this field, several modifications and improvements can be made without departing from the concept of the present invention, which all belong to the protection scope of the present invention.

Claims

1. A large language model KV Cache optimization method based on recent query attention information, characterized in that: The following steps are involved: 1) Segment the text to be input into the selected large language model to obtain the split subwords; 2) Input the split subwords into the large language model and perform attention calculation to obtain the attention information of each query during the large language model inference process; Then, the average attention score of each key in the query is used as the threshold to build a record table of attention information of the most recent r queries; 3) Enter the loop reasoning stage of the large language model to infer the next subword. After the reasoning is completed, the attention information of the current query is added 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 is retained. If the number of queries is less than r, directly enter step 5); 4) Perform an OR operation on the attention information of each query recorded to determine whether the key-value cache in the query is important, and delete the key-value cache that is determined to be unimportant; 5) Determine whether the number of all generated subwords has reached the preset inference number. If not, return to step 3) to generate the next subword; if it has reached, decode the generated subword to obtain the generated text.

2. The large language model KV Cache optimization method according to claim 1, characterized in that: The attention calculation in step 2) is specifically as follows: Among them, t is the number of steps of large language model reasoning, q t is the query vector, k i represents the key vector of step i, 1≤i≤t, q t and k i are both d-dimensional vectors, α t,i is the query vector q t The attention score assigned to the i-th key vector.

3. The large language model KV Cache optimization method according to claim 1, characterized in that: In the step 2), when the record table of attention information of the recent r queries is constructed with the average attention score as the threshold, the 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 of all queries is recorded.

4. The large language model KV Cache optimization method according to claim 1, characterized in that: The most recent r key-value pairs in step 3) will be marked as 1.

5. The large language model KV Cache optimization method according to claim 1, characterized in that: The specific OR operation performed on the attention information of each query record in step 4) is: decision=m t-r+1 ∨m t-r+2 ∨…∨m t Wherein, decision is the result of the OR operation for the key-value cache that is not important for decision making and KV Cache optimization; m t-r+1 ,m t-r+2 ,…,m t is the attention information of the corresponding r queries, and "∨" represents the logical OR operator.

6. The large language model KV Cache optimization method according to claim 5, characterized in that: After the OR operation, the key-value pairs whose marks are still 0 will be regarded as unimportant key-value pairs and deleted directly, and the remaining key-value pairs will still be saved and used for the next step of reasoning generation.

7. A system based on the large language model KV Cache optimization method according to claim 1, characterized in that: include: A tokenizer, which is used to split the text to be input into the selected large language model into subwords and to decode the subwords generated by the large language model to obtain the generated text; The attention calculation module is used to perform attention calculation on subwords input into the large language model; A record table construction module is used to construct a record table of attention information of the most recent r queries using the average attention score as a threshold; An attention information OR operation module, used to perform an OR operation on the attention information of each query recorded, and delete the key-value cache that is determined to be unimportant; The subword quantity determination module is used to determine whether the number of all generated subwords has reached a preset number.

8. An electronic device, characterized in that: It includes 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 large language model KV Cache optimization method according to any one of claims 1 to 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 a processor, are used to implement the large language model KV Cache optimization method according to any one of claims 1 to 6.

Images