Method and apparatus for optimizing inference cache for generative language model

Abstract

Embodiments of the present application relate to the technical field of natural language processing, and provide a method and apparatus for optimizing an inference cache for a generative language model, an electronic device, and a readable medium. The method comprises: acquiring a kv cache for a generative language model during the nth inference process; segmenting a second kv cache on the basis of predicted tokens recording punctuation marks; calculating an attention weight of each predicted token; on the basis of the attention weights, performing segmented compression on each segment of the second kv cache according to different lengths of cache segments; combining a first kv cache with the compressed second kv cache to obtain an updated kv cache; and using the updated kv cache to reset the kv cache for the generative language model during the (n+1)th inference process. The present application effectively reduces the video memory usage of inference caches while ensuring inference accuracy, thereby increasing the model inference speed, and avoiding the problem of out-of-video memory.

Description

A method and apparatus for optimizing inference caching in generative language models

[0001] Cross-reference to related applications

[0002] This application claims priority to Chinese Patent Application No. 202411782521.7, filed on December 5, 2025, entitled "A Method and Apparatus for Optimizing Inference Caching of Generative Language Models", the entire contents of which are incorporated herein by reference. Technical Field

[0003] This application relates to the field of natural language processing technology, and in particular to a method for optimizing the inference cache of a generative language model, a device for optimizing the inference cache of a generative language model, an electronic device, and a computer-readable medium. Background Technology

[0004] Key-value caching, or kv caching, is an important engineering technique for optimizing Transformer inference performance. By trading space for time, it caches the key and value vectors calculated by the model. When the model performs autoregressive calculations next time, it reuses the kv cache from the previous inference, avoiding redundant calculations and thus improving inference performance.

[0005] Generative language models, especially the currently popular large language models, have extremely high requirements for GPU memory. On the one hand, the number of parameters in the models themselves is increasing, consuming more and more GPU memory. On the other hand, to improve inference speed, models commonly use key-value caching technology. Key-value caching is an important engineering technique for optimizing Transformer inference performance. It uses the idea of ​​trading space for time to cache the key and value vectors calculated by the model. The model reuses the key-value cache from the previous inference in the next autoregression calculation, avoiding redundant calculations. As the text to be generated becomes longer and longer, the key-value cache data used by the model during inference is also increasing. How to reduce the GPU memory consumption during the inference process of generative language models to improve model inference performance is a challenge in the industry.

[0006] However, as the number of samples and the length of the output sequence increase in the parallel generation strategy, the key-value cache of the existing generative language model consumes too much GPU memory and the GPU memory consumption increases too rapidly, even exceeding the memory consumption of the generative language model's own parameters. This easily leads to problems such as GPU memory exhaustion and excessive inference latency, reducing the user experience. Summary of the Invention

[0007] This application provides a method, apparatus, electronic device, and computer non-volatile readable storage medium for optimizing the inference cache of a generative language model. The aim is to address the problem in the prior art where, with the increase in the number of samples and the length of the output sequence in the parallel generation strategy, the key-value cache of existing generative language models occupies excessive amounts of video memory, and the rate of increase in video memory usage is too rapid, even exceeding the memory usage of the generative language model's own parameters. This easily leads to video memory exhaustion and excessively high inference latency, thus reducing the user experience.

[0008] This application discloses, in some embodiments, a method for optimizing inference caching in generative language models, including:

[0009] S1: Obtain the kv cache of the generative language model in the nth inference process, where the kv cache includes the first kv cache of the input text and the second kv cache in the process of generating each prediction token;

[0010] S2: Segment the second key-value cache according to the predicted token that records punctuation marks;

[0011] S3: Calculate the attention weight for each predicted token;

[0012] S4: Combining attention weights, the second key-value cache of each segment is segmented and compressed according to the different lengths of the cache segments;

[0013] S5: Combine the first kv cache with the compressed second kv cache to obtain the updated kv cache;

[0014] S6: Reset the kv cache of the generative language model in the (n+1)th inference process by updating the kv cache.

[0015] Some embodiments of this application disclose an inference caching optimization device for generative language models, including:

[0016] The acquisition module is used to acquire the key-value cache of the generative language model in the nth inference process, wherein the key-value cache includes the first key-value cache of the input text and the second key-value cache in the process of generating each prediction token;

[0017] The segmentation module is used to segment the second key-value cache according to the predicted token that records punctuation marks;

[0018] The calculation module is used to calculate the attention weight of each predicted token;

[0019] The compression module is used to combine attention weights to segment and compress each segment of the second key-value cache according to the different lengths of the cache segments.

[0020] The combination module combines the first key-value cache with the compressed second key-value cache to obtain the updated key-value cache;

[0021] The reset module is used to reset the kv cache of the generative language model during the (n+1)th inference process by updating the kv cache.

[0022] Some embodiments of this application also disclose an electronic device, including a processor, a communication interface, a memory, and a communication bus, wherein the processor, the communication interface, and the memory communicate with each other through the communication bus;

[0023] Memory, used to store computer programs;

[0024] When a processor executes a program stored in memory, it implements an inference cache optimization method for a generative language model as described in the embodiments of this application.

[0025] Some embodiments of this application also disclose one or more computer-readable media having instructions stored thereon that, when executed by one or more processors, cause the processors to perform an inference cache optimization method for a generative language model as described in the embodiments of this application.

