Method and device for reasoning cache optimization of generative language model

Abstract

The embodiment of the invention provides a reasoning cache optimization method and device for a generative language model, electronic equipment and a readable medium, and relates to the technical field of natural language process.The method comprises the steps that kv cache of the generative language model in the nth reasoning process is obtained; segmenting the second kv cache according to the prediction token for recording the punctuation marks; calculating the attention weight of each prediction token; in combination with the attention weight, performing segmented compression on each segment of second kv cache according to different lengths of cache segments; combining the first kv cache with the compressed second kv cache to obtain an updated kv cache; and resetting the kv cache of the generative language model in the (n + 1) th reasoning process by utilizing the updated kv cache. The video memory occupation of the reasoning cache is effectively reduced while the reasoning precision is ensured, the model reasoning speed is improved, and the problem of video memory depletion is avoided.

Description

Technical Field

[0001] The present invention relates to the technical field of natural language processing, 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 Art

[0002] KV cache stands for key-value cache, an important engineering technology for optimizing Transformer reasoning performance. It caches the key vectors and value vectors calculated by the model by exchanging space for time. The model reuses the KV cache of the last reasoning during the next autoregressive calculation to avoid repeated calculations and improve reasoning performance.

[0003] Generative language models, especially the currently popular large language models, have extremely high requirements for the size of video memory. On the one hand, the number of model parameters itself occupies more and more video memory. On the other hand, in order to improve the reasoning speed, the model generally adopts kv cache technology. Kv cache is an important engineering technology for optimizing Transformer reasoning performance. By exchanging space for time, the key vector and value vector calculated by the model are cached. The model reuses the kv cache of the last reasoning during the next autoregressive calculation to avoid repeated calculations. As the text required to be generated becomes longer and longer, the kv cache data occupied in the model reasoning process is also increasing. How to reduce the video memory usage in the generative language model reasoning process to improve the model reasoning performance is a difficulty in the industry.

[0004] However, as the number of samples and the length of the output sequence in the parallel generation strategy increase, the kv cache of the existing generative language model occupies too much video memory and the video memory occupancy rate increases too quickly, even exceeding the occupancy of the generative language model's own parameters. It is very easy to have problems such as video memory exhaustion and excessive inference delay, which reduces the user experience. Summary of the invention

[0005] The embodiment of the present invention provides an inference cache optimization method for a generative language model, an inference cache optimization device for a generative language model, an electronic device and a computer-readable storage medium to solve the problem in the prior art that as the number of samples and the output sequence length in the parallel generation strategy increase, the kv cache of the existing generative language model occupies too much video memory and the video memory occupancy rate increases too fast, even exceeding the occupancy of the parameters of the generative language model itself, which easily leads to the problem of video memory exhaustion and excessive inference delay, thereby reducing the user experience.

[0006] The embodiment of the present invention discloses a method for optimizing the inference cache of a generative language model, comprising:

[0007] S1: Obtain the kv cache of the generative language model in the nth reasoning 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 predicted token;

[0008] S2: segmenting the second kv cache according to the predicted tokens of the recorded punctuation marks;

[0009] S3: Calculate the attention weight of each predicted token;

[0010] S4: Combining the attention weight, compressing each segment of the second kv cache in segments according to the different cache segment lengths;

[0011] S5: combining the first kv cache and the compressed second kv cache to obtain an updated kv cache;

[0012] S6: Reset the kv cache of the generative language model in the (n+1)th reasoning process by using the updated kv cache.

[0013] The embodiment of the present invention discloses a device for optimizing inference cache of a generative language model, including:

[0014] An acquisition module, used to acquire a kv cache of the generative language model in the nth reasoning 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 predicted token;

[0015] A segmentation module, used for segmenting the second kv cache according to the predicted tokens of the recorded punctuation marks;

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

[0017] A compression module, used to compress each segment of the second kv cache in sections according to the length of the cache segment in combination with the attention weight;

[0018] A combining module, combining the first kv cache with the compressed second kv cache to obtain an updated kv cache;

[0019] A reset module is used to reset the kv cache of the generative language model in the (n+1)th reasoning process by using the updated kv cache.

[0020] The embodiment of the present invention further discloses an electronic device, comprising a processor, a communication interface, a memory and a communication bus, wherein the processor, the communication interface and the memory communicate with each other via the communication bus;

[0021] The memory is used to store computer programs;

[0022] The processor is used to implement the inference cache optimization method of the generative language model as described in the embodiment of the present invention when executing the program stored in the memory.

[0023] The embodiment of the present invention further discloses one or more computer-readable media having instructions stored thereon, which, when executed by one or more processors, enable the processors to execute the inference cache optimization method for the generative language model as described in the embodiment of the present invention.

