A kv cache compression summary compensation method and system for LLM inference

Abstract

The application discloses a KV cache compression summary compensation method and system for LLM reasoning. Firstly, the application obtains the key and value vectors of the to-be-calculated layer query vector and the reserved set, and calculates the attention score of each historical prefix token. Secondly, in the cache compression stage, the discarded set is determined by the discard strategy, and the compensation summary is generated for the key-value pair of the discarded token through feature mapping. Then, in the autoregressive decoding reasoning stage, the calculation is performed on the current query to obtain the numerator and denominator estimates of the discarded set, and the stabilization factor scaling is performed. Finally, the numerators and denominators of the reserved set and the discarded set are synthesized to obtain the attention result after compensation. The application does not require an auxiliary small model, realizes the approximate recovery of the attention contribution of the discarded token and the error suppression, is decoupled from the discard strategy, can improve the long context reasoning accuracy and stability, and is irrelevant to the number of discarded tokens, so as to be convenient for integrated deployment.

Description

Technical Field

[0001] This invention relates to the field of large language model inference and long context attention computation technology, and particularly to a KV cache compression summary compensation method and system for LLM inference. Background Technology

[0002] Autoregressive large language models based on the Transformer architecture require self-attention computation on historical tokens during the inference phase. To avoid redundant computation of historical context, engineering implementations typically cache the key and value vectors corresponding to historical tokens as a KV cache, which is then reused repeatedly during subsequent decoding. As the context length increases, the storage size of the KV cache usually grows approximately linearly with the number of tokens, thus putting continuous pressure on GPU memory capacity, bandwidth, and memory access overhead, becoming one of the key bottlenecks in long-context inference and low-cost deployment.

[0003] To alleviate the resource constraints caused by the growth of KV cache, existing technologies have proposed various methods to reduce KV cache usage and attention computation overhead without significantly reducing generation quality. These methods mainly include, but are not limited to:

[0004] (1) Streaming inference based on sliding window / fixed retention: Only the most recent window token is retained, and a small number of "anchor / convergence" tokens are retained to maintain decoding stability, so that the cache size remains bounded under long sequences, which is suitable for streaming or infinitely long input scenarios.

[0005] (2) Cache eviction or selection based on token importance / sparseness: Based on attention scoring or QK correlation, a small number of key KV entries are selectively retained for each head of each layer to participate in attention calculation; some works have observed that attention exhibits non-continuous sparse characteristics in long contexts, and propose dynamic token-level KV selection to reduce the number of KVs involved in the calculation while maintaining performance. Other works utilize the cue tail observation window to extract the attention pattern of each head and perform clustering / selective compression on the KVs in the prefill stage.

[0006] (3) Cache merging / aggregation: Instead of directly discarding key-value entries, the key-value entries to be evicted are merged into the retained entries according to a certain strategy, in an attempt to retain information while reducing the size of the cache; for example, Cache Merging (CaM) proposes to adaptively merge the cache to be evicted into the retained cache and to mitigate the output disturbance caused by eviction through a sampling strategy. Related directions also include key-value merging schemes based on token-level similarity to identify mergeable sets.

[0007] (4) Quantization, low-rank decomposition and system-level external memory / tiered caching: reduce memory usage or increase throughput by quantizing and decomposing KV values ​​or by using paging / block management, cross-request reuse and swap-in / swap-out mechanisms at the system level; for example, some inference systems use block-based KV management and eviction strategies to improve reuse and memory efficiency.

[0008] However, the aforementioned methods still face several common problems in terms of engineering implementation and performance stability. First, once the KV entries are discarded, the numerator and denominator of the attention will undergo irreversible changes, thus introducing output perturbations. Moreover, as the compression ratio increases, these perturbations tend to become more significant, affecting the inference quality. Second, there is a phenomenon of attention saliency shift during long context decoding: historical tokens that were considered unimportant in the early stages may become critical in subsequent steps, and permanent eviction makes it difficult to recover the relevant contributions in later stages. Furthermore, many methods treat marginal tokens and truly unimportant tokens almost equally, which can easily lead to over-compression of marginal information under high compression ratios, meaning that the weight of a single token is not high, but its overall cumulative contribution cannot be ignored.

[0009] To mitigate the errors caused by expulsion, some studies have introduced compensation / correction mechanisms. For example, they utilize the similarity of attention patterns among models of different sizes to assist in estimating the attention contribution of the compressed / ignored parts using a smaller model, thereby improving the attention matching and inference performance of large models in KV compression scenarios. However, introducing auxiliary models or additional predictors usually introduces additional parameters, inference latency, and deployment complexity, and may make it more difficult to conduct unified comparisons and reproducibility under different model types, token sizes, and hardware conditions.