[0026] The embodiments of this application have the following advantages:

[0027] In the generative language model inference process, this application introduces a key-value (kV) cache segmentation mechanism. Key information is selected from each kV cache segment based on attention weights and relative positions, and compression is performed at different ratios. Throughout the process, the first kV cache of the input text is fully preserved, while some non-critical kV cache vectors in the second kV cache are discarded, thereby reducing the impact on the model's output accuracy. This effectively avoids excessive memory usage by the kV cache, preventing memory exhaustion. While ensuring the preservation of key inference information and maintaining model inference accuracy, it also effectively improves model inference speed. Attached Figure Description

[0028] Figure 1 is a flowchart of the steps of an inference caching optimization method for a generative language model provided in some embodiments of this application;

[0029] Figure 2 is a schematic diagram of a model inference process based on key-value caching technology provided in some embodiments of this application;

[0030] Figure 3 is a schematic diagram of a kv cache compression process provided in some embodiments of this application;

[0031] Figure 4 is a schematic diagram of the calculation process of attention weights for each predicted token provided in some embodiments of this application;

[0032] Figure 5 is a schematic diagram of the structure of an inference cache optimization device for a generative language model provided in some embodiments of this application;

[0033] Figure 6 is a block diagram of an electronic device provided in some embodiments of this application;

[0034] Figure 7 is a schematic diagram of a computer-readable medium provided in some embodiments of this application. Detailed Implementation

[0035] To make the above-mentioned objectives, features and advantages of this application more apparent and understandable, the application will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0036] Referring to Figure 2, a schematic diagram of a model inference process based on key-value caching technology provided in some embodiments of this application is shown. After the disaster recovery system solution is confirmed and implemented, fault drills need to be conducted to verify the effectiveness and robustness of the disaster recovery system solution. After repeated fault injections, fault problems are located, thereby continuously improving and strengthening the disaster recovery system.

[0037] Figure 2 illustrates the flowchart of model inference based on key-value caching technology. The processing flow is as follows:

[0038] In the first stage, based on the user-input token, each transformer layer calculates and stores a key cache and a value cache corresponding to the length of the token. After the first forward pass, the model generates the first token. In the second stage, based on the previously generated token, the model calculates its key vector and value vector, appends them to the key cache and value cache, and performs self-attention calculation based on the updated key-value cache. After the model completes the forward pass, it generates the next token.

[0039] For example, in Figure 2, in the first stage, upon receiving the initial tokens input by the user, the model processes these tokens, calculates key-value vectors and stores them in a cache, and simultaneously generates the first output token. The input token embeddings undergo a linear transformation to generate a Query matrix of shape 4×64, where 4 represents the number of initial input tokens and 64 represents the hidden dimension. The input token embeddings undergo a linear transformation to generate a 4×64 Key matrix, representing the feature representation of each token. The Key matrix is ​​transposed and stored in the Key cache to support subsequent attention calculations. The input token embeddings undergo a linear transformation to generate a 4×64 Value matrix, representing the specific information of each token. The Value matrix is ​​stored in the Value cache. The Query matrix and the Key cache matrix are multiplied by a dot product to obtain the attention score. The attention score is then processed by Softmax to obtain weights, which are multiplied by the Value cache matrix to generate a 4×64 result matrix. After further decoding or linear mapping, the result is mapped to a vocabulary probability distribution, ultimately generating the first output token. In the second stage, after generating the first token, the model progressively updates the key-value cache based on the previously generated tokens, performing autoregressive generation. Based on the previously generated token, calculate the corresponding Query matrix, which is currently 1×64 in shape because there is only one token. Based on the current token, calculate the Key matrix (1×64) and append it to the cached Key matrix, making the Key cache shape 5×64. Based on the current token, calculate the Value matrix (1×64) and append it to the cached Value matrix, making the Value cache shape 5×64. Perform a dot product between the updated Query matrix (1×64) and the Key cache matrix (5×64) to generate an attention score. Multiply the attention score by the updated Value cache matrix (5×64) to generate a 1×64 result matrix. The result is decoded or linearly mapped to generate the probability distribution of the next token, and the token is output. Repeat the above process to generate subsequent tokens step by step, while updating the Key and Value caches. The introduction of KV caching avoids repeatedly calculating the KV matrix of historical tokens, greatly improving inference efficiency.

[0040] Referring to Figure 3, a schematic diagram of a kv cache compression process provided in some embodiments of this application is shown.

[0041] In Figure 3, the generated tokens are segmented according to punctuation marks, and each segment is compressed to varying degrees. The purpose is to optimize key-value caching. In the generative language model inference process, by introducing a key-value caching segmentation mechanism and selecting key information for each key-value caching segment, different compression ratios are set for compression, thereby reducing the key-value caching of the entire model inference and ultimately improving the model inference performance.

[0042] Referring to Figure 1, a flowchart illustrating the steps of an inference cache optimization method for a generative language model provided in some embodiments of this application is shown. The inference cache optimization method for a generative language model may specifically include the following steps:

[0043] S1: Obtain the kv cache of the generative language model during the nth inference process.

[0044] The key-value cache includes a first key-value cache for the input text and a second key-value cache for the process of generating each predicted token.