[0024] The embodiments of the present invention include the following advantages:

[0025] In the process of generative language model reasoning, the present invention introduces a kv cache segmentation mechanism, selects key information for each kv cache according to attention weight and relative position, and sets different compression ratios for compression. During the whole process, the first kv cache of the input text is completely retained, and some non-critical kv cache vectors in the second kv cache are discarded, thereby reducing the impact on the output accuracy of the model. In addition, the kv cache effectively avoids the occupation of video memory, effectively avoids the problem of video memory exhaustion, and effectively improves the model reasoning speed while ensuring that key reasoning information is retained and the model reasoning accuracy is ensured. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 It is a flowchart of the steps of a method for optimizing the inference cache of a generative language model provided in an embodiment of the present invention;

[0027] Figure 2 It is a flow chart of a model reasoning based on kv cache technology provided in an embodiment of the present invention;

[0028] Figure 3 It is a schematic diagram of a process of kv cache compression provided in an embodiment of the present invention;

[0029] Figure 4 It is a schematic diagram of a calculation flow of the attention weight of each predicted token provided in an embodiment of the present invention;

[0030] Figure 5 It is a structural diagram of an inference cache optimization device for a generative language model provided in an embodiment of the present invention;

[0031] Figure 6 is a block diagram of an electronic device provided in an embodiment of the present invention;

[0032] Figure 7 is a schematic diagram of a computer-readable medium provided in an embodiment of the present invention. DETAILED DESCRIPTION

[0033] In order to make the above-mentioned objects, features and advantages of the present invention more obvious and easy to understand, the present invention is further described in detail below with reference to the accompanying drawings and specific embodiments.

[0034] Reference Figure 1 , shows a step flow chart of a method for optimizing the inference cache of a generative language model provided in an embodiment of the present invention.

[0035] Reference Figure 2 , shows a flow chart of model reasoning based on kv cache technology provided in an embodiment of the present invention. 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, the fault problem is located, thereby continuously improving and strengthening the disaster recovery system.

[0036] Figure 2 The figure shows the process flow of model reasoning based on kv cache technology. The processing flow is as follows. In the first stage, the model calculates and saves the key cache and value cache of the corresponding token length in each transformer layer according to the token input by the user. The model generates the first token after the first forward calculation. In the second stage, the model calculates the key vector and value vector of the last generated token, appends them to the key cache and value cache, and performs self-attention calculation based on the updated kv cache. The model generates the next token after the forward calculation.

[0037] For example, Figure 2In the first stage, when the initial token input by the user is received, the model processes these tokens, calculates the KV vector and stores it in the cache, and generates the first output token. The input token embedding is linearly transformed to generate the Query matrix, which is 4×64 in shape, where 4 represents the number of initial input tokens and 64 represents the hidden dimension. The input token embedding is linearly transformed to generate the Key matrix 4×64, which represents 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 embedding is linearly transformed to generate the Value matrix 4×64, which represents the specific information of each token. The Value matrix is ​​stored in the Value cache. The Query matrix and the Key cache matrix are dot-producted to obtain the attention score. The attention score is weighted by Softmax and multiplied with the Value cache matrix to generate the result matrix 4×64. After further decoding or linear mapping, the result is mapped to the vocabulary probability distribution, and finally the first output token is generated. In the second stage, after the first token is generated, the model gradually updates the KV cache according to the previously generated token for autoregressive generation. According to the last generated token, calculate the corresponding Query matrix. At this time, the shape of the Query matrix is ​​1×64, because there is only one token at present. According to the current token, calculate the Key matrix 1×64 and append it to the cached Key matrix, so that the shape of the Key cache becomes 5×64. According to the current token, calculate the Value matrix 1×64 and append it to the cached Value matrix, so that the shape of the Value cache becomes 5×64. Use the updated Query matrix 1×64 to perform dot product calculation with the Key cache matrix 5×64 to generate the attention score. Use the attention score to multiply the updated Value cache matrix 5×64 to generate the result matrix 1×64. The result is decoded or linearly mapped to generate the probability distribution of the next token and output the token. Repeat the above process to gradually generate subsequent tokens and update the Key and Value caches at the same time. The introduction of the KV cache avoids repeated calculation of the KV matrix of historical tokens, greatly improving the reasoning efficiency.

[0038] Reference Figure 3 , showing a schematic diagram of a process of kv cache compression provided in an embodiment of the present invention.

[0039] Figure 3In the example, the generated tokens are segmented by punctuation marks, and each segment is compressed to different degrees. The purpose is to optimize the kv cache. In the process of generative language model reasoning, by introducing the kv cache segmentation mechanism, selecting key information for each kv cache segment, and setting different compression ratios for compression, the kv cache of the entire model reasoning is reduced, and finally the model reasoning performance is improved.