[0010] Therefore, a compensation mechanism is still urgently needed: after KV cache compression (especially eviction), it should be able to recover or approximate the key-value contribution of the discarded tokens in a way that is decoupled from the specific discarding strategy without introducing a small model; at the same time, the mechanism should ensure that the additional computation and storage overhead does not increase linearly with the number of discarded tokens, so as to meet the efficiency and stability requirements of ultra-long context reasoning. Summary of the Invention

[0011] To address the shortcomings of existing technologies, this invention aims to solve the problem of attention output errors caused by discarding tokens in KV cache compression, and provides a KV cache compression summary compensation method and system for LLM inference. This method does not require the introduction of a small model, is decoupled from the discarding strategy, can run online, and its compensation overhead is independent of the number of discarded tokens, effectively improving the inference accuracy of large language models after token discarding.

[0012] In a first aspect, the present invention provides a KV cache compressed digest compensation method for LLM inference, comprising the following steps:

[0013] Get KV cache: Get the query vector corresponding to the attention head of the layer to be computed, as well as the key vector and value vector in the retention set;

[0014] Token attention scoring: Based on the SnapKV method, the consecutive tokens at the end of the attention head cue sequence of the layer to be calculated are selected as the observation window. The attention weight of the query token to the prefix token key vector within the observation window is calculated. Then, the tokens are aggregated along the observation window dimension to obtain the importance score of each historical prefix token. The tokens are then further pooled to obtain the smoothed score.

[0015] KV cache compression: Based on the smoothed score, the Top-K prefix tokens are selected as the retained index set, and the unselected prefix tokens are assigned to the discard set;

[0016] Compensation calculation: For each discarded prefix token, calculate the feature mapping between its key vector and value vector, and update the summary vector and summary matrix of the discard set;

[0017] Inference synthesis: stabilized numerators and denominators are calculated for the retained set and the discarded set respectively, and then the compensated output is synthesized after being scaled to a uniform scale.

[0018] Secondly, the present invention also provides a KV cache compressed digest compensation system for LLM inference, comprising:

[0019] The feature mapping module is used to generate and store the random projection matrix for each layer and each head, and output the corresponding features when the input is a key vector or a query vector;

[0020] The digest update module is used to read the key-value vector of the evicted token when the KV cache compression eviction is triggered, divide it into the retention set and the discard set, and calculate and update the digest for the discard set.

[0021] The compensation calculation module is used to read the query and perform calculations during the decoding stage, and at the same time complete the calculation of the stabilization scaling term;

[0022] The output synthesis module is used to synthesize the stable numerator and denominator of the retained set and the estimated numerator and denominator of the discarded set, and obtain the attention output.

[0023] The beneficial effects of this invention are as follows:

[0024] First, this invention decomposes the contribution of discarded tokens to attention output into numerator and denominator terms, and estimates and recovers them using feature mapping and summary accumulation. This can effectively reduce the attention output bias caused by discarding without changing the existing token selection / eviction strategy, and improve the accuracy and stability of long context reasoning.

[0025] Secondly, this invention can achieve compensation without relying on any auxiliary small model or additional prediction network, avoiding the additional parameter overhead, inference delay and deployment complexity brought about by introducing a small model, and making it easy to make unified comparison and reproduction under different model types and different context scales;

[0026] Next, this invention only needs to maintain constant-scale state variables such as the summary vector and the summary matrix. The additional computation and storage overhead in the inference phase is mainly related to the feature dimension and the attention head dimension, and is not related to the number of discarded tokens. It is suitable for efficient inference on devices with very long contexts and limited resources.

[0027] Finally, this invention employs numerical stabilization strategies such as maximum stabilization to synthesize the compensation and retention terms on a consistent scale, reducing the risk of numerical instability caused by exponential operations, and is compatible with efficient attention mechanisms such as FlashAttention, facilitating engineering integration. Attached Figure Description

[0028] Figure 1 The schematic diagram of the overall method flow provided by the present invention includes the discard digest update in the KV cache compression stage and the compensation output calculation in the decoding inference stage;

[0029] Figure 2 This is a schematic diagram of the discarded token digest update and decoding calculation provided by the present invention;

[0030] Figure 3 The present invention provides a schematic diagram of a KV cache compression digest compensation system for LLM inference, which includes a feature mapping module, a digest update module, a compensation calculation module and an output synthesis module. Detailed Implementation

[0031] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments. The following embodiments are used to explain the technical solutions of the present invention, facilitating understanding and implementation by those skilled in the art, but do not limit the scope of the present invention. Where there is no conflict, the technical features in the embodiments of the present invention can be combined with each other to form new implementation schemes, and these new implementation schemes should also be considered within the scope of protection of the present invention.

[0032] This invention proposes a KV cache compression digest compensation method for LLM inference. Please refer to [link / reference]. Figure 1This method comprises two parts: discard digest update during the KV cache compression stage and compensation output calculation during the decoding inference stage. Please refer to [reference needed]. Figure 2 This method approximates the contribution of discarded tokens to the attention output by performing feature mapping on the key-value pairs of the discarded tokens and maintaining a constant-scale summary.

[0033] Specifically, in this embodiment of the invention, taking any attention head at any layer as an example, let the query vector of this head be... The key vector and value vector of the historical token are respectively , The scoring of attention in scaling dot product Complete attention output is .in, The result is the calculation result of the original attention formula for Transformer.

[0034] Furthermore, when KV cache compression divides historical tokens into a retention set... With discarded sets When this happens, the output can be written in fractional form. The set of reserves With discarded sets Substitute into the attention formula Then, calculate separately according to the numerator and denominator, and divide it into... Therefore, directly discarding the set... This will reduce both the denominator and numerator, thus introducing output deviation. Therefore, this invention eliminates the need to retain... Under the premise of, Make approximate estimates and compensate.

[0035] Furthermore, the present invention is achieved through the following technical solution, the specific steps of which include:

[0036] S1: Get KV cache

[0037] This embodiment first obtains the query vector corresponding to the attention head of the layer to be computed. and retain the KV cache set key vectors in Sum value vector For the initial inference, since there is no compressed content, all key-value caches are treated as a reserved set. .

[0038] S2: Token Attention Score

[0039] Furthermore, this embodiment uses the SnapKV method to score the importance of historical tokens based on the attention voting method of the observation window.

[0040] Specifically, for any attention head of the layer to be computed, this embodiment divides the current cue sequence into a prefix part and an observation window part. Let the length of the cue sequence be... The length of the observation window is The prefix length is Then it satisfies

[0041]

[0042] The observation window indicates that the ends of the sequence are continuous. Each token represents the model's attention pattern to the historical context before generation; the prefix part is the set of historical tokens located before the observation window.

[0043] Furthermore, let the query matrix corresponding to the observation window be... The key matrix corresponding to the prefix part is Then, for each attention head, calculate the scaling dot product attention weight of the observation window query with respect to the prefix key, as shown in the following formula:

[0044]

[0045] in, For attention head dimension; For attention masking; This indicates the attention distribution of each query token to the prefix token within the observation window.

[0046] Subsequently, the attention weights along the observation window dimension are aggregated to obtain an importance score for each prefix token:

[0047]

[0048] Furthermore, in this embodiment, to avoid disrupting contextual continuity by selecting only discrete high-scoring tokens, importance scoring is performed. Further one-dimensional pooling is performed to aggregate the local importance of adjacent tokens, resulting in a smoothed score:

[0049]

[0050] in, This represents the pooling kernel size; This corresponds to the fill size.

[0051] Furthermore, in this embodiment, pooling can be either max pooling or average pooling.

[0052] S3: KV Cache Compression

[0053] Further, please refer to Figure 2 In this embodiment, the smoothing score is obtained based on S2. Choose the one with the highest rating The positions of the prefixed tokens are used as a set of reserved indices:

[0054]

[0055] in, It can be determined based on the maximum capacity of the KV cache and the length of the observation window, that is:

[0056]

[0057] in This indicates the maximum KV cache length that can be retained after compression.

[0058] Furthermore, the reserved set includes the reserved index set. The corresponding prefix token and all tokens in the observation window; the prefix part not included The selected tokens constitute the discarded set. .

[0059] S4: Generate Compensation Summary

[0060] Furthermore, when the token is discarded, this embodiment no longer stores it. Instead, it updates the summary vector and summary matrix.

[0061]

[0062]

[0063] Discard set of the same layer and the same head The abstract satisfies ,in , The storage size of both is only related to the storage size factor. Related to the number of tokens discarded Irrelevant; for Feature Mapping Used to approximate the calculation of attention .

[0064] Furthermore, this embodiment selects feature mapping. This makes the exponential kernel satisfy the approximate relationship:

[0065]

[0066] Preferably, Positive Random Features (PRF) are used to construct the output, ensuring its non-negativity and numerical stability in the denominator estimation. Specifically, a Gaussian sampling random projection matrix is ​​applied to each attention head. Scaling the input vector x Then calculate Its exponent and square norm are implemented using element-wise / scalar operations. Because And since it is a positive value, it can be used to approximate the cumulative contribution of the exponential scoring. Here, x can be substituted into... or Calculate separately and .

[0067] In implementation, the abstract is maintained separately for each "layer-head-batch sample", that is, one set is maintained for each layer and each head. Batch inference can be accumulated independently according to the sample dimension. It is recommended to use FP32 accumulation to reduce rounding errors for long sequences; the rest of the model can remain FP16 / BF16. If multi-query attention or Grouped-Query Attention is used, the summary can be maintained along the KV header dimension, and the corresponding summary can be reused when sharing Q headers.

[0068] At any point during the prefill phase or inference process, after KV cache compression is triggered, this invention is decoupled from the "discarding strategy": the discard set output by any compression strategy can be directly connected to the summary update process of this invention.

[0069] Specifically, the compression stage may include the following steps:

[0070] a. Determine the retention set based on existing compression strategies. With discarded sets ,in It can be composed of "the most recent window token + a small number of important anchor tokens"; For the remaining historical tokens;

[0071] b. For each token to be evicted Retrieve from the current cache ,calculate and execute .

[0072] c. After the update is complete, release and discard the key-value entries corresponding to the tokens, retaining only the set. The original KV cache and digest .

[0073] Furthermore, to reduce update overhead, this embodiment can merge discarded tokens into blocks for updating: for a block of discarded tokens... Perform matrix summation and matrix multiplication, and update all at once. and In order to utilize GPU GEMM acceleration ,in It is to discard the token Substitution Corresponding positive random features The result obtained from the calculation.

[0074] Furthermore, if compression occurs multiple times, this invention allows for multiple cumulative summaries, i.e. It consistently refers to the "discarded parts of history" as a summary of contributions.

[0075] S5: Reasoning

[0076] In the autoregressive decoding stage, a common scenario is single-step decoding ( =1). In this embodiment, the "stabilized numerator and denominator" are calculated separately for the retained set and the discarded set, and the output is synthesized under the same stabilization scale. Further, for the retained set... First calculate the maximum value stabilization term. Define a stable denominator The output obtained by the conventional implementation of the attention layer at this time It can be understood as In engineering, it can be achieved through To restore and stabilize molecules.

[0077] Furthermore, the contribution to the discarded set is obtained according to the definition of the summary. and The estimated value , : .because At the "unstable scale", in order to be consistent with the retained set, according to the same Scaling: The final compensation output is: To prevent abnormal values ​​from being caused by an excessively small denominator, it is preferable to set a lower threshold for the denominator in the implementation. : .in Desirable (It is recommended to use a larger value for FP16 / BF16), and then convert back to the model calculation accuracy after completing the numerator and denominator calculations in FP32.

[0078] In engineering systems, compensation calculations can be consistent with compression triggering: if and only if the compression strategy results in expulsion (a non-empty condition exists). After updating the digest, the decoding phase enables compensation for that layer's header; if no expulsion occurs or the digest is empty, it degenerates into normal attention output. Furthermore, this invention is decoupled from the discarding strategy: the same compensation operator can be used regardless of whether the discarding originates from a sliding window, importance selection, or merging strategy.

[0079] Please refer to Figure 3 In one exemplary embodiment, a KV cache compressed digest compensation system for LLM inference is also provided, including a feature mapping module, a digest update module, a compensation calculation module, and an output synthesis module;

[0080] The feature mapping module is used to generate and store the random projection matrix for each layer and each head. It outputs the corresponding features when the input is a key vector or a query vector. or ;

[0081] The digest update module is used to read the evicted token when a KV cache compression eviction occurs. Divide the set into retainable and discardable sets, and calculate the value of the discardable set. And update the summary and ;

[0082] The compensation calculation module is used to read the query during the decoding stage. ,calculate and obtain , Simultaneously calculate the stabilization scaling term ;

[0083] The output synthesis module is used to retain the stable numerator and denominator of the set. Compared with the estimated discard set Synthesized It performs numerical stabilization processing such as denominator threshold truncation and precision conversion, and synthesizes the compensated output after uniform scaling.

[0084] Preferably, the above modules can be integrated into the attention layer or uniformly scheduled by an external KV cache manager; in distributed inference or tensor parallel scenarios, the summary can be split in accordance with the parallel sharding rules and consistent with the KV or reduced and synchronized across devices to ensure compensation consistency.

[0085] In summary, this application achieves approximate recovery and error suppression of the attention contribution of discarded tokens during KV cache compression without introducing an auxiliary small model, through a closed-loop processing flow of KV cache compression, discarded summary update, decoding stage compensation, and numerical stabilization synthesis. At the same time, since the compensation only requires maintaining a constant-sized summary vector and summary matrix, and the additional computation and storage overhead during the inference stage is independent of the number of discarded tokens, the system resource consumption is controllable, making it suitable for integration and application in long-context large language model inference engines and related deployment platforms.

[0086] The above description is merely a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any person skilled in the art can easily conceive of various equivalent modifications or substitutions within the technical scope disclosed in this invention, and these modifications or substitutions should all be covered within the scope of protection of this invention. Therefore, the scope of protection of this invention should be determined by the scope of the claims.

Claims

1. A KV cache compression digest compensation method for LLM inference, characterized in that, Includes the following steps: Get KV cache: Get the query vector corresponding to the attention head of the layer to be computed, as well as the key vector and value vector in the retention set; Token attention scoring: Based on the SnapKV method, the consecutive tokens at the end of the attention head cue sequence of the layer to be calculated are selected as the observation window. The attention weight of the query token to the prefix token key vector within the observation window is calculated. Then, the tokens are aggregated along the observation window dimension to obtain the importance score of each historical prefix token. The tokens are then further pooled to obtain the smoothed score. KV cache compression: Based on the smoothed score, the Top-K prefix tokens are selected as the retained index set, and the unselected prefix tokens are assigned to the discard set; Compensation calculation: For each discarded prefix token, calculate the feature mapping between its key vector and value vector, and update the summary vector and summary matrix of the discard set; Inference synthesis: stabilized numerators and denominators are calculated for the retained set and the discarded set respectively, and then the compensated output is synthesized after being scaled to a uniform scale.

2. The KV cache compression digest compensation method for LLM inference according to claim 1, characterized in that, The prompt sequence is divided into a prefix part and an observation window part; The prefix part consists of consecutive tokens whose end length is equal to the observation window, used to characterize the attention pattern to historical context before model generation; The observation window portion is the set of historical tokens located before the observation window.

3. A KV cache compression digest compensation method for LLM inference according to claim 1 or 2, characterized in that, The pooling can be either max pooling or average pooling.

4. The KV cache compression digest compensation method for LLM inference according to claim 3, characterized in that, The reserved set includes the prefix token corresponding to the reserved index set and all tokens within the observation window.

5. The KV cache compression digest compensation method for LLM inference according to claim 1, characterized in that, The feature mapping is constructed using positive random features to ensure that the output is non-negative, thereby guaranteeing the non-negativity and numerical stability of the denominator term estimation.

6. A KV cache compressed digest compensation method for LLM inference according to claim 1 or 5, characterized in that, The execution steps of the KV cache compression are as follows: Determine the retention set and the discard set based on the existing compression strategy; For each token to be discarded, retrieve the key vector and value vector from the current cache, and update the digest. After the update is complete, release the key-value entry corresponding to the discarded token, and retain only the original key-value cache and digest.

7. A KV cache compression digest compensation method for LLM inference according to claim 6, characterized in that, The discarded tokens can be updated by merging blocks, and the digest vector and digest matrix can be updated synchronously through matrix summation and multiplication operations.

8. A KV cache compression digest compensation method for LLM inference according to claim 6, characterized in that, Compensation calculations are performed on the retained set and the discarded set, wherein: For the retained set, calculate its maximum value stabilization term and define the stabilization denominator; For the discard set, calculate its estimated numerator and denominator values, and then scale it uniformly with the retained set.

9. A KV cache compressed digest compensation method for LLM inference according to claim 1 or 8, characterized in that, The compensation calculation is consistent with the KV cache compression trigger: When the compression strategy evicts and the digest is updated, compensation is enabled for the attention head of the computation layer during the decoding phase; If no expulsion occurs or the summary is empty, it degenerates into a normal attention output.

10. A KV cache compressed digest compensation system for LLM inference, characterized in that, include: The feature mapping module is used to generate and store the random projection matrix for each layer and each head, and output the corresponding features when the input is a key vector or a query vector; The digest update module is used to read the key-value vector of the evicted token when the KV cache compression eviction is triggered, divide it into the retention set and the discard set, and calculate and update the digest for the discard set. The compensation calculation module is used to read the query and perform calculations during the decoding stage, and at the same time complete the calculation of the stabilization scaling term; The output synthesis module is used to synthesize the stable numerator and denominator of the retained set and the estimated numerator and denominator of the discarded set, and obtain the attention output.

Images