[0045] It's important to note that the key-value (KV) caches used by the generative language model during the nth inference iteration are divided into two parts: a first KV cache and a second KV cache. The first KV cache is a key-value vector calculated for each token in the input text, representing the initial context information. The second KV cache is a key-value vector that is progressively updated during generation, recording the KV pairs corresponding to each generated predicted token. These two caches together constitute the historical information that the model needs to refer to during inference, ensuring that the generated tokens can dynamically adjust the output based on the complete context.

[0046] In some embodiments of this application, S1 specifically includes:

[0047] S101: Get the input text.

[0048] S102: Calculate attention triples for each input token in the input text through each transformer layer in the generative language model. The attention triples include a query vector, a key vector, and a value vector.

[0049] S103: Cache each key vector and its corresponding value vector to obtain the first key-value cache.

[0050] S104: Input the input text, query vector, and first key-value cache into the generative language model and output the first predicted token.

[0051] S105: Compute the attention triplet of the first predicted token for each transformer layer in the generative language model.

[0052] S106: Add the first predicted token and the corresponding attention triple to the input text to update the input text.

[0053] S107: Re-input the updated input text into the generative language model and output the second predicted token.

[0054] S108: Compute the attention triplet of the second predicted token for each transformer layer in the generative language model.

[0055] S109: Use the second predicted token as the first predicted token, return to step S106, until the generated predicted token is the terminator token.

[0056] S111: Cache the key and value vectors in the attention triplet of the first predicted token and the key and value vectors in the attention triplet of the second predicted token to obtain the second key-value cache.

[0057] It's important to note that this process describes the detailed workflow of a generative language model acquiring a key-value (KV) cache. First, the input text is acquired, and then, in each Transformer layer of the model, an attention triple (containing a query vector, key vector, and value vector) is calculated for each token in the input text. These key-value vectors are then cached to obtain the initial first KV cache. Next, the input text, query vector, and the first KV cache are input into the model to generate the first predicted token, and the attention triple for this token is calculated. Subsequently, this predicted token and its attention triple are added to the input text to update the text. The updated text is then input into the model again to generate the next predicted token. By calculating the attention triple for the new predicted token and using it as the current predicted token, this process is repeated until a terminator token is generated. Finally, the key-value vectors of all generated tokens are cached to form the second KV cache, recording the dynamic context information generated during the generation process. This process ensures that the model can progressively update and utilize the KV cache during inference, efficiently generating the output sequence.

[0058] S2: Segment the second key-value cache according to the predicted token of the punctuation mark.

[0059] It's important to note that the second key-value cache generated during the generation process is segmented according to the predicted tokens corresponding to punctuation marks. This means that by identifying punctuation marks in the generated sequence, the key-value cache is divided into multiple segments, each representing a context structure. Segmentation provides a foundation for subsequent importance-based compression and cache optimization, while also reducing unnecessary data redundancy.

[0060] In some embodiments of this application, S2 specifically includes:

[0061] S201: Extract the predicted token representing the punctuation mark from the predicted token.

[0062] S202: Segment the second key-value cache according to the predicted token representing the punctuation mark.

[0063] S203: Determine whether the cache length of the second kv cache of each segment is less than the preset cache length. If the cache length of the second kv cache of each segment is less than the preset cache length, merge the second kv cache segments with a segment number of m that are less than the preset cache length with the second kv cache segments with a segment number of m-1 and update the segment number. Otherwise, proceed to the next step.

[0064] S204: Output the second kv buffer for each segment after segmentation.

[0065] Specifically, this process describes the steps for segmenting the second KV cache. First, tokens representing punctuation marks are extracted from the generated prediction tokens; these punctuation marks serve as the basis for segmentation. Next, the second KV cache is divided into multiple segments according to these punctuation marks, each segment corresponding to a portion of the generated sequence. In step S203, it is determined whether the cache length of each segment is less than a preset cache length standard. If some segments are insufficient in length, these shorter segments are merged with the adjacent previous segment, and the number of segments is updated. Finally, the adjusted segmented structure of the second KV cache is output. This segmentation method provides a good foundation for subsequent KV cache compression, while reducing unnecessary redundant segments and improving cache management efficiency.

[0066] Referring to Figure 4, a schematic diagram of the calculation process of attention weights for various predicted tokens provided in some embodiments of this application is shown.

[0067] In Figure 4, each token generated by the model is first assigned a corresponding weight value, initially set to 0. For each input token, the model calculates an attention score for each preceding token using self-attention. This attention score is then accumulated into the corresponding weight value. As shown in Figure 4, during the autoregressive output of tokens t1 to t7, a new attention score is added to the weight value of each output token after each iteration. Ultimately, each vector in each key-value cache has a corresponding weight value. The current weight value for each token at the bottom of Figure 4 is the total weight value obtained by accumulating the weight values ​​after multiple iterations. At this point, a larger weight value indicates that the information from that vector is more critical.

[0068] S3: Calculate the attention weights for each predicted token.