[0040] It should be noted that, for the sake of simplicity, the method embodiments are described as a series of action combinations, but those skilled in the art should be aware that the embodiments of the present invention are not limited by the order of the actions described, because according to the embodiments of the present invention, certain steps can be performed in other orders or simultaneously. Secondly, those skilled in the art should also be aware that the embodiments described in the specification are all preferred embodiments, and the actions involved are not necessarily required by the embodiments of the present invention.

[0041] The inference cache optimization method of the generative language model may specifically include the following steps:

[0042] S1: Get the kv cache of the generative language model during the nth reasoning process.

[0043] The kv cache includes a first kv cache for the input text and a second kv cache in the process of generating each predicted token.

[0044] It should be noted that the KV cache used by the generative language model in the nth reasoning process is obtained. These caches are divided into two parts: the first KV cache and the second KV cache. The first KV cache is a key-value vector calculated based on each token of the input text, representing the initial context information. The second KV cache is a key-value vector that is gradually updated during the generation process, recording the KV pairs corresponding to each predicted token generated. These two caches together constitute the historical information that the model needs to refer to during reasoning, ensuring that the generated token can dynamically adjust the output based on the complete context.

[0045] In a possible implementation, S1 specifically includes:

[0046] S101: Obtain input text.

[0047] S102: Calculate the attention triplet of each input token in the input text through each transformer layer in the generative language model, where the attention triplet includes a query vector, a key vector, and a value vector.

[0048] S103: Cache each key vector and the corresponding value vector to obtain a first kv cache.

[0049] S104: Input the input text, the query vector and the first kv cache into the generative language model, and output the first predicted token.

[0050] S105: Calculate the attention triplet of the first predicted token through each transformer layer in the generative language model.

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

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

[0053] S108: Calculate the attention triplet of the second predicted token through each transformer layer in the generative language model.

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

[0055] 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 a second kv cache.

[0056] It should be noted that this process describes the detailed process of obtaining the KV cache for the generative language model. First, the input text is obtained, and then the attention triplet (including query vector, key vector and value vector) of each token of the input text is calculated in each Transformer layer of the model. These key and value vectors are then cached to obtain the initial first KV cache. Next, the input text, query vector and first kv cache are input into the model to generate the first predicted token, and the attention triplet of the token is calculated. Subsequently, the predicted token and its attention triplet are added to the input text to update the text, and the updated text is input into the model again to generate the next predicted token. The attention triplet of the new predicted token is calculated and used as the current predicted token, and the process is repeated until the terminator token is generated. Finally, the key and value vectors of all generated tokens are cached to form the second kv cache, which records the dynamic context information generated during the generation process. This process ensures that the model can gradually update and utilize the kv cache during the inference process to efficiently generate the output sequence.

[0057] S2: Segment the second kv cache according to the predicted tokens of the recorded punctuation marks.

[0058] It should be noted that the second KV cache generated during the generation process is segmented according to the predicted tokens corresponding to the punctuation marks. This means that by identifying the punctuation marks in the generated sequence, the KV cache is divided into multiple segments, each of which represents a context structure. After segmentation, it provides a basis for subsequent compression and optimization of the cache by importance, while reducing unnecessary data redundancy.

[0059] In a possible implementation, S2 specifically includes:

[0060] S201: extracting a prediction token representing a punctuation mark from the prediction tokens.

[0061] S202: Segment the second kv cache according to the predicted tokens representing punctuation marks.

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

[0063] S204: Output the second kv cache of each segment after segmentation.

[0064] Specifically, the process describes the specific steps for segmenting the second KV cache. First, tokens representing punctuation marks are extracted from the generated prediction tokens, and these punctuation marks are used as the basis for segmentation. Then, the second KV cache is divided into multiple segments according to these punctuation marks, and each segment corresponds to a part of the generated sequence. In S203, it is determined whether the cache length of each segment is less than the preset cache length standard. If the length of some segments is insufficient, these short segments are merged with the adjacent previous segment, and the number of segments is updated at the same time. Finally, the second KV cache with the adjusted segment structure is output. This segmentation method provides a good foundation for subsequent KV cache compression, while reducing unnecessary redundant segments and improving the efficiency of cache management.

[0065] Reference Figure 4 , showing a schematic diagram of the calculation process of the attention weights of each predicted token provided in an embodiment of the present invention.

[0066] Figure 4In the model, a corresponding weight value is first introduced for each token generated by the model. The initial value is 0. Each time the model inputs a token, the self-attention will calculate the attention score of each token above, and the attention score of each token will be accumulated to the respective weight value, as shown in the attached figure. Figure 4 , in the process of the model autoregressively outputting tokens t1 to t7, each iteration accumulates a new attention score to the weight value of the output token. Finally, each vector in each kv cache has a corresponding weight value. Figure 4 The current weight value of each token at the bottom is the total weight value obtained by accumulating the weight values ​​after multiple iterations. The larger the weight value, the more critical the information of the vector.

[0067] S3: Calculate the attention weight of each predicted token.

[0068] Among them, the attention weight is a key concept in the self-attention mechanism, which indicates the degree of attention or importance of the model to each position in the sequence when processing sequence data. It measures the similarity between the query and the key to generate a probability distribution to determine the weighted proportion of the value. Calculate the attention weight of each predicted token. Through the self-attention mechanism of the model, the correlation between the currently generated token and all previously generated tokens is calculated to obtain the attention score, and these scores are converted into attention weights through Softmax. These weights reflect the importance of each predicted token to the generated sequence, providing a basis for subsequent KV cache compression.

[0069] In a possible implementation, S3 specifically includes:

[0070] S301: Initialize the initial attention weights of each existing predicted token.

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

[0072] S302: In response to the instruction to generate a new prediction token, the attention weights of each existing prediction token are calculated through self-attention.

[0073] S303: Update the corresponding initial attention weights of the attention weights of each existing predicted token obtained through calculation.

[0074] In a possible implementation manner, S303 specifically includes:

[0075] The calculated attention weights of each existing predicted token are added to the initial attention weight to update the initial attention weights of each existing predicted token.

[0076] It can be understood that the current attention weights of each existing predicted token calculated by the self-attention mechanism are accumulated with the initial attention weights of these tokens. Through this accumulation operation, the weight values ​​of each predicted token are continuously updated, thereby dynamically reflecting their accumulated importance during the generation process.

[0077] S304: Use the updated initial attention weight as the attention weight of each predicted token.

[0078] It should be noted that this process describes in detail the steps of calculating the attention weights of each predicted token. First, the initial attention weights of all generated predicted tokens are initialized, and the initial value is set to 0. Then, each time a new predicted token is generated, the attention weights of the token to all generated tokens are calculated through the self-attention mechanism, which represents its relevance to the context. Then, the calculated attention weights are added to the initial weights of the corresponding tokens and updated. Finally, the updated weights are used as the final attention weights of each predicted token. This process dynamically accumulates the importance of each predicted token and provides accurate weight information for subsequent cache optimization.

[0079] S4: Combined with the attention weight, the second kv cache of each segment is segmented and compressed according to the length of the cache segment.

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

[0081] In a possible implementation, S4 specifically includes:

[0082] S401: Determine the compression ratio of the second kv cache of each segment.

[0083] In a possible implementation manner, the compression ratio is determined as follows:

[0084]

[0085] Among them, o represents the compression ratio, d represents the distance between the position of the last token in the cache segment of the second kv cache whose compression ratio is to be calculated and the position of the first token in the last kv cache segment in the 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 position of the last token in the cache segment whose compression ratio is to be calculated and the position of the first token in the last cache segment in the current second kv cache is less than B, the compression ratio of the cache segment whose compression ratio is to be calculated is set to 1, that is, no compression is performed, indicating that the cache segment is relatively close to the position of the token to be generated and does not need to be compressed; E is the second preset distance. If the distance between the position of the last token in the cache segment whose compression ratio is to be calculated and the position of the first token in the last cache segment in the current second kv cache is greater than E, the compression ratio of the cache segment whose compression ratio is to be calculated is set to p, that is, compression is performed at the minimum compression ratio, indicating that the cache segment is relatively far from the position of the token to be generated and requires the maximum degree of compression. Therefore, the E value is greater than the B value, and the d value range is between B and E, that is, B≤d≤E, so for the cache segments between B and E, the value range of the compression ratio calculated by them 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 vector of the current token and the first vector of the second KV cache. By setting two preset distance ranges, the compression ratio changes linearly between the maximum and minimum values ​​as d changes, ensuring that important contextual information is retained with higher priority, while secondary information far from the first segment is compressed at a higher ratio. Parameter d defines the minimum compression ratio, ensuring that a small amount of necessary information is retained even at long distances.

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

[0089] S403: retain the key vectors and the corresponding value vectors in the sorted second kv cache in descending order of attention weights, wherein the ratio of the retained vectors to all vectors in the second kv cache is the compression ratio.

[0090] S404: removing unreserved vectors, so as to perform segmented compression on each segment of the second kv cache according to different cache segment lengths.

[0091] Specifically, the process describes the specific steps for segmented compression of the second KV cache. First, the compression ratio of each segment of the second KV cache is obtained to determine the ratio range of the retained vector. Next, the attention weights in each segment of the cache are sorted so that the key-value vectors can be screened according to importance. In S403, according to the order of attention weights from large to small, the key-value pairs (Key and Value vectors) with a compression ratio after sorting are retained to ensure that high-weight information is retained first. Finally, the vectors that are not in the retained range are removed to complete the segmented compression of each segment of the second KV cache. This method effectively reduces the storage occupancy of the KV cache while retaining the most critical contextual information in the generation task.

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

[0093] S6: Reset the kv cache of the generative language model during the n+1th reasoning process by updating the kv cache.

[0094] It is understandable that the bottleneck of generative language models, especially large language models, currently seems to be more limited by video memory. During the inference process, the kv cache occupies a large part of the video memory. As the batch size (the number of samples in the parallel generation strategy) and the length of the output sequence increase, the video memory overhead occupied by the kv cache will increase rapidly, and may even exceed the video memory size occupied by the model's own weights. Therefore, it is very necessary to optimize the kv cache. By introducing a segmentation mechanism, the output tokens are segmented by punctuation marks, thereby segmenting the kv cache. Secondly, the kv cache of each segment of tokens is selected for key information, and some non-critical kv cache vectors are discarded to achieve the effect of compressing the kv cache. In addition, according to different location information, the kv cache of each segment of tokens is compressed at different compression ratios to reduce the impact on the output accuracy of the model.

[0095] The purpose of the present invention is to optimize the kv cache. In the process of generative language model reasoning, the kv cache segmentation mechanism is introduced, and key information is selected for each kv cache segment, and different compression ratios are set for compression, thereby reducing the kv cache of the entire model reasoning (see Appendix Figure 2 ), ultimately improving the model’s reasoning performance.

[0096] For example, the specific process of obtaining the kv cache is as follows: Step 1: The user inputs m tokens, and 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 will be saved to form a kv cache with a cache length of 2m for subsequent model autoregression generation calculation. After the entire model is inferred, the first token is generated and output. Step 2: The generated first token is used as input, and each transformer layer in the model calculates the corresponding 1 query vector, 1 key vector, and 1 value vector. The key vector and value vector are saved in the kv cache respectively. At this time, the length of the kv cache becomes 2(m+1). The current 1 query vector and the cached m+1 key vectors and m+1 value vectors are used for self-attention calculation. After the entire model is inferred, the second token is generated and output. Step 3: The generated second token is used as input, and step 2 is repeated to generate the next token. At the same time, the length of the kv cache increases by 2. Step 4: The model continues to autoregressively generate the next token until it generates an output token representing the terminator. The model generation is completed. If a total of n tokens are generated at this time, the kv cache length during the inference process is 2(m+n).

[0097] For the obtained kv cache, examples of kv cache segmentation and key information selection mechanism are implemented as follows:

[0098] In the process of the model autoregression to generate each token, the position of the token representing the punctuation mark "," or "." is recorded. Since each generated token has a corresponding relationship with each vector in the kv cache, the kv cache is segmented according to the position of the recorded punctuation mark. Here, an adjustable parameter C is set to control the granularity of the segmentation. When the length of one cache segment is less than C, the cache segment is merged with the previous cache segment. A corresponding weight value is introduced for each token generated by the model, and the initial value is 0. For each token input into the model, the self-attention will calculate the attention score for each token in the previous text, and the attention score of each token will be accumulated to the respective weight value. In the process of the model autoregression outputting tokens t1 to t7, a new attention score is accumulated to the weight value of the output token for each iteration. Finally, each vector in each kv cache segment has a corresponding weight value. At this time, the larger the weight value, the more critical the information of the vector is. If we want to compress a certain section of kv cache with a compression ratio of 50%, then we can sort the weight values ​​of all vectors in this section of kv cache from large to small, retain only the vectors corresponding to the first 50% of the weight values, and discard the remaining vectors in this section of cache, so as to achieve the effect of kv cache compression.

[0099] During segment compression, the compression ratio is determined as follows:

[0100] First, the kv cache corresponding to the tokens input by the user is more important, so this part of the cache is not compressed. Secondly, the model's self-attention focuses on the most recent output token. We do not compress the most recently saved kv cache segment. We can set a distance threshold B (adjustable parameter). For each kv cache segment, if the distance between the position of the last vector of the cache segment and the first vector of the kv cache segment where the token is currently generated is less than B, then the kv cache segment is not compressed. Then, set the minimum compression ratio P (adjustable parameter) and the distance threshold E (adjustable parameter). For each kv cache segment, when the distance between the position of the last vector of the cache segment and the first vector of the kv cache segment where the token is currently generated is greater than E, the cache segment is compressed according to the minimum compression ratio P. Finally, the kv cache segment between B and E is compressed. The farther the distance from the current generated token is, the smaller the compression ratio is. If the distance between the last vector of the kv cache segment and the first vector of the kv cache segment where the current generated token is located is d, then the compression ratio of the kv cache segment is In general, on the one hand, we first segmented the kv cache, and retained the key vector for each kv cache segment, so that the useful information in the entire kv cache can be more evenly preserved. On the other hand, we do not compress the kv cache segments corresponding to the tokens input by the user and the tokens output recently, and only compress the kv cache segments far away from the current output token position. The farther the distance, the smaller the compression ratio, and set the minimum compression ratio, which reduces the overall impact of kv cache compression on the accuracy of model reasoning and maintains the stability of model output.

[0101] In the actual application process, the process of S1-S6 describes a KV cache optimization method for a generative language model. First, S1 obtains the KV cache used in the reasoning process, including the first KV cache of the initial input text and the second KV cache that is gradually updated during the generation process. Then in S2, by identifying punctuation marks in the generated sequence, the second KV cache is managed in segments, and each segment corresponds to a context structure. In S3, the model calculates the attention weight of each generated token through the self-attention mechanism to reflect its importance to the generated sequence. Then, S4 combines the attention weights and segment lengths in each segment of the KV cache, sorts and compresses them, and only retains 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, S6 uses the updated KV cache for the next step of reasoning to efficiently support the generation process, reduce video memory usage, and ensure generation quality.

[0102] The embodiments of the present invention include the following advantages:

[0103] In the process of generative language model reasoning, the present invention introduces a kv cache segmentation mechanism, selects key information for each kv cache according to attention weight and relative position, and sets different compression ratios for compression. During the whole process, the first kv cache of the input text is completely retained, and some non-critical kv cache vectors in the second kv cache are discarded, thereby reducing the impact on the output accuracy of the model. In addition, the kv cache effectively avoids the occupation of video memory, effectively avoids the problem of video memory exhaustion, and effectively improves the model reasoning speed while ensuring that key reasoning information is retained and the model reasoning accuracy is ensured.

[0104] It should be noted that, for the sake of simplicity, the method embodiments are described as a series of action combinations, but those skilled in the art should be aware that the embodiments of the present invention are not limited by the order of the actions described, because according to the embodiments of the present invention, certain steps can be performed in other orders or simultaneously. Secondly, those skilled in the art should also be aware that the embodiments described in the specification are all preferred embodiments, and the actions involved are not necessarily required by the embodiments of the present invention.

[0105] In addition, refer to Figure 5 , shows a block diagram of an inference cache optimization device for a generative language model provided in an embodiment of the present invention.

[0106] The embodiment of the present invention discloses a generative language model inference cache optimization device 20, comprising:

[0107] The acquisition module 201 is used to obtain the kv cache of the generative language model in the nth reasoning 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 predicted token.

[0108] The segmentation module 202 is used to segment the second kv cache according to the predicted tokens of the recorded punctuation marks.

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

[0110] The compression module 204 is used to perform segmented compression on each segment of the second kv cache according to the different cache segment lengths in combination with the attention weights.

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

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

[0113] In a possible implementation, the acquisition module 201 is specifically used for:

[0114] Get the input text.

[0115] The attention triplet of each input token in the input text is calculated through each transformer layer in the generative language model, where the attention triplet includes the query vector, key vector and value vector.

[0116] Cache each key vector and the corresponding value vector to obtain the first kv cache.

[0117] Input the input text, query vector and first kv cache to the generative language model and output the first predicted token.

[0118] The attention triplet of the first predicted token is computed by each transformer layer in the generative language model.

[0119] Add the first predicted token and the corresponding attention triplet to the input text to update the input text.

[0120] The updated input text is re-input into the generative language model and the second predicted token is output.

[0121] The attention triplet of the second predicted token is calculated by each transformer layer in the generative language model.

[0122] The second predicted token is used as the first predicted token, and the process returns to step S106 until the generated predicted token is a terminator token.

[0123] 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 are cached to obtain a second kv cache.

[0124] In a possible implementation, the segmentation module 202 is specifically configured to:

[0125] Extract the predicted token representing the punctuation mark from the predicted tokens.

[0126] The second kv cache is segmented by the predicted tokens representing punctuation marks.

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

[0128] Output the second kv cache of each segment after segmentation.

[0129] In a possible implementation, the calculation module 203 is specifically configured to:

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

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

[0132] The attention weights of each existing predicted token are calculated to update the corresponding initial attention weights.

[0133] The updated initial attention weight is used as the attention weight of each predicted token.

[0134] In a possible implementation, the attention weights of each existing predicted token obtained by calculation are updated with the corresponding initial attention weights, specifically:

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

[0136] In a possible implementation, the compression module 204 is specifically configured to:

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

[0138] Sort the attention weights in the second kv cache of each segment.

[0139] The key vectors and the corresponding value vectors in the sorted second kv cache are retained in descending order of attention weights, wherein the ratio of the retained vectors to all vectors in the second kv cache is the compression ratio.

[0140] The unreserved vectors are removed to perform segmented compression on each segment of the second kv cache according to the different lengths of the cache segments.

[0141] In a possible implementation manner, the compression ratio is determined as follows:

[0142]

[0143] Among them, o represents the compression ratio, d represents the distance between the position of the last token in the cache segment of the second kv cache whose compression ratio is to be calculated and the position of the first token in the last kv cache segment in the second kv cache, B and E represent the first preset distance and the second preset distance respectively, and p represents the minimum compression ratio.

[0144] As for the device embodiment, since it is basically similar to the method embodiment, the description is relatively simple, and the relevant parts can be referred to the partial description of the method embodiment.

[0145] The embodiments of the present invention include the following advantages:

[0146] In the process of generative language model reasoning, the present invention introduces a kv cache segmentation mechanism, selects key information for each kv cache according to attention weight and relative position, and sets different compression ratios for compression. During the whole process, the first kv cache of the input text is completely retained, and some non-critical kv cache vectors in the second kv cache are discarded, thereby reducing the impact on the output accuracy of the model. In addition, the kv cache effectively avoids the occupation of video memory, effectively avoids the problem of video memory exhaustion, and effectively improves the model reasoning speed while ensuring that key reasoning information is retained and the model reasoning accuracy is ensured.

[0147] In addition, refer to Figure 6 , shows a block diagram of an electronic device provided in an embodiment of the present invention. The embodiment of the present invention also provides 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 through the communication bus 1304,

[0148] Memory 1303, used for storing computer programs;

[0149] The processor 1301 is used to implement the following steps when executing the program stored in the memory 1303:

[0150] S1: Obtain the kv cache of the generative language model in the nth reasoning 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 predicted token;

[0151] S2: segmenting the second kv cache according to the predicted tokens of the recorded punctuation marks;

[0152] S3: Calculate the attention weight of each predicted token;

[0153] S4: Combining the attention weight, compressing each segment of the second kv cache in segments according to the length of the cache segment;

[0154] S5: combining the first kv cache and the compressed second kv cache to obtain an updated kv cache;

[0155] S6: Reset the kv cache of the generative language model in the (n+1)th reasoning process by using the updated kv cache.

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

[0157] The communication interface is used for communication between the above terminal and other devices.

[0158] The memory may include a random access memory (RAM) or a non-volatile memory, such as at least one disk memory. Optionally, the memory may also be at least one storage device located away from the aforementioned processor.

[0159] The above-mentioned processor can be a general-purpose processor, including a central processing unit (CPU), a network processor (NP), etc.; it can also be a digital signal processor (DSP), an application specific integrated circuit (ASIC), a field programmable gate array (FPGA) or other programmable logic devices, discrete gates or transistor logic devices, discrete hardware components.

[0160] Reference Figure 7 , showing a schematic diagram of a computer-readable medium provided in an embodiment of the present invention. In another embodiment provided by the present invention, a computer-readable storage medium 1401 is also provided, in which instructions are stored, and when the instructions are run on a computer, the computer executes the inference cache optimization method of the generative language model described in the above embodiment.

[0161] In another embodiment provided by the present invention, a computer program product including instructions is also provided. When the computer program product is run on a computer, the computer executes the inference cache optimization method of the generative language model described in the above embodiment.

[0162] In the above embodiments, it can be implemented in whole or in part by software, hardware, firmware or any combination thereof. When implemented by software, it can be implemented in whole or in part in the form of a computer program product. The computer program product includes one or more computer instructions. When the computer program instructions are loaded and executed on a computer, the process or function described in the embodiment of the present invention is generated in whole or in part. 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-readable storage medium, or transmitted from one computer-readable storage medium to another computer-readable storage medium. For example, the computer instructions can be transmitted from a website site, computer, server or data center to another website site, computer, server or data center by wired (e.g., coaxial cable, optical fiber, digital subscriber line (DSL)) or wireless (e.g., infrared, wireless, microwave, etc.). The computer-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 includes one or more available media integrated. The available medium can be a magnetic medium (e.g., a floppy disk, a hard disk, a tape), an optical medium (e.g., a DVD), or a semiconductor medium (e.g., a solid-state hard disk Solid State Disk (SSD)), etc.

[0163] It should be noted that, in this article, relational terms such as first and second, etc. are only used to distinguish one entity or operation from another entity or operation, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Moreover, the terms "include", "comprise" or any other variants thereof are intended to cover non-exclusive inclusion, so that a process, method, article or device including a series of elements includes not only those elements, but also other elements not explicitly listed, or also includes elements inherent to such process, method, article or device. In the absence of further restrictions, the elements defined by the sentence "comprise a ..." do not exclude the existence of other identical elements in the process, method, article or device including the elements.

[0164] Each embodiment in this specification is described in a related manner, and the same or similar parts between the embodiments can be referred to each other, and each embodiment focuses on the differences from other embodiments. In particular, for the system embodiment, since it is basically similar to the method embodiment, the description is relatively simple, and the relevant parts can be referred to the partial description of the method embodiment.

[0165] The above description is only a preferred embodiment of the present invention and is not intended to limit the protection scope of the present invention. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present invention are included in the protection scope of the present invention.

Claims

1. A method for optimizing the inference cache of a generative language model, characterized in that: include: S1: Obtain the kv cache of the generative language model in the nth reasoning 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 predicted token; S2: segmenting the second kv cache according to the predicted tokens of the recorded punctuation marks; S3: Calculate the attention weight of each predicted token; S4: Combining the attention weight, compressing each segment of the second kv cache in segments according to the different cache segment lengths; S5: combining the first kv cache and the compressed second kv cache to obtain an updated kv cache; S6: Reset the kv cache of the generative language model in the (n+1)th reasoning process by using the updated kv cache.

2. The inference cache optimization method of the generative language model according to claim 1, characterized in that: The S1 specifically includes: S101: Acquire the input text; S102: calculating the attention triplet of 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 the corresponding value vector to obtain the first kv cache; S104: Input the input text, the query vector and the first kv cache into the generative language model, and output a first predicted token; S105: Calculate the attention triplet of the first predicted token through each transformer layer in the generative language model; S106: adding the first predicted token and the corresponding attention triplet to the input text to update the input text; S107: re-inputting the updated input text into the generative language model and outputting a 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 predicted token as the first predicted token, and return to step S106 until the generated predicted 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 cache optimization method of the generative language model according to claim 1, characterized in that: The S2 specifically includes: S201: extracting a prediction token representing the punctuation mark from the prediction token; S202: segmenting the second kv cache according to the predicted token representing the punctuation mark; S203: Determine whether the cache length of each second kv cache segment is less than the preset cache length. If so, merge the second kv cache segment with the number of segments m less than the preset cache length with the second kv cache segment with the number of segments m-1, and update the number of segments. Otherwise, proceed to the next step. S204: Output the second kv cache of each segment after segmentation.

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

5. The method for optimizing the inference cache of a generative language model according to claim 4, characterized in that: The S303 is specifically: The calculated attention weights of each existing prediction token are added to the initial attention weight to update the initial attention weights of each existing prediction token.

6. The method for optimizing the inference cache of a generative language model according to claim 1, characterized in that: The S4 specifically includes: S401: Obtain the compression ratio of the second kv cache of each segment; S402: sorting the attention weights in the second kv buffer of each segment; S403: retaining the key vectors and the corresponding value vectors in the sorted second kv cache in descending order of the attention weights, wherein the ratio of the retained vectors to all the vectors in the second kv cache is the compression ratio; S404: removing unreserved vectors, so as to perform segmented compression on each segment of the second kv cache according to different cache segment lengths.

7. The method for optimizing the inference cache of a generative language model according to claim 6, characterized in that: The compression ratio is determined in the following manner: Among them, o represents the compression ratio, d represents the distance between the position of the last token in the cache segment of the second kv cache whose compression ratio is to be calculated and the position of the first token in the last kv cache segment in the 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 the inference cache of a generative language model, characterized in that: include: An acquisition module, used to acquire a kv cache of the generative language model in the nth reasoning 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 predicted token; A segmentation module, used for segmenting the second kv cache according to the predicted tokens of the recorded punctuation marks; The calculation module is used to calculate the attention weight of each predicted token; A compression module, used to compress each segment of the second kv cache in sections according to the length of the cache segment in combination with the attention weight; A combining module, combining the first kv cache with the compressed second kv cache to obtain an updated kv cache; A reset module is used to reset the kv cache of the generative language model in the (n+1)th reasoning process by 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; The processor is used to implement the inference cache optimization method of the generative language model as described in any one of claims 1 to 7 when executing the program stored in the memory.

10. A computer-readable medium, characterized in that Instructions are stored thereon, which, when executed by one or more processors, enable the processors to execute the inference cache optimization method for a generative language model as described in any one of claims 1-7.

Images