Method and system for kv cache compression and out-of-vocabulary token recovery for large language model inference

Abstract

The present disclosure relates to the technical field of artificial intelligence and natural language processing, in particular to a KV cache compression and word piece recovery method and system for large language model inference, the method comprising: calculating the attention score of the query vector of the i-th word piece calculated by the current Transformer layer with all key vectors; performing a V cache dynamic update operation based on the attention score; for the i-th word piece, after the V cache dynamic update operation is completed, performing attention calculation using the updated value vector set in the V cache storage pool and the pre-calculated attention score partial product P. The above technical solution solves the technical problem that the existing technology is difficult to optimize the memory occupation and the calculation efficiency, and has the advantages of dynamically managing the KV cache memory occupation, maintaining the model inference quality, and improving the calculation efficiency.

Description

TECHNICAL FIELD

[0001] The present disclosure relates to the technical field of artificial intelligence and natural language processing, in particular to a KV cache compression and evicted token recovery method and system for large language model inference. BACKGROUND

[0002] With the wide deployment of large-scale language models in various application scenarios, the memory bottleneck problem in their inference process is increasingly prominent. As a key data structure for storing historical key-value pairs in the autoregressive generation process, the memory occupancy of KV Cache increases linearly with the increase of context length, which seriously restricts the deployment ability of the model in resource-constrained environments.

[0003] The current mainstream KV Cache compression technology mainly adopts three technical routes: a static truncation method based on a sliding window, such as Streaming LLM, which simply discards historical tokens that exceed the window, resulting in permanent loss of long-range dependency information; a dynamic evicting strategy based on attention scores, such as H2O, which can filter tokens according to their importance, but lacks an effective recovery mechanism; and a Paged-Attention scheme based on page management, which implements page management of KV Cache by borrowing the paging mechanism of the operating system. It evicts the entire page of KV pairs according to the strategy, rather than a single token. Its evicting strategy can be compatible with two kinds, one is direct eviction (delete page); the other is to swap the evicted page to the CPU memory to save space for another task, but with certain limitations.

[0004] In terms of recovery technology, existing solutions generally face two core defects: the traditional recalculation method needs to reload the complete projection matrix and , causing serious calculation delay; and the offload solution based on CPU memory is limited by limited memory bandwidth. It is particularly noteworthy that the matrix transformation theory proposed by Slim-Attention has compression potential, but in actual quantitative deployment scenarios, the matrix inversion operation will introduce unacceptable precision loss. These technical limitations collectively lead to the key contradiction faced by current KV Cache management: it is difficult to achieve the coordinated optimization of memory occupancy and computational efficiency under the premise of maintaining model inference quality. SUMMARY

[0005] To solve the problems in the related art, the embodiments of the present disclosure provide a KV cache compression and evicted token recovery method and system for large language model inference.

[0006] In a first aspect, the embodiments of the present disclosure provide a KV cache compression and token recovery method for large language model inference, comprising:

[0007] In the LLM inference process, the following operations are performed for at least one Transformer layer:

[0008] For the query vector of the i-th token calculated by the current Transformer layer , the attention score of the query vector with all key vectors calculated by the current Transformer layer is calculated :

[0009] where j = 1, 2, …, n; n is the total number of tokens calculated by the current Transformer layer;

[0010] Performing a V cache dynamic update operation based on the attention score, comprising:

[0011] When , performing an eviction operation to remove the value vector of the j-th token from the V cache storage pool;

[0012] When and the value vector of the j-th token has been evicted, performing a recovery operation, comprising: regenerating the value vector of the j-th token by calculating the formula and adding it to the V cache storage pool; wherein is a pre-computed fusion weight matrix, and are the key projection matrix and value projection matrix of the current Transformer layer, respectively; T is a preset threshold.

[0013] In an embodiment of the present disclosure, the preset threshold T is set as:

[0014] wherein is a configurable hyperparameter, satisfying 0 < a < 1.

[0015] In an embodiment of the present disclosure, the method further comprises:

[0016] configuring a fixed reserved window;

[0017] wherein the size of the fixed reserved window is a preset constant, and the tokens located in the fixed reserved window do not participate in the eviction operation and the recovery operation, but their attention scores are used to calculate the preset threshold.

[0018] ​In an embodiment of the present disclosure, the pre-computed fusion weight matrix is calculated and resident in the GPU memory once in the model initialization phase and loaded only when the recovery operation is performed Instead of being loaded separately With .

[0019] In an embodiment of the present disclosure, the maximum capacity of the V cache storage pool is configured as a preset parameter m , for storing up to m value vectors;

[0020] Correspondingly, the recovery operation is performed prior to the eviction operation, and when the number of value vectors in the V cache storage pool exceeds m after the recovery operation is performed, a forced eviction operation is triggered to evict value vectors with low attention scores, so that the number of finally retained value vectors does not exceed m.

[0021] In an embodiment of the present disclosure, the method performs an eviction operation and a recovery operation once in the pre-filling phase to compress the long context KV cache and improve the response efficiency to new problems.

[0022] In an embodiment of the present disclosure, the method performs an eviction operation and a recovery operation once after wp inferences are performed in the decoding phase, where wp is a preset window period length.

[0023] In an embodiment of the present disclosure, the method further comprises:

[0024] For the ith token, the attention value is calculated based on all After the V cache dynamic update operation is completed, the updated value vector set in the V cache storage pool and the pre-computed attention score partial product P are used to calculate the attention value :

[0025] V new , wherein , V new is the updated value vector set, .

[0026] In a second aspect, the present disclosure provides a large language model inference system, comprising:

[0027] An attention score calculation module configured to calculate the attention score of the ith token query vector calculated by the current Transformer layer with all key vectors calculated by the current Transformer layer :

[0028] wherein, j = 1, 2, …, n; n is the total number of word pieces calculated by the current Transformer layer;

[0029] a V cache dynamic management module configured to perform a V cache dynamic update operation based on the attention score, including:

[0030] when , a eviction operation is performed to remove the value vector of the jth word piece from the V cache storage pool ;

[0031] when and the value vector of the jth word piece has been evicted, a recovery operation is performed, including: regenerating the value vector of the jth word piece by calculating formula and adding it to the V cache storage pool; wherein, is a pre-computed fusion weight matrix, and are the key projection matrix and the value projection matrix of the current Transformer layer, respectively; T is a preset threshold.

[0032] In a third aspect, the embodiments of the present disclosure provide an electronic device, including a memory and a processor, wherein the memory is configured to store one or more computer instructions, wherein the one or more computer instructions are executed by the processor to implement the method of any one of the first aspect.

[0033] In a fourth aspect, the embodiments of the present disclosure provide a computer-readable storage medium having computer instructions stored thereon, which are executed by a processor to implement the method of any one of the first aspect.

[0034] The technical effects provided by the embodiments of the present disclosure can include the following beneficial effects:

[0035] According to the technical scheme provided by the embodiments of the present disclosure, the KV cache compression and evicted word piece recovery method for large language model inference, by presetting a threshold to control the eviction and recovery of the KV cache, and by pre-computing a matrix to reduce the recovery delay, solves the technical problem that the memory occupation and the calculation efficiency are difficult to be optimized simultaneously in the prior art, and has the advantages of dynamically managing the memory occupation of the KV cache, maintaining the model inference quality, and improving the calculation efficiency.

[0036] It should be understood that the foregoing general description and the following detailed description are only exemplary and explanatory, and cannot limit the present disclosure. BRIEF DESCRIPTION OF DRAWINGS

[0037] Other features, objects, and advantages of the present disclosure will become more apparent from the following detailed description of the non-limiting embodiments, combined with the attached drawings. The following is a description of the drawings.

[0038] Figure 1 A flowchart of a KV cache compression and token recovery method for large language model inference according to an embodiment of the present disclosure is shown.

[0039] Figure 2 A V cache dynamic management flowchart of a Transformer decoding stage according to an embodiment of the present disclosure is shown.

[0040] Figure 3 An attention score and value vector evicting reservation diagram when decoding to different tokens is shown.

[0041] Figure 4 A structural block diagram of a large language model inference system according to an embodiment of the present disclosure is shown.

[0042] Figure 5 A structural block diagram of an electronic device according to an embodiment of the present disclosure is shown.

[0043] Figure 6 A structural diagram of a computer system suitable for implementing the method according to an embodiment of the present disclosure is shown. DETAILED DESCRIPTION

[0044] Hereinafter, exemplary embodiments of the present disclosure will be described in detail with reference to the accompanying drawings so as to be easily implemented by those skilled in the art. Also, parts irrelevant to the description of the exemplary embodiments are omitted in the accompanying drawings for the sake of clarity.

[0045] In the present disclosure, it should be understood that terms such as "include" or "have" are intended to indicate that there are features, numbers, steps, actions, components, parts or combinations thereof disclosed in the specification, and do not exclude the possibility of adding one or more other features, numbers, steps, actions, components, parts or combinations thereof.

[0046] It should also be further noted that the embodiments in the present disclosure and the features in the embodiments can be combined with each other without conflict. The present disclosure will be described in detail below with reference to the accompanying drawings and in conjunction with the embodiments.

[0047] In the prior art, the current mainstream KV Cache compression technology mainly adopts three technical routes: a static truncation method based on a sliding window, such as StreamingLLM, which simply discards the history token exceeding the window, resulting in permanent loss of long-range dependency information; a dynamic eviction strategy based on attention scores, such as H2O, which can filter tokens according to their importance, but lacks an effective recovery mechanism; and a Paged-Attention scheme based on page management, which implements page management of KV Cache by borrowing the paging mechanism of an operating system, and evicts the entire page of KV pairs according to the strategy, rather than a single token. Its eviction strategy can be compatible with two kinds, one is direct eviction (delete page); the other is to swap the evicted page to the CPU memory to save space for another task, but has certain limitations. In terms of recovery technology, the existing schemes generally face two core defects: the traditional recalculation method needs to reload the complete projection matrix and , causing serious calculation delay; and the offload scheme based on CPU memory is limited by limited memory bandwidth. It is particularly worth noting that the matrix transformation theory proposed by Slim-Attention has compression potential, but in actual quantitative deployment scenarios, the matrix inversion operation will introduce unacceptable precision loss. These technical limitations have collectively led to the key contradiction faced by current KV Cache management: it is difficult to achieve the coordinated optimization of video memory occupation and calculation efficiency on the premise of maintaining the quality of model inference.

[0048] Considering the above defects, the KV Cache compression and evicted token recovery method for large language model inference provided by the present disclosure solves the technical problem that the video memory occupation and calculation efficiency are difficult to be synergistically optimized in the prior art by controlling the eviction and recovery of the KV Cache through a preset threshold, precomputing the matrix to reduce the recovery delay, and has the advantages of dynamically managing the KV Cache video memory occupation, maintaining the quality of model inference, and improving the calculation efficiency.

[0049] Figure 1 A flowchart of a KV Cache compression and evicted token recovery method for large language model inference according to an embodiment of the present disclosure is shown.

[0050] As shown in Figure 1 , the KV Cache compression and evicted token recovery method for large language model inference specifically includes the following steps S110-S130:

[0051] In step S110, for the i-th token (Token) of the query vector calculated for the current layer (for example, the current Transformer layer), the attention score of the token with all key vectors calculated for the current layer is calculated;

[0052] In step S120, a V-cache dynamic updating operation is performed based on the attention score;

[0053] In step S130, for the ith token, an attention calculation is performed based on the updated value vector set in the V-cache storage pool and the pre-computed attention score partial product P. After the V-cache dynamic updating operation is completed, an attention calculation is performed using the updated value vector set in the V-cache storage pool and the pre-computed attention score partial product P.

[0054] In step S110, the attention score is calculated using the following formula :

[0055] wherein, is the dimension of the key vector, j = 1, 2, …, n; n is the total number of tokens calculated in the current layer;

[0056] In step S120, when , a eviction operation is performed to remove the value vector of the jth token from the V-cache storage pool.

[0057] When and the value vector of the jth token has been evicted, a recovery operation is performed, including: regenerating the value vector of the jth token by calculating the formula and adding it to the V-cache storage pool; wherein, is a pre-computed fusion weight matrix, and are the key projection matrix and value projection matrix of the current layer, respectively; T is a pre-set threshold.

[0058] In step S130, the attention calculation formula is as follows:

[0059] V new , wherein, , V new is the updated value vector set, is the attention score set of the ith token, .

[0060] The inventors found that the prior art cannot effectively balance the video memory compression and model performance. By analyzing the attention score distribution law, it was found that part of the value vector of the evicted token still has potential importance. Further research found that directly recalculating the value vector requires loading the original weight matrix at the same time, which increases the video memory bandwidth pressure. Therefore, by pre-computing the matrix to eliminate the weight loading bottleneck and combining a pre-set threshold mechanism to achieve selective recovery, the model inference quality is maintained under the condition of limited video memory.

[0061] Therefore, the present application proposes a scheme of performing the above method steps for at least one Transformer layer and preferably for each Transformer layer in the LLM inference process. In this scheme, the attention score is first calculated for each Token, then the value vector of the low-score token is discarded according to the preset threshold, the value vector of the high-score token that is discarded is regenerated by the precomputed fusion weight matrix and recovered, and finally the attention calculation is performed based on the updated value vector set.

[0062] In the present disclosure, attention score calculation refers to a numerical evaluation index generated by the multi-head attention mechanism, which can be implemented by scaled dot-product attention, and is used to quantify the importance of each token to the current inference task.

[0063] The preset threshold refers to a dynamically adjusted critical value, which can be set in proportion to the maximum value of the current layer attention score, for example, the preset threshold is set to 30% of the maximum value, which is used to screen the token set that needs to be managed.

[0064] Precomputed fusion weight matrix is a composite matrix obtained by multiplying the inverse matrix of the weight matrix and , which can be calculated offline and stored during model initialization, and is used to directly load during the recovery process, avoiding repeated loading of the original weight matrix and reducing the calculation delay.

[0065] Specifically, in the model inference process, after each layer of Transformer performs standard attention calculation on the current sequence, the attention score set of all tokens is extracted. The system dynamically generates a preset threshold according to the maximum score value calculated in real time, for example, when the maximum score is 0.85, the preset threshold is set to 0.255. Tokens with scores higher than the preset threshold are added to the recovery queue. During the recovery operation, the pre-stored fusion weight matrix is multiplied by the reserved Key vector to regenerate the value vector and write it back to the video memory, completing the reconstruction of the key information.

[0066] Compared with the prior art, the dynamic discarding scheme such as H2O only performs one-way elimination operation, while the present scheme realizes reversible management of token state through the preset threshold mechanism. Compared with the weight offloading scheme of KVPR, the present method eliminates the weight loading link through precomputation, reducing the time consumption of the recovery operation.

[0067] Through the above technical solutions, this application reduces KV cache usage while maintaining model inference accuracy. Fine-grained token management is achieved through a dynamic threshold mechanism, and the application of pre-calculated fusion weight matrices reduces recovery operation time, significantly improving inference efficiency in long-context scenarios.

[0068] In one embodiment of this disclosure, the preset threshold T is set as follows:

[0069] ,in, For configurable hyperparameters, satisfying 0 < <1.

[0070] In this disclosure method, Defined as an expulsion ratio coefficient, this coefficient is used to dynamically adjust the expulsion threshold based on the current attention score. It can be implemented using a preset fixed value or an adaptive adjustment algorithm. This coefficient forms a dynamic threshold benchmark by proportionally scaling down the highest attention score. Hyperparameter The numerical range is limited to the interval (0, 1) to ensure the validity of the threshold. This refers to the maximum attention score of the current layer, which is used to establish a baseline reference value for threshold calculation. Specifically, it can be achieved by iterating through the attention scores of all tokens in the current layer and taking the maximum value.

[0071] Traditional methods like H2O employ fixed thresholds or static scaling factors based on cumulative attention, which cannot adapt to the differences in attention distribution across different layers. For example, StreamingLLM uses a fixed window length for eviction, and its threshold is implicitly set in the window boundary, lacking dynamic awareness of the semantic importance of the current context. In contrast, this solution calculates the current layer's attention peak in real time and dynamically generates a scaling threshold, enabling the compression strategy to accurately match the semantic features of the current layer, reducing the false eviction of important tokens while maintaining computational efficiency.

[0072] Through the above technical solution, this application achieves the technical effect of dynamically adjusting the compression intensity based on the current context, solving the problem of model performance fluctuation caused by static thresholds. By correlating the threshold with the peak value of the attention score calculated in real time, the system can automatically adapt to the length changes and semantic density differences of different input sequences, maintaining model inference accuracy while compressing GPU memory usage.

[0073] In one embodiment of this disclosure, the method further includes:

[0074] Configure a fixed window to be retained;

[0075] The size of the fixed retention window is a preset constant. Tokens located within the fixed retention window do not participate in the eviction and restoration operations, but their attention scores are used to calculate the preset threshold.

[0076] In the present disclosure, the fixed reserved window refers to a protected area composed of a plurality of continuous tokens generated recently, and can be implemented by a sliding window mechanism. The starting position of the window is dynamically updated with sequence generation, and the window length remains fixed. The window is used to protect tokens that have an important influence on subsequent reasoning, so as to avoid loss of key information due to misjudgment.

[0077] The preset constant refers to a fixed value of the window size that is set in advance according to the hardware configuration before model reasoning. The value can be determined by experience or offline experiment, for example, set to 2, 4 or other values. This setting avoids additional computational overhead caused by dynamic adjustment of the window, while ensuring that the memory usage is controllable.

[0078] Specifically, in the model reasoning process, the tokens in the fixed reserved window are always in a protected state, and the corresponding value vectors will not be evicted and will not participate in the recovery operation. The tokens outside the window are managed according to the preset threshold dynamically calculated, wherein the threshold calculation process includes the attention score data of the tokens in the window. When the sequence length exceeds the window size, only the tokens outside the window range are executed for eviction and recovery operations. The window size remains constant and does not change with the sequence length, effectively reducing the number of dynamic memory allocation.

[0079] Through the above technical solutions, the present application effectively protects important tokens that have a lasting impact on model reasoning while achieving dynamic compression of KV cache, avoiding semantic coherence destruction caused by excessive eviction. The fixed window mechanism reduces the complexity of dynamic memory management, and the preset constant window size ensures that the memory usage is predictable, providing stable resource occupation guarantee for end-side deployment. The attention data of the tokens in the window participates in the threshold calculation, making the eviction and recovery decisions more consistent with the current context state, and improving the accuracy of the model in long sequence generation tasks.

[0080] In an embodiment of the present disclosure, the pre-computed fusion weight matrix is calculated once in the model initialization stage and resides in the memory, and only is loaded during the recovery operation instead of being loaded separately .

[0081] In the present disclosure, the model initialization stage refers to the preparation stage before the model starts processing the input sequence after being loaded into the computing device. The pre-computed fusion weight matrix can be generated by calling a one-time matrix operation function, so that subsequent recovery operations do not need to be repeated. The memory resident refers to the pre-computed fusion weight matrix The matrix is continuously stored in the video memory space of the graphics processor, and specifically, the video memory allocation interface can be used to lock the matrix data in a specific area of the video memory, so that the matrix data can be quickly accessed during the recovery operation.

[0082] In an embodiment of the present disclosure, the maximum capacity of the V cache storage pool is configured as a preset parameter m , for storing up to m value vectors;

[0083] Correspondingly, the recovery operation is performed before the eviction operation, and when the number of value vectors in the V cache storage pool exceeds m after the recovery operation is performed, a forced eviction operation is triggered to evict value vectors with low attention scores, so that the number of finally retained value vectors does not exceed m.

[0084] In the present disclosure, the preset parameter m refers to a pre-set video memory capacity limit, which can be determined by system configuration parameters or hardware resource constraints, and is used to control the total amount of video memory occupied by the KV cache. This feature avoids inference interruption caused by resource exhaustion by forcibly restricting the use boundary of video memory.

[0085] The operation sequence of recovery first and then eviction refers to processing the value vectors of the tokens to be recovered first, and then performing eviction according to the updated attention scores. This feature can be implemented through a state flag queue. This feature ensures that important tokens are recovered first, and ensures that the eviction decision is based on the latest attention distribution, avoiding the elimination of important tokens.

[0086] Specifically, when the video memory capacity reaches a fixed upper limit k, the value vectors of the tokens to be recovered are first screened according to a preset threshold, and the value vectors are regenerated using the precomputed matrix . After the recovery is completed, the top m value vectors with the highest scores are retained, and the remaining value vectors are evicted. For example, when m is set to 4000, 300 value vectors may be added in the recovery stage, and at this time, the lowest 300 value vectors are eliminated through the eviction stage, and finally 4000 value vectors are retained. This sequence ensures that the attention scores of the recovered tokens are involved in the eviction decision, avoiding video memory overflow caused by the recovery operation.

[0087] Through the above technical solutions, the present application realizes better KV cache utilization under fixed video memory capacity, optimizes the execution order of the recovery and eviction operations, preferentially ensures the retention of high-value tokens, and avoids the decline of model performance caused by video memory limitations. At the same time, this scheme can dynamically adapt to changes in attention distribution at different stages, and maintain the balance between inference quality and efficiency in resource-constrained scenarios.

[0088] In one embodiment of this disclosure, the method performs an evicting and recovery operation during the pre-filling phase to compress the long context key-value cache and improve response efficiency to new problems.

[0089] In this disclosure, the pre-filling stage refers to the stage where the model receives user input and generates the initial sequence embedding, at which time the KV cache of all input tokens is fully computed and stored.

[0090] Specifically, after calculating the attention score of the initial sequence during the pre-filling phase, a preset threshold is dynamically determined based on the attention scores of all tokens in the current layer. Among tokens outside the fixed retention window, V-cache tokens with attention scores below the eviction threshold are marked as releaseable regions, while evictiond value vectors with attention scores above the recovery threshold are processed through a pre-calculated matrix. A new key-value (KV) cache is generated. This operation is performed only once, pre-compressing the memory usage of non-critical tokens in long contexts while preserving the core tokens within the window for full-precision storage. Therefore, subsequent decoding stages can directly perform inference based on the compressed KV cache, reducing memory resource contention.

[0091] In one embodiment of this disclosure, the method performs an evicting and recovery operation after each wp inference iteration during the decoding phase, where wp is a preset window length.

[0092] In this disclosed method, the window length refers to the number of inference operations required between two eviction and recovery operations. It can be implemented using a fixed value or a dynamic adjustment strategy. This parameter controls the operation execution frequency to balance memory usage and computational overhead. The eviction and recovery operations refer to the process of dynamically managing the key-value cache in video memory based on attention scores. Periodic execution of these operations avoids performance degradation caused by frequent triggering.

[0093] Specifically, during the decoding phase, when generating new tokens, a memory management process is triggered after each preset number of inference iterations. When the window length is reached, the system statistically calculates the attention scores of all current tokens and filters the key-value caches that need to be retained or restored based on dynamic thresholds. This periodic operation mechanism allows the model to effectively control memory resource consumption while maintaining inference coherence.

[0094] Through the above technical solution, this application effectively reduces the memory management overhead in long sequence generation scenarios and balances computational efficiency and resource utilization by reasonably setting the operation trigger frequency. This mechanism is particularly suitable for application scenarios that require the continuous generation of a large number of tokens, such as question-answering scenarios, reducing inference latency caused by frequent memory operations while maintaining generation quality.

[0095] Figure 2 The V cache dynamic management flow of the Transformer decoding stage is shown.

[0096] As shown in Figure 2 , in the decoding stage of the current Transformer layer, the V cache is dynamically managed in the multi-head attention (MHA) calculation process of the self-attention mechanism after a window period (wp), and the flow includes the following steps:

[0097] Input sequence: the user inputs a question, and generates a sequence embedding as the current token input sequence according to the pre-filling of the question.

[0098] Calculate QKV and : Calculate Query (Query), Key (Key), and Value (Value) vectors, and use the pre-computed fusion weight matrix to prepare for the subsequent possible Value vector recovery.

[0099] Update KV cache: update the Key and Value vectors calculated in the current step to the KV cache, where the K cache is completely retained, and the V cache is dynamically managed.

[0100] Calculate attention score: calculate the attention score of the current Query vector (Query) and all Key vectors.

[0101] Evict V cache: according to the attention score and the set threshold, evict the Value vectors whose scores are lower than the threshold T (remove them from the V cache storage pool).

[0102] Restore V cache: for the Value vectors whose attention scores reach the threshold T but have been evicted, use the formula to recalculate and restore them to the V cache storage pool.

[0103] After the restoration operation, if the number of Value vectors in the V cache storage pool exceeds the preset capacity m, the Value vector with the lowest score is eliminated to ensure that the capacity does not exceed m.

[0104] Calculate MHA: use the updated V cache (V new ) and the complete K cache to perform multi-head attention calculation.

[0105] Feedforward network (FFN): input the results of the attention calculation into the feedforward network for further processing to obtain the output token of the current layer, which is added to the Sequence Embedding.

[0106] Figure 3 The attention score and the evicting reservation diagram of the value vector when decoding to different tokens are shown. Among them, the vertical axis is Query Embedding, and the horizontal axis is Key Embedding. The green square represents the window, the window size is 2, the yellow represents the selected value vector to be reserved, the blue represents the selected value vector to be recovered, the white represents the selected value vector to be evicted, the green is the reserved value vector, and actually only the corresponding V cache is evicted. The black is the mask area.

[0107] Threshold hyperparameter = 0.6, taking the fourth token as an example, the current is 0.4, at this time the preset threshold T = 0.6 x 0.4 = 0.24, according to the preset threshold, the first value vector (0.03) is evicted and the second value vector (0.32) is reserved.

[0108] Taking the fifth token as an example, the current is 0.3, at this time the preset threshold T = 0.6 x 0.3 = 0.18, according to the threshold, the second value vector (0.06) is eliminated, the third value vector (0.24) is reserved, and the first value vector (0.18) is recovered.

[0109] Figure 4 The structural block diagram of the large language model inference system according to the embodiment of the present disclosure is shown. Among them, the device can be realized as part or all of the electronic equipment through software, hardware or combination of the two.

[0110] As shown in Figure 4 , the large language model inference system 400 comprises:

[0111] An attention score calculation module 410 is configured to calculate the attention score of the query vector of the i-th token calculated by the current Transformer layer with all key vectors calculated by the current Transformer layer: :

[0112] , wherein, is the dimension of the key vector, j = 1, 2, …, n; n is the total number of tokens calculated by the current Transformer layer;

[0113] A V cache dynamic management module 420 is configured to perform V cache dynamic update operation based on the attention score, comprising:

[0114] When , the evicting operation is performed to remove the value vector of the j-th token from the V cache storage pool ;

[0115] when And the value vector of the j-th word element If expelled, perform a recovery operation, including: via calculation. Regenerate the value vector of the j-th word. And add it to the V cache storage pool; among which, For the pre-computed fusion weight matrix, and These are the key projection matrix and value projection matrix of the current Transformer layer, respectively; T is a preset threshold.

[0116] Attention calculation module 430 is configured to perform attention calculation on the i-th word based on all After completing the V-cache dynamic update operation, the attention value is calculated using the updated set of value vectors in the V-cache storage pool and the pre-computed attention score partial product P. :

[0117] V new ,in, V new For the updated set of value vectors, .

[0118] According to the technical solution provided in the embodiments of this disclosure, by controlling the eviction and restoration of the KV cache by preset threshold and reducing the restoration latency by pre-computing the matrix, the technical problem of difficulty in coordinating the optimization of video memory usage and computing efficiency in the prior art is solved. It has the advantages of dynamically managing the video memory usage of the KV cache, maintaining the model inference quality, and improving computing efficiency.

[0119] In one embodiment of this disclosure, the preset threshold T is set as follows:

[0120] ,in, For configurable hyperparameters, satisfying 0 < <1.

[0121] In one embodiment of this disclosure, the system further includes:

[0122] The configuration module is configured to keep the window permanently open.

[0123] The size of the fixed retention window is a preset constant. The words located in the fixed retention window do not participate in the eviction and recovery operations, but their attention scores are used to calculate the preset threshold.

[0124] In one embodiment of this disclosure, the pre-computed fusion weight matrix The maximum capacity of the V cache storage pool is configured as a preset parameter , and the recovery operation is performed only by loading . .

[0125] In an embodiment of the present disclosure, the maximum capacity of the V cache storage pool is configured as a preset parameter m , and is used to store a maximum of m value vectors;

[0126] Correspondingly, the recovery operation is performed before the eviction operation, and when the number of value vectors in the V cache storage pool exceeds m after the recovery operation is performed, a forced eviction operation is triggered to evict value vectors with low attention scores, so that the number of finally retained value vectors does not exceed m.

[0127] In an embodiment of the present disclosure, the system performs an eviction operation and a recovery operation once in the pre-filling stage to compress the long context KV cache and improve the response efficiency to new problems.

[0128] In an embodiment of the present disclosure, the system performs an eviction operation and a recovery operation once after wp inferences are performed in the decoding stage, where wp is a preset window period length.

[0129] The present disclosure also discloses an electronic device, Figure 5 a structural block diagram of an electronic device according to an embodiment of the present disclosure is shown.

[0130] As Figure 5 shown, the electronic device includes a memory and a processor, wherein the memory is used to store one or more computer instructions, and the one or more computer instructions are executed by the processor to implement the method according to the embodiment of the present disclosure.

[0131] The KV cache compression and evicted token recovery method for large language model inference includes:

[0132] In the LLM inference process, the following operations are performed for at least one Transformer layer:

[0133] For the query vector of the i-th token calculated by the current Transformer layer , the attention score of the query vector with all key vectors calculated by the current Transformer layer is calculated :

[0134] , wherein is the dimension of the key vector, j = 1, 2, …, n; n is the total number of tokens calculated by the current Transformer layer.

[0135] performing a V-cache dynamic updating operation based on the attention score, comprising:

[0136] when , performing an eviction operation to remove the value vector of the j-th token from the V-cache storage pool ;

[0137] when and the value vector of the j-th token has been evicted, performing a recovery operation, comprising: regenerating the value vector of the j-th token by calculating and adding it to the V-cache storage pool; wherein, is a pre-computed fusion weight matrix, and are respectively the key projection matrix and the value projection matrix of the current Transformer layer; T is a preset threshold value; for the i-th token, after completing the V-cache dynamic updating operation based on all

[0138] , the attention value is calculated using the updated value vector set in the V-cache storage pool and the pre-computed attention score partial product P:

[0139] V new , wherein, , V new is the updated value vector set, .

[0140] In an embodiment of the present disclosure, the preset threshold value T is set as:

[0141] , wherein, is a configurable hyperparameter satisfying 0 < T < 1.

[0142] In an embodiment of the present disclosure, the method further comprises:

[0143] configuring a fixed reserved window;

[0144] wherein the size of the fixed reserved window is a preset constant, the tokens located in the fixed reserved window do not participate in the eviction operation and the recovery operation, but their attention scores are used to calculate the preset threshold value.

[0145] In an embodiment of the present disclosure, the pre-computed fusion weight matrix is calculated once in the model initialization phase and resides in the GPU memory, and only is loaded in the recovery operation without loading​​ With .

[0146] In an embodiment of the present disclosure, the maximum capacity of the V-cache storage pool is configured as a preset parameter m for storing at most m value vectors;

[0147] Correspondingly, the recovery operation is performed prior to the eviction operation, and when the number of value vectors in the V-cache storage pool exceeds m after the recovery operation is performed, a forced eviction operation is triggered to evict value vectors with low attention scores, so that the number of finally retained value vectors does not exceed m.

[0148] In an embodiment of the present disclosure, the method performs an eviction operation and a recovery operation once in the pre-filling stage to compress the long-context KV cache and improve the response efficiency to new problems.

[0149] In an embodiment of the present disclosure, the method performs an eviction operation and a recovery operation once after wp inferences are performed in the decoding stage, where wp is a preset window period length.

[0150] Figure 6 A structural schematic diagram of a computer system suitable for implementing the method according to an embodiment of the present disclosure is shown.

[0151] As shown in Figure 6 , the computer system includes a processing unit that can execute various methods in the above embodiments according to programs stored in a read-only memory (ROM) or loaded from a storage section into a random access memory (RAM). Various programs and data required for the operation of the computer system are also stored in the RAM. The processing unit, the ROM, and the RAM are connected to each other through a bus. An input / output (I / O) interface is also connected to the bus.

[0152] The following components are connected to the I / O interface: an input section including a keyboard, a mouse, etc.; an output section including a cathode ray tube (CRT), a liquid crystal display (LCD), etc., and a speaker, etc.; a storage section including a hard disk, etc.; and a communication section including a network interface card such as a LAN card, a modem, etc. The communication section performs communication processes via a network such as the Internet. A drive is also connected to the I / O interface as necessary. A removable medium, such as a magnetic disk, an optical disk, a magneto-optical disk, a semiconductor memory, etc., is attached to the drive as necessary, so that a computer program read therefrom is installed into the storage section as necessary. The processing unit can be implemented as a CPU, a GPU, a TPU, a FPGA, a NPU, etc.

[0153] In particular, the method described above can be implemented as a computer software program according to embodiments of the present disclosure. For example, embodiments of the present disclosure include a computer program product comprising a computer program tangibly embodied on a machine-readable medium, the computer program containing program code for executing the methods described above. In such embodiments, the computer program can be downloaded and installed from a network via a communication part, and / or installed from a removable medium.

[0154] The flow and block diagrams in the drawings show the architectural, functional and operational views of possible implementations of systems, methods and computer program products according to various embodiments of the present disclosure. In this regard, each block in the flow and block diagrams can represent a module, a segment, or a portion of code, which contains one or more executable instructions for implementing the specified logical function. It should also be noted that in some alternative implementations, the functions noted in the blocks can occur out of the order noted in the figures. For example, two blocks shown in succession can in fact be executed substantially concurrently, or the blocks can sometimes be executed in the reverse order, depending upon the functionality involved. It will also be noted that each block in the block diagrams and / or flowchart illustrations, and combinations of blocks in the block diagrams and / or flowchart illustrations, can be implemented by special purpose hardware-based systems that perform the specified functions or operations, or combinations of special purpose hardware and computer instructions.

[0155] The units or modules described in the embodiments of the present disclosure can be implemented by means of software, or by means of programmable hardware. The described units or modules can also be provided in a processor, and the names of these units or modules do not constitute a limitation on the units or modules themselves in some cases.

[0156] As another aspect, the present disclosure also provides a computer readable storage medium, which can be the computer readable storage medium contained in the electronic device or computer system in the above embodiments; or can exist separately, and is not assembled into the device. The computer readable storage medium stores one or more programs, which are used by one or more processors to execute the methods described in the present disclosure.

[0157] The above description is merely that of the preferred embodiments of the present disclosure and a description of the technical principles of the present disclosure. It should be understood by those skilled in the art that the inventive scope involved in the present disclosure is not limited to the technical solutions formed by the specific combinations of the above technical features, and should also cover other technical solutions formed by the combinations of the above technical features or equivalent features without departing from the inventive concept. For example, the technical solutions formed by the mutual replacement of the above features and the technical features with similar functions disclosed in the present disclosure (but not limited to) can be used.

Claims

1. A KV cache compression and out-of-vocabulary token restoration method for large language model inference, characterized in that, Comprise: In the LLM inference process, the following operations are performed for at least one Transformer layer: for the i-th wordpiece computed by the current Transformer layer , computes its attention scores with all the key vectors computed by the current Transformer layer : Wherein, j=1, 2……, n; n is the total number of word tokens calculated by the current Transformer layer; The V cache dynamic update operation is performed based on the attention score, comprising: When the eviction operation is performed, removing the value vector of the jth lexeme from the V-cache storage pool ; When and the value vector of the jth token has been evicted, performing a recovery operation, comprising: regenerating the value vector of the jth token by calculating the formula and adding to the V cache storage pool; wherein, is a pre-computed fusion weight matrix, and are the key projection matrix and the value projection matrix of the current Transformer layer, respectively; T is a preset threshold.

2. The method of claim 1, wherein, The preset threshold T is set to: wherein, is a configurable hyperparameter satisfying 0 < a < 1. < 1.

3. The method according to claim 1 or 2, characterized in that, The method further comprises: Configure a fixed reserved window; Wherein, the size of the fixed reserved window is a preset constant, the word tokens in the fixed reserved window do not participate in the evicting operation and the restoring operation, but their attention scores are used to calculate the preset threshold.

4. The method of claim 1, wherein, The pre-computed fusion weight matrix Computed once at model initialization phase and resident in GPU memory, load only at restore operation Instead of loading separately With .

5. The method of claim 1, wherein, The maximum capacity of the V cache storage pool is configured as a preset parameter m for storing up to m value vectors; Correspondingly, the restoring operation is performed before the evicting operation, and when the number of value vectors in the V cache storage pool exceeds m after the restoring operation is performed, a forced evicting operation is triggered to evict value vectors with low attention scores, so that the number of finally retained value vectors does not exceed m.

6. The method of claim 1, wherein, The method performs an evicting operation and a restoring operation once in the prefilling stage to compress the long context KV cache and improve the response efficiency to new problems.

7. The method of claim 1, wherein, The method performs an evicting operation and a restoring operation once every wp times of inference in the decoding stage, where wp is a preset window period length.

8. The method of claim 1, wherein, The method further comprises: For the i-th token, the attention value is computed based on the whole set of updated value vectors and the pre-computed attention score partial product P After the V-cache dynamic update operation is completed, the attention value is computed using the updated set of value vectors in the V-cache storage pool and the pre-computed attention score partial product P : V new wherein, , V new is the updated value vector set, .

9. A large language model inference system, comprising: Comprise: a attention score calculation module configured to calculate, for the i-th wordpiece for which the current Transformer layer is computed, an attention score of the query vector computed for the current Transformer layer with all the key vectors computed for the current Transformer layer :​ Wherein, j=1, 2……, n; n is the total number of word tokens calculated by the current Transformer layer; The V cache dynamic management module is configured to perform a V cache dynamic update operation based on the attention score, comprising: When the eviction operation is performed, removing the value vector of the jth lexeme from the V-cache storage pool ; When and the value vector of the jth token has been evicted, performing a recovery operation, comprising: regenerating the value vector of the jth token by calculating the formula and adding to the V cache storage pool; wherein, is a pre-computed fusion weight matrix, and are the key projection matrix and the value projection matrix of the current Transformer layer respectively; T is a preset threshold.

10. A computer readable storage medium having stored thereon computer instructions, wherein, The computer instructions are executed by the processor to implement the method of any one of claims 1-7.

Images