[0069] Attention weight is a key concept in the self-attention mechanism, representing the degree of attention or importance the model pays to each position in the sequence when processing sequential data. It generates a probability distribution by measuring the similarity between the query and the key, which is used to determine the weighting ratio of the values. The attention weight is calculated for each predicted token. Through the model's self-attention mechanism, the relevance of the currently generated token to all previously generated tokens is calculated, resulting in attention scores. These scores are then converted into attention weights using Softmax. These weights reflect the importance of each predicted token to the generated sequence, providing a basis for subsequent key-value caching and compression.

[0070] In some embodiments of this application, S3 specifically includes:

[0071] S301: Initialize the initial attention weights for each existing prediction token.

[0072] Specifically, it can be initialized to 0.

[0073] S302: In response to the instruction to generate a new predicted token, calculate the attention weights of each existing predicted token through self-attention.

[0074] S303: Update the corresponding initial attention weights by calculating the attention weights of each existing predicted token.

[0075] In some embodiments of this application, S303 specifically refers to:

[0076] The calculated attention weights of each existing predicted token are summed with the initial attention weights to update the initial attention weights of each existing predicted token.

[0077] Understandably, the current attention weights of each existing predicted token, calculated through a self-attention mechanism, are accumulated with their initial attention weights. This accumulation process continuously updates the weight values ​​of each predicted token, dynamically reflecting their importance accumulated during the generation process.

[0078] S304: Use the updated initial attention weights as the attention weights for each predicted token.

[0079] It's important to note that this process details the steps for calculating the attention weights of each predicted token. First, the initial attention weights of all generated predicted tokens are initialized to 0. Next, each time a new predicted token is generated, its attention weight relative to all generated tokens is calculated using a self-attention mechanism, representing its relevance to the context. Then, the calculated attention weights are added to the corresponding token's initial weights, updating the token's status. Finally, the updated weights are used as the final attention weights for each predicted token. This process dynamically accumulates the importance of each predicted token, providing precise weight information for subsequent caching optimizations.

[0080] S4: Combining attention weights, the second key-value cache of each segment is segmented and compressed according to the different lengths of the cache segments.

[0081] Specifically, based on the attention weights of each token in each second KV cache segment, the cache is segmented and compressed according to its segment length. Specifically, by sorting the attention weights, key-value pairs with higher weights are retained first, while less important parts are discarded according to a preset compression ratio, thereby reducing the storage space occupied by the KV cache. This segmented compression strategy ensures that while reducing GPU memory requirements, it retains as much information as possible that is most important to the generation task.

[0082] In some embodiments of this application, S4 specifically includes:

[0083] S401: Determine the compression ratio of the second kv buffer for each segment.

[0084] In some embodiments of this application, the compression ratio is determined as follows:

[0085] Where o represents the compression ratio, d represents the distance between the position of the last token in the current cache segment of the second kv cache and the position of the first token in the last kv cache segment of the current second kv cache, B and E represent the first preset distance and the second preset distance, respectively, and p represents the minimum compression ratio.

[0086] More specifically, B is the first preset distance. If the distance between the last token in the cache segment whose compression ratio is to be calculated and the first token in the last cache segment of the current second key-value cache is less than B, then the compression ratio of this cache segment is set to 1, meaning no compression is performed. This indicates that the cache segment is relatively close to the position of the token to be generated and does not need compression. E is the second preset distance. If the distance between the last token in the cache segment whose compression ratio is to be calculated and the first token in the last cache segment of the current second key-value cache is greater than E, then the compression ratio of this cache segment is set to p, meaning compression is performed at the minimum compression ratio. This indicates that the cache segment is relatively far from the position of the token to be generated and needs maximum compression. Therefore, the value of E is greater than the value of B, and the range of the value of d is between B and E, i.e., B≤d≤E. So, for cache segments between B and E, the calculated compression ratio values ​​will be between 1 and p, where 1>p.

[0087] It should be noted that this formula describes a dynamic calculation method for the KV cache compression ratio, where the compression ratio is calculated based on the distance between the current token's vector and the first vector of the second KV cache. By setting two preset distance ranges, the compression ratio changes linearly between its maximum and minimum values ​​as 'd' changes, ensuring that important contextual information is preserved with higher priority, while less important information far from the first segment is compressed at a higher ratio. The parameter 'd' defines the minimum compression ratio, ensuring that even at long distances, a small amount of necessary information is retained.

[0088] S402: Sort the attention weights in the second key-value cache of each segment.

[0089] S403: Retain the key vectors and corresponding value vectors in the second key-value cache after sorting them in descending order of attention weight. The proportion of retained vectors relative to all vectors in the second key-value cache is the compression ratio.

[0090] S404: Remove unretained vectors to perform segmented compression of each segment of the second kv cache according to the different cache segment lengths.

[0091] Specifically, this process describes the steps for segmented compression of the second KV cache. First, the compression ratio for each segment of the second KV cache is obtained to determine the range of vectors to be retained. Next, the attention weights in each cache segment are sorted to filter key-value vectors according to their importance. In step S403, key-value pairs (Key and Value vectors) that represent the sorted percentage of the compression ratio are retained according to the attention weights from largest to smallest, ensuring that high-weight information is retained first. Finally, vectors outside the retention range are removed, thus completing the segmented compression of each segment of the second KV cache. This approach effectively reduces the storage footprint of the KV cache while preserving the most critical contextual information in the generation task.

[0092] S5: Combine the first kv cache with the compressed second kv cache to obtain the updated kv cache.

[0093] S6: Reset the kv cache of the generative language model in the (n+1)th inference process by updating the kv cache.

[0094] Understandably, the bottleneck of generative language models, especially large language models, currently appears to be primarily limited by GPU memory. During inference, key-value (kV) caches consume a significant portion of GPU memory. As the batch size (the number of samples in a parallel generation strategy) and the length of the output sequence increase, the GPU memory overhead of the kV cache will rapidly increase, potentially exceeding the GPU memory occupied by the model's weights themselves. Therefore, optimizing the kV cache is crucial. This can be addressed by introducing a segmentation mechanism, dividing the output tokens into segments based on punctuation marks, thus segmenting the kV cache. Secondly, for each segment of the tokens' kV cache, key information is selected, discarding some non-critical kV cache vectors to compress the kV cache. Furthermore, based on different positional information, different compression ratios are applied to the kV cache of each segment of tokens to minimize the impact on the model's output accuracy.

[0095] The purpose of this application is to optimize key-value caching. In the process of generative language model inference, a key-value caching segmentation mechanism is introduced, and key information is selected for each key-value caching segment and compressed at different compression ratios, thereby reducing the key-value caching of the entire model inference (Figure 2), and ultimately improving the performance of model inference.

[0096] For example, the key-value (kV) cache acquisition process is as follows: Step 1: The user inputs m tokens. Each transformer layer in the model calculates m query vectors, m key vectors, and m value vectors for self-attention calculation. The m key vectors and m value vectors are saved to form a kV cache with a length of 2m, which is used for subsequent autoregressive generation calculations. After the entire model inference is complete, the first token is generated and output. Step 2: Using the generated first token as input, each transformer layer in the model calculates a corresponding query vector, a key vector, and a value vector. The key vector and value vector are saved to the kV cache respectively. At this time, the length of the kV cache becomes 2(m+1). Self-attention is calculated using the current query vector and the cached m+1 key vectors and m+1 value vectors. After the entire model inference is complete, the second token is generated and output. Step 3: Using the generated second token as input, Step 2 is repeated to generate the next token, while the kV cache length increases by 2. Step 4: The model continues to autoregressively generate the next token until a token representing the terminator is generated. The model generation is complete. If a total of n tokens are generated at this point, then the key-value cache length during this inference process is 2(m+n).

[0097] The following is an example of the implementation of the key-value cache segmentation and key information selection mechanism:

[0098] During the autoregressive generation of each token, the position of the generated token representing the punctuation mark "," or "." is recorded. Since each generated token corresponds to each vector in the key-value cache, the key-value cache is segmented according to the recorded punctuation mark position. An adjustable parameter C is set to control the fine-grainedness of segmentation. When the length of a segment is less than C, the segment is merged with the previous segment. A corresponding weight value is introduced for each token generated by the model, initially set to 0. For each input token, the model calculates the attention score for each preceding token through self-attention, and adds the attention score of each token to its respective weight value. During the autoregressive output of tokens t1 to t7, a new attention score is added to the weight value of the output tokens for each iteration. Finally, each vector in each segment of the key-value cache has a corresponding weight value. At this point, the larger the weight value, the more critical the information of the vector. If we want to compress a certain key-value cache with a compression ratio of 50%, we can sort all the vector weights in the key-value cache from largest to smallest, keep only the vectors corresponding to the top 50% of the weight values, and discard the rest of the vectors in the cache, thus achieving the effect of key-value cache compression.

[0099] The compression ratio is determined as follows during segmented compression:

[0100] First, the key-value cache corresponding to the user-input tokens is relatively important, so this part of the cache will not be compressed. Second, the model's self-attention focuses on the most recently output token, so we will not compress the most recently saved key-value cache segments. We can set a distance threshold B (adjustable parameter). For each key-value cache segment, if the distance between the position of the last vector in that segment and the position of the first vector in the key-value cache segment containing the currently generated token is less than B, then that segment of the key-value cache will not be compressed. Then, we set a minimum compression ratio P (adjustable parameter) and a distance threshold E (adjustable parameter). For each key-value cache segment, if the distance between the position of the last vector in that segment and the position of the first vector in the key-value cache segment containing the currently generated token is greater than E, then that segment of the cache will be compressed according to the minimum compression ratio P. Finally, the key-value cache segments with a location distance between B and E are compressed. Following the principle that the greater the distance from the location of the currently generated token, the smaller the compression ratio, if the distance between the last vector in this key-value cache segment and the first vector in the key-value cache segment containing the currently generated token is d, then the compression ratio of this key-value cache segment is [value missing]. In summary, on the one hand, we segmented the key-value cache, retaining the key vector in each segment to ensure a more even distribution of useful information throughout the cache. On the other hand, we did not compress the key-value cache segments corresponding to user-input tokens and the most recently output tokens. Instead, we compressed key-value cache segments that were farthest from the current output token, with a smaller compression ratio for greater distances. We also set a minimum compression ratio to minimize the impact of key-value cache compression on model inference accuracy and maintain the stability of model output.

[0101] In practical applications, steps S1-S6 describe a key-value (KV) caching optimization method for generative language models. First, S1 acquires the KV cache used during inference, including a first KV cache of the initial input text and a second KV cache that is progressively updated during generation. Next, in S2, the second KV cache is segmented by identifying punctuation marks in the generated sequence, with each segment corresponding to a context structure. In S3, the model calculates the attention weight of each generated token using a self-attention mechanism, reflecting its importance to the generated sequence. Then, in S4, the attention weights and segment lengths in each KV cache segment are combined to sort and compress them, retaining only high-weight key-value pairs to reduce storage requirements. In S5, the compressed second KV cache is combined with the first KV cache to generate a new, updated KV cache. Finally, in S6, the updated KV cache is used for the next inference step to efficiently support the generation process, reducing GPU memory usage while ensuring generation quality.

[0102] The embodiments of this application have the following advantages:

[0103] In the generative language model inference process, this application introduces a key-value (kV) cache segmentation mechanism. Key information is selected from each kV cache segment based on attention weights and relative positions, and compression is performed at different ratios. Throughout the process, the first kV cache of the input text is fully preserved, while some non-critical kV cache vectors in the second kV cache are discarded, thereby reducing the impact on the model's output accuracy. This effectively avoids excessive memory usage by the kV cache, preventing memory exhaustion. While ensuring the preservation of key inference information and maintaining model inference accuracy, it also effectively improves model inference speed.

[0104] It should be noted that, for the sake of simplicity, the method embodiments are all described as a series of actions. However, those skilled in the art should understand that the embodiments of this application are not limited to the described order of actions, because according to the embodiments of this application, some steps can be performed in other orders or simultaneously. Secondly, those skilled in the art should also understand that the embodiments described in the specification are all embodiments of this application, and the actions involved are not necessarily required by the embodiments of this application.

[0105] Additionally, referring to FIG5, a block diagram of an inference cache optimization device for a generative language model provided in some embodiments of this application is shown.

[0106] Some embodiments of this application disclose an inference cache optimization device 20 for generative language models, including:

[0107] The acquisition module 201 is used to acquire the key-value cache of the generative language model in the nth inference process, wherein the key-value cache includes a first key-value cache of the input text and a second key-value cache in the process of generating each prediction token.

[0108] Segmentation module 202 is used to segment the second key-value cache according to the predicted token of the punctuation mark.

[0109] The calculation module 203 is used to calculate the attention weights of each predicted token.

[0110] Compression module 204 is used to combine attention weights to segment and compress each segment of the second kv cache according to the different lengths of the cache segments.

[0111] The combination module 205 combines the first kv cache with the compressed second kv cache to obtain the updated kv cache.

[0112] Reset module 206 is used to reset the kv cache of the generative language model in the (n+1)th inference process by updating the kv cache.

[0113] In some embodiments of this application, the acquisition module 201 is specifically used for:

[0114] Get the input text.

[0115] Attention triples for each input token in the input text are computed at each transformer layer in the generative language model. Each attention triple consists of a query vector, a key vector, and a value vector.

[0116] Each key vector and its corresponding value vector are cached to obtain the first key-value cache.

[0117] Input the input text, query vector, and first key-value cache into the generative language model, and output the first predicted token.

[0118] The attention triples for the first predicted token are computed for each transformer layer in the generative language model.

[0119] The first predicted token and its corresponding attention triple are added to the input text to update the input text.

[0120] The updated input text is re-inputted into the generative language model, which outputs a second predicted token.

[0121] The attention triples for the second predicted token are computed for each transformer layer in the generative language model.

[0122] Use the second predicted token as the first predicted token, and return to step S106 until the generated predicted token is the terminator token.

[0123] The key and value vectors from the attention triples of the first predicted token and the key and value vectors from the attention triples of the second predicted token are cached to obtain the second key-value cache.

[0124] In some embodiments of this application, the segmentation module 202 is specifically used for:

[0125] Extract the prediction tokens representing punctuation marks from the prediction tokens.

[0126] The second key-value cache is segmented based on the predicted token representing punctuation marks.

[0127] Determine whether the cache length of the second kv cache of each segment is less than the preset cache length. If the cache length of the second kv cache of each segment is less than the preset cache length, merge the first kv cache segment with a segment number of m that is less than the preset cache length with the second kv cache segment with a segment number of m-1 and update the segment number. Otherwise, proceed to the next step.

[0128] Output the second key-value buffer for each segment after segmentation.

[0129] In some embodiments of this application, the calculation module 203 is specifically used for:

[0130] Initialize the initial attention weights for each existing prediction token.

[0131] In response to the instruction to generate a new predicted token, the attention weights of each existing predicted token are calculated through self-attention.

[0132] The initial attention weights are updated by calculating the attention weights of each existing predicted token.

[0133] The updated initial attention weights are used as the attention weights for each predicted token.

[0134] In some embodiments of this application, the initial attention weights are updated by calculating the attention weights of each existing predicted token, specifically as follows:

[0135] The calculated attention weights of each existing predicted token are summed with the initial attention weights to update the initial attention weights of each existing predicted token.

[0136] In some embodiments of this application, the compression module 204 is specifically used for:

[0137] Determine the compression ratio of the second kv buffer for each segment.

[0138] Sort the attention weights in the second key-value cache of each segment.

[0139] The key vectors and their corresponding value vectors in the second key-value cache are retained in descending order of attention weight. The proportion of retained vectors relative to all vectors in the second key-value cache is the compression ratio.

[0140] Remove unretained vectors to perform segmented compression on each segment of the second kv cache according to the different cache segment lengths.

[0141] In some embodiments of this application, the compression ratio is determined as follows:

[0142] Where o represents the compression ratio, d represents the distance between the position of the last token in the current cache segment of the second kv cache and the position of the first token in the last kv cache segment of the current second kv cache, B and E represent the first preset distance and the second preset distance, respectively, and p represents the minimum compression ratio.

[0143] As the device embodiment is basically similar to the method embodiment, the description is relatively simple, and relevant parts can be found in the description of the method embodiment.

[0144] The embodiments of this application have the following advantages:

[0145] In the generative language model inference process, this application introduces a key-value (kV) cache segmentation mechanism. Key information is selected from each kV cache segment based on attention weights and relative positions, and compression is performed at different ratios. Throughout the process, the first kV cache of the input text is fully preserved, while some non-critical kV cache vectors in the second kV cache are discarded, thereby reducing the impact on the model's output accuracy. This effectively avoids excessive memory usage by the kV cache, preventing memory exhaustion. While ensuring the preservation of key inference information and maintaining model inference accuracy, it also effectively improves model inference speed.

[0146] Additionally, referring to FIG6, a block diagram of an electronic device provided in some embodiments of this application is shown. Some embodiments of this application also provide an electronic device including a processor 1301, a communication interface 1302, a memory 1303, and a communication bus 1304, wherein the processor 1301, the communication interface 1302, and the memory 1303 communicate with each other via the communication bus 1304.

[0147] Memory 1303 is used to store computer programs;

[0148] When processor 1301 executes a program stored in memory 1303, it performs the following steps:

[0149] S1: Obtain the kv cache of the generative language model in the nth inference process, where the kv cache includes the first kv cache of the input text and the second kv cache in the process of generating each prediction token;

[0150] S2: Segment the second key-value cache according to the predicted token that records punctuation marks;

[0151] S3: Calculate the attention weight for each predicted token;

[0152] S4: Combining attention weights, the second key-value cache of each segment is segmented and compressed according to the different lengths of the cache segments;

[0153] S5: Combine the first kv cache with the compressed second kv cache to obtain the updated kv cache;

[0154] S6: Reset the kv cache of the generative language model in the (n+1)th inference process by updating the kv cache.

[0155] The communication bus mentioned above can be a Peripheral Component Interconnect (PCI) bus or an Extended Industry Standard Architecture (EISA) bus, etc. This communication bus can be divided into address bus, data bus, control bus, etc. For ease of illustration, only one thick line is used to represent it in the diagram, but this does not mean that there is only one bus or one type of bus.

[0156] The communication interface is used for communication between the aforementioned terminal and other devices.

[0157] The memory may include random access memory (RAM) or non-volatile memory, such as at least one disk storage device. In some embodiments, the memory may also be at least one storage device located remotely from the aforementioned processor.

[0158] The processors mentioned above can be general-purpose processors, including central processing units (CPUs), network processors (NPs), etc.; they can also be digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs), or other programmable logic devices, discrete gate or transistor logic devices, or discrete hardware components.

[0159] Referring to FIG7, a schematic diagram of a computer non-volatile readable medium provided in some embodiments of this application is shown. In some embodiments of this application, a computer non-volatile readable storage medium 1401 is also provided, which stores instructions that, when executed on a computer, cause the computer to perform the inference cache optimization method for generative language models in the above embodiments.

[0160] In some embodiments of this application, a computer program product containing instructions is also provided, which, when run on a computer, causes the computer to execute the inference caching optimization method for the generative language model described in the above embodiments.

[0161] In the above embodiments, implementation can be achieved, in whole or in part, through software, hardware, firmware, or any combination thereof. When implemented in software, it can be implemented, in whole or in part, as a computer program product. A computer program product includes one or more computer instructions. When the computer program instructions are loaded and executed on a computer, all or part of the flow or function according to the embodiments of this application is generated. The computer can be a general-purpose computer, a special-purpose computer, a computer network, or other programmable device. The computer instructions can be stored in a computer non-volatile readable storage medium or transmitted from one computer non-volatile readable storage medium to another. For example, computer instructions can be transmitted from one website, computer, server, or data center to another website, computer, server, or data center via wired (e.g., coaxial cable, fiber optic, digital subscriber line (DSL)) or wireless (e.g., infrared, wireless, microwave, etc.) means. The computer non-volatile readable storage medium can be any available medium that a computer can access or a data storage device such as a server or data center that integrates one or more available media. The available media can be magnetic media (e.g., floppy disks, hard disks, magnetic tapes), optical media (e.g., DVDs), or semiconductor media (e.g., solid state disks (SSDs)).

[0162] It should be noted that, in this document, relational terms such as "first" and "second" are used merely to distinguish one entity or operation from another, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Furthermore, the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Without further limitations, an element defined by the phrase "comprising one..." does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes the element.

[0163] The various embodiments in this specification are described in a related manner. Similar or identical parts between embodiments can be referred to mutually. Each embodiment focuses on describing the differences from other embodiments. In particular, the system embodiments are basically similar to the method embodiments, so the description is relatively simple; relevant parts can be referred to the descriptions of the method embodiments.

[0164] The above are merely some embodiments of this application and are not intended to limit the scope of protection of this application. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of this application are included within the scope of protection of this application.

Claims

1. A method for optimizing inference caching in generative language models, characterized in that, include: S1: Obtain the kv cache of the generative language model in the nth inference process, wherein the kv cache includes a first kv cache of the input text and a second kv cache in the process of generating each prediction token; S2: Segment the second kv cache according to the predicted token that records punctuation marks; S3: Calculate the attention weight for each predicted token; S4: Combining the attention weights, the second kv cache of each segment is segmented and compressed according to the different lengths of the cache segments; S5: Combine the first kv cache with the compressed second kv cache to obtain the updated kv cache; S6: Use the updated kv cache to reset the kv cache of the generative language model in the (n+1)th inference process.

2. The inference caching optimization method for generative language models according to claim 1, characterized in that, S1 specifically includes: S101: Obtain the input text; S102: Calculate the attention triplet for each input token in the input text through each transformer layer in the generative language model, wherein the attention triplet includes a query vector, a key vector, and a value vector; S103: Cache each key vector and its corresponding value vector to obtain the first key-value cache; S104: Input the input text, the query vector, and the first key-value cache into the generative language model, and output the first predicted token; S105: Calculate the attention triplet of the first predicted token through each transformer layer in the generative language model; S106: Add the first predicted token and the corresponding attention triple to the input text to update the input text; S107: Re-input the updated input text into the generative language model and output the second predicted token; S108: Calculate the attention triplet of the second predicted token through each transformer layer in the generative language model; S109: Use the second prediction token as the first prediction token, return to step S106, until the generated prediction token is a terminator token; S111: Cache the key vector and value vector in the attention triplet of the first predicted token and the key vector and value vector in the attention triplet of the second predicted token to obtain the second kv cache.

3. The inference caching optimization method for generative language models according to claim 1, characterized in that, S2 specifically includes: S201: Extract the prediction token representing the punctuation mark from the prediction token; S202: Segment the second kv cache according to the predicted token representing the punctuation mark; S203: Determine whether the cache length of each segment of the second kv cache is less than the preset cache length. If the cache length of each segment of the second kv cache is less than the preset cache length, merge the second kv cache segments with a number of segments m that are less than the preset cache length with the second kv cache segments with a number of segments m-1, and update the number of segments. Otherwise, proceed to the next step. S204: Output the second kv buffer for each segment after segmentation.

4. The inference caching optimization method for generative language models according to claim 1, characterized in that, S3 specifically includes: S301: Initialize the initial attention weights for each existing prediction token; S302: In response to the instruction to generate a new predicted token, calculate the attention weights of each existing predicted token through self-attention; S303: Update the corresponding initial attention weights by calculating the attention weights of each existing predicted token; S304: Use the updated initial attention weights as the attention weights for each predicted token.

5. The inference caching optimization method for generative language models according to claim 4, characterized in that, Specifically, S303 is: The calculated attention weights of each existing predicted token are summed with the initial attention weights to update the initial attention weights of each existing predicted token.

6. The inference caching optimization method for generative language models according to claim 1, characterized in that, S4 specifically includes: S401: Obtain the compression ratio of the second kv buffer for each segment; S402: Sort the attention weights in the second key-value cache of each segment; S403: Retain the key vectors and corresponding value vectors in the sorted second kv cache according to the order of attention weights from largest to smallest, wherein the proportion of retained vectors relative to all vectors in the second kv cache is the compression ratio; S404: Remove unretained vectors to perform segmented compression of each segment of the second kv cache according to the different cache segment lengths.

7. The inference caching optimization method for generative language models according to claim 6, characterized in that, The method for determining the compression ratio is as follows: Where o represents the compression ratio, d represents the distance between the position of the last token in the current cache segment of the second kv cache and the position of the first token in the last kv cache segment of the current second kv cache, B and E represent the first preset distance and the second preset distance, respectively, and p represents the minimum compression ratio.

8. A device for optimizing inference caching in a generative language model, characterized in that, include: The acquisition module is used to acquire the key-value cache of the generative language model in the nth inference process, wherein the key-value cache includes a first key-value cache of the input text and a second key-value cache in the process of generating each prediction token; The segmentation module is used to segment the second key-value cache according to the predicted token that records punctuation marks; The calculation module is used to calculate the attention weights of each predicted token; The compression module is used to combine the attention weights to segment and compress each segment of the second kv cache according to the different lengths of the cache segments; The combination module combines the first kv cache with the compressed second kv cache to obtain the updated kv cache; The reset module is used to reset the kv cache of the generative language model in the (n+1)th inference process using the updated kv cache.

9. An electronic device, characterized in that, It includes a processor, a communication interface, a memory, and a communication bus, wherein the processor, the communication interface, and the memory communicate with each other through the communication bus; The memory is used to store computer programs; When the processor executes the program stored in the memory, it implements the inference cache optimization method for generative language models as described in any one of claims 1-7.

10. A computer-readable medium, characterized in that, It stores instructions that, when executed by one or more processors, cause the processors to perform the inference cache optimization method for generative language models as described in any one of claims 1-7.

Images