Multi-level KV cache optimization method based on compensation token

Abstract

The invention relates to the technical field of computers, in particular to a multi-level KV cache optimization method based on a compensation token. The method comprises the following specific steps that in the pre-filling stage, a simplified KV cache is constructed, and the simplified KV cache comprises an initial token, a first compensation token set and a nearest token; in the decoding stage, according to whether the simplified KV cache reaches the upper limit value or not and when the nearest window starts to slide out the oldest token, updating operation is activated, a second compensation token set is generated according to the abandoned token, the second compensation token set is inserted into the position between the first compensation token set of the simplified KV cache and the nearest token, and the KV cache is dynamically updated. By adopting the optimization method, the precision loss caused by discarding the long-distance token is effectively recovered by introducing the compensation token mechanism, the memory occupation of the KV cache is reduced, and the storage cost is reduced.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of computer, and in particular to a multi-level KV cache optimization method based on compensation token. BACKGROUND

[0002] The inference process of a large language model (LLM) refers to using a trained LLM to perform inference. In this process, the user inputs problem information, and the LLM can perform inference based on the problem information and output corresponding answer information. The input problem information and the output answer information are collectively referred to as context. Streaming LLM is an efficient framework that enables training of LLMs with limited attention window to generalize to infinite sequence length without any fine-tuning. It keeps the latest tokens and attention focus points and discards intermediate tokens, enabling the model to generate coherent text from the latest tokens without cache reset. The working principle of Streaming LLM is to keep the latest tokens and attention focus points and discard the intermediate tokens. The model can only process the latest tokens, and the context window is still limited by its initial pre-training. Keeping a fixed window size while discarding too many historical tokens may lead to performance degradation in most cases. Meanwhile, the existing technology uses a cache eviction strategy to improve the throughput of the inference system, but it also evicts too many historical tokens, which also leads to performance degradation. SUMMARY

[0003] The purpose of the present application is to provide a multi-level KV cache optimization method based on compensation token, which solves the above technical problems.

[0004] To achieve the above purpose, the present application provides a multi-level KV cache optimization method based on compensation token, the specific steps are as follows: Step S1: set initial parameters; Step S2: in the Prefilling stage, process the user Prompt text and build a simplified KV cache, the simplified KV cache includes initial tokens, a first compensation token set and the latest tokens; Step S3: in the Decoding stage, when the simplified KV cache reaches the upper limit value and the latest window starts to slide out the oldest token, activate the update operation, generate a second compensation token set according to the discarded tokens, and insert the second compensation token set between the first compensation token set and the latest tokens in the simplified KV cache to realize dynamic updating of the KV cache.

[0005] Preferably, the initial parameters in step S1 include the total token budget M, the locality rate r, and the initial fixed number of tokens. .

[0006] Preferably, the specific step S2 is as follows: Step S21: Convert the user's Prompt text into a token sequence using a token segmenter, and extract the key vector and value vector for each token in the token sequence; Step S22: Based on the key vector and value vector of the token, use Top-k filtering based on the self-attention score to obtain the first set of compensation tokens; Step S23: Concatenate the first compensation token set with the fixed-retention portion of the initial token and the most recent token combination to obtain a simplified KV cache. The size of the most recent token is... ; ; in, To streamline KV caching, This is a concatenation function.

[0007] The preferred formula for calculating self-attention score is as follows: ; in, For the first The key vector of each token Attention score calculated by scaling the dot product with the query vector of the last token in the token sequence. For the first The dimension of attention head; Each attention point independently calculates a score and selects the corresponding Top-k Key-Value pairs. All Top-k Key-Value pairs are then merged to obtain the first set of compensation tokens.

[0008] Preferably, step S3 is as follows: Step S31: Determine whether the simplified KV cache has reached its upper limit based on the size of the first compensation token set; When the size of the first compensation token set is greater than or equal to At this time, the key-value cache is not updated; When the size of the first compensation token set is less than When necessary, update the key-value cache; Step S32: For discarded tokens, for every N tokens generated, set the N tokens into a calculation group and put them into the calculation set; Step S33: calculating a second compensation token according to the key vector and the value vector of the N tokens; Step S34: adding the obtained second compensation token into the second compensation token set and inserting it into the simplified KV cache for updating; The updated KV cache is as follows: .

[0009] Preferably, the second compensation token is calculated according to the average of the key vector and the value vector of the N tokens, and the calculation formula is as follows: ; ; wherein, and are the key vector and the value vector of the second compensation token respectively; and are the i-th key vector and the value vector in the calculation set .

[0010] Therefore, the present application adopts the above-mentioned multi-level KV cache optimization method based on compensation token, which has the beneficial effects that: the first compensation token set is obtained by screening key tokens, which improves the recall rate of the intermediate area tokens, improves the recognition accuracy performance of the model, the second compensation token performs semantic average compression on the discarded tokens, and one token represents a piece of content, preventing information loss, the present application effectively recovers the accuracy loss caused by discarding long-distance tokens by introducing the first compensation token set and the second compensation token set, and to a certain extent, reduces the memory occupation of the KV cache and the storage cost.

[0011] The technical solutions of the present application will be further described in detail below through the accompanying drawings and embodiments. BRIEF DESCRIPTION OF DRAWINGS

[0012] Figure 1 is a flow chart of the multi-level KV cache optimization method based on compensation token of the present application; Figure 2 is a flow chart of the simplified KV cache establishment of the present application. DETAILED DESCRIPTION

[0013] ​In the description of this invention, it should be noted that the terms "upper," "lower," "inner," and "outer," etc., indicate the orientation or positional relationship based on the orientation or positional relationship shown in the accompanying drawings, or the orientation or positional relationship commonly used when the product is in use. These terms are used only for the convenience of describing the invention and for simplifying the description, and do not indicate or imply that the device or element referred to must have a specific orientation, or be constructed and operated in a specific orientation. Therefore, they should not be construed as limitations on the invention. In the description of this invention, it should also be noted that, unless otherwise explicitly specified and limited, the terms "set," "install," and "connect" should be interpreted broadly. For example, they can refer to a fixed connection, a detachable connection, or an integral connection; they can refer to a mechanical connection or an electrical connection; they can refer to a direct connection or an indirect connection through an intermediate medium; and they can refer to the internal communication between two components. Those skilled in the art can understand the specific meaning of the above terms in this invention based on the specific circumstances.

[0014] The embodiments of the present invention will now be described in detail with reference to the accompanying drawings.

[0015] like Figure 1 As shown, a multi-level key-value caching optimization method based on compensation tokens is presented in this embodiment. This method is used to accelerate inference and reduce GPU memory usage during large model inference. For example, models such as Qwen and Llama show reduced GPU memory usage and improved speed during the inference stage under different parameter values ​​(e.g., 3B, 7B, 8B) and different benchmark tests (e.g., the Needle in the Ocean test, Math, AIME, etc.). The benchmark test sets used in the process are: NiaH (Needle in the Ocean), AIME (American Mathematical Olympiad), Math (Mathematical Olympiad), and LongBench.

[0016] The specific steps are as follows: Step S1: Set initial parameters; initial parameters include total token budget M, locality rate r, and initial fixed number of tokens. .

[0017] Step S2: In the Prefilling phase, process the user's Prompt text and build a simplified KV cache, which includes the initial token, the first set of compensation tokens, and the most recent token.

[0018] like Figure 2 As shown, the specific steps S2 are as follows: Step S21: Convert the user's Prompt text into a token sequence using a token segmenter, and extract the key vector and value vector for each token in the token sequence; Step S22: Based on the key vector and value vector of the token, use Top-k filtering based on the self-attention score to obtain the first set of compensation tokens; The formula for calculating self-attention score is as follows: ; in, For the first The key vector of each token Attention score calculated by scaling the dot product with the query vector of the last token in the token sequence. For the first The dimension of the attention head. Using the query vector of the last token in the input sequence, the hidden state of the last token often encodes global semantics and is suitable as a filtering anchor.

[0019] Each attention point independently calculates a score and selects the corresponding Top-k Key-Value pairs. All Top-k Key-Value pairs are then merged to obtain the first set of compensation tokens.

[0020] Step S23: Concatenate the first set of compensation tokens with the fixed retention portion of the initial token and the most recent token combination to obtain a simplified KV cache; ; in, To streamline KV caching, For the concatenation function, absolute position encoding is applied to the tokens in C to maintain the original order information and avoid relative position misalignment caused by window sliding in dynamic methods. The initial token receives high attention due to its positional bias and is used as the sentence-first token to stabilize the attention distribution. It has a size of N0 and is fixed in place. The most recent token is closely related to the current generation position and its context. The size of the most recent token is... The tokens scattered in the middle region are fixed and retained. A small number of tokens are crucial to the task. This embodiment obtains a first compensation token set by filtering key tokens. The recall rate of tokens in the middle region is improved by using the first compensation token set, thereby improving the recognition accuracy performance of the model.

[0021] Step S3: In the Decoding stage, based on whether the simplified KV cache has reached its upper limit and when the latest window starts to slide out the oldest token, the update operation is activated. A second compensation token set is generated based on the discarded tokens, and the second compensation token set is inserted between the first compensation token set of the simplified KV cache and the latest token, thereby realizing dynamic updating of the KV cache.

[0022] Step S3 is as follows: Step S31: Determine whether the simplified KV cache has reached its upper limit based on the size of the first compensation token set.

[0023] When the size of the first compensation token set is greater than or equal to At this time, the key-value cache is not updated; When the size of the first compensation token set is less than When necessary, update the key-value cache; Step S32: For discarded tokens, for every N tokens generated, set the N tokens into a calculation group and put them into the calculation set; Step S33: Calculate the second compensation token based on the key vectors and value vectors of the N tokens; calculate the second compensation token based on the average of the key vectors and value vectors of the N tokens, using the following formula: ; ; in, and These are the key vector and value vector of the second compensation token, respectively; and respectively compute sets The Middle Each set of tokens consists of a key vector and a value vector. This approximates the overall statistical information of the discarded token set. A second compensation token performs semantic average compression on the discarded tokens, using one token to represent a segment of content to prevent information loss.

[0024] Step S34: Add the obtained second compensation token to the second compensation token set and update it by inserting it into the simplified KV cache; The updated KV cache is as follows: .

[0025] The technical solution in this embodiment, based on existing sparse attention operations, uses a compensation token set mechanism to ensure the accuracy of large models and effectively optimize overall inference efficiency. It significantly reduces context caching requirements, enables long-context deployment, supports hardware acceleration, and can be integrated with high-performance engines such as Flash Attention and Triton.

[0026] It should be pointed out finally that the above examples are only used to illustrate the technical solutions of the present application but not to limit it, and although the present application has been described in detail with reference to the preferred embodiments, it should be understood by those skilled in the art that the technical solutions of the present application can still be modified or replaced equivalently, and these modifications or equivalent replacements should not make the modified technical solutions deviate from the spirit and scope of the technical solutions of the present application.

Claims

1. A multi-level key-value caching optimization method based on compensation tokens, characterized in that, The specific steps are as follows: Step S1: Set initial parameters; Step S2: In the Prefilling phase, process the user's Prompt text and build a simplified KV cache, which includes the initial token, the first set of compensation tokens, and the most recent token; Step S3: In the Decoding stage, based on whether the simplified KV cache has reached its upper limit and when the latest window starts to slide out the oldest token, the update operation is activated. A second compensation token set is generated based on the discarded tokens, and the second compensation token set is inserted between the first compensation token set of the simplified KV cache and the latest token, thereby realizing dynamic updating of the KV cache.

2. The multi-level KV cache optimization method based on compensation tokens according to claim 1, characterized in that: In step S1, the initial parameters include the total token budget M, the locality rate r, and the initial fixed number of tokens. .

3. The multi-level KV cache optimization method based on compensation tokens according to claim 2, characterized in that, The specific steps S2 are as follows: Step S21: Convert the user's Prompt text into a token sequence using a token segmenter, and extract the key vector and value vector for each token in the token sequence; Step S22: Based on the key vector and value vector of the token, use Top-k filtering based on the self-attention score to obtain the first set of compensation tokens; Step S23: Concatenate the first compensation token set with the fixed-retention portion of the initial token and the most recent token combination to obtain a simplified KV cache. The size of the most recent token is... ; ； in, To streamline KV caching, This is a concatenation function.

4. The multi-level KV cache optimization method based on compensation tokens according to claim 3, characterized in that, The formula for calculating self-attention score is as follows: ； in, For the first The key vector of each token Attention score calculated by scaling the dot product with the query vector of the last token in the token sequence. For the first The dimension of attention head; Each attention point independently calculates a score and selects the corresponding Top-k Key-Value pairs. All Top-k Key-Value pairs are then merged to obtain the first set of compensation tokens.

5. The multi-level KV cache optimization method based on compensation tokens according to claim 4, characterized in that, Step S3 is as follows: Step S31: Determine whether the simplified KV cache has reached its upper limit based on the size of the first compensation token set; When the size of the first compensation token set is greater than or equal to At this time, the key-value cache is not updated; When the size of the first compensation token set is less than When necessary, update the key-value cache; Step S32: For discarded tokens, for every N tokens generated, set the N tokens into a calculation group and put them into the calculation set; Step S33: Calculate the second compensation token based on the key vectors and value vectors of the N tokens; Step S34: Add the obtained second compensation token to the second compensation token set and update it by inserting it into the simplified KV cache; The updated KV cache is as follows: 。 6. The multi-level KV cache optimization method based on compensation tokens according to claim 5, characterized in that, The second compensation token is calculated based on the average of the key vectors and value vectors of the N tokens, using the following formula: ； ； in, and These are the key vector and value vector of the second compensation token, respectively; and Compute sets respectively The Middle Each has a key vector and a value vector.

Images