Key-value cache method, device, equipment and medium

Abstract

The application relates to the technical field of artificial intelligence, and discloses a key value cache method and device, equipment and a medium, which comprise the following steps: acquiring Token key value cache vector sets of each attention layer in a large model; adjusting a KV cache budget according to the density distribution of attention weights, and obtaining a to-be-expelled Token key value vector set and a reserved Token key value vector set based on the KV cache budget, the Token marked in the key value cache vector set of each layer, and the accumulated attention scores of the remaining Tokens; detecting the current similarity of Tokens in the to-be-expelled Token and reserved Token key vector sets of the current time step by using a weighted moving average method, so as to recall the current to-be-expelled Token based on a moving average threshold and the current similarity; and merging the recalled Token into the reserved Token key value vector set according to the current similarity. The technology can be applied to models of business systems such as medical health and pension, and financial technology.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and in particular to a key-value caching method, apparatus, device, and medium. Background Technology

[0002] With the continuous development of artificial intelligence technology, large language models (LLMs), as an important branch of this field, are widely used across various industries due to their outstanding performance in various natural language processing tasks (such as text generation, translation, question answering, and summarization). For example, in the field of smart healthcare, large language models can be used to enable online consultations or to generate corresponding treatment plans based on patients' personalized information, thereby improving the efficiency of medical consultations. In the field of fintech, large language models can assist banks in business recommendations, marketing, and debt collection outbound calls, enabling better customer maintenance and rapid business expansion.

[0003] In related technologies, key-value (KV) caching is a commonly used optimization technique in large-scale model inference. By caching the results of previous calculations, it reduces redundant computations, thereby accelerating inference and reducing computational costs. Key-value caching requires storing a large number of intermediate results. As the scale of large models increases, the memory requirements of key-value caching grow superlinearly with the sequence length. To alleviate computational and resource costs, existing technologies typically evict "non-critical" keys based on metrics such as attention scores to compress the cache. However, this ignores the hierarchical and heterogeneous characteristics of large models, treating them indiscriminately at the layer level and applying the same eviction strategy to different layers. This leads to significant information loss, impairs deep long-context understanding, exacerbates the risk of information loss, and severely affects the generative inference capabilities of large language models in long-context scenarios. Summary of the Invention

[0004] This invention provides a key-value caching method, apparatus, device, and medium to solve the technical problem that existing key-value caching methods apply the same eviction policy to different layers of a large model, resulting in significant information loss and seriously affecting the generative reasoning ability of large language models in long context scenarios.

[0005] Firstly, a key-value caching method is provided, including:

[0006] Obtain the token key-value cache vector set for each attention layer in the large model;

[0007] The KV cache budget is adjusted according to the density distribution of attention weights in each attention layer, and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer and the remaining tokens are obtained based on the KV cache budget and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer.

[0008] The weighted moving average method is used to detect the current similarity between the tokens to be evicted in the token key vector set and the tokens in the token key vector set to be retained in the current time step, so as to recall the tokens to be evicted based on the moving average threshold and the current similarity.

[0009] Based on the current similarity, the recalled currently expellable tokens are merged into the set of key-value vectors of retained tokens to obtain the final retained key-value vectors.

[0010] Secondly, a key-value caching device is provided, comprising:

[0011] The acquisition module is used to obtain the key-value cache vector set of tokens for each attention layer in the large model;

[0012] The cache adjustment module is used to adjust the KV cache budget according to the density distribution of attention weights in each attention layer, and based on the KV cache budget and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer and the remaining tokens, obtain the set of token key-value vectors to be evicted and the set of token key-value vectors to be retained.

[0013] The recall module is used to detect the current similarity between the tokens to be expelled in the token key vector set and the tokens in the token key vector set to be retained in the current time step using a weighted moving average method, so as to recall the tokens to be expelled based on the moving average threshold and the current similarity.

[0014] The merging module is used to merge the recalled currently expellable tokens into the retained token key-value vector set according to the current similarity, so as to obtain the final retained key-value vector.

[0015] Thirdly, a computer device is provided, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor implements the steps of the above-described key-value caching method when executing the computer program.

[0016] Fourthly, a computer-readable storage medium is provided, which stores a computer program that, when executed by a processor, implements the steps of the above-described key-value caching method.

[0017] The scheme implemented by the above-mentioned key-value caching method, apparatus, device, and medium can obtain the token key-value cache vector set of each attention layer in a large model; adjust the KV cache budget according to the density distribution of attention weights in each attention layer, and obtain the token key-value vector set to be evicted and the token key-value vector set to be retained based on the KV cache budget and the cumulative attention scores of the tokens marked in the token key-value cache vector set of each attention layer and the accumulated attention scores of the remaining tokens; use the weighted moving average method to detect the current similarity between the tokens to be evicted in the token key-value vector set to be evicted and the tokens in the token key-value vector set to be retained at the current time step, so as to recall the current tokens to be evicted based on the moving average threshold and the current similarity; finally, merge the recalled current tokens to be evicted into the token key-value vector set to be retained according to the current similarity to obtain the final retained key-value vector. In this invention, it is applicable to online consultation services in smart healthcare or financial services. When developing large-scale online service models for business promotion in the technology sector, the key-value cache budget can be adjusted based on the density distribution of attention weights in each attention layer of the large model. This allows key-value cache management to adapt to the attention weight distribution characteristics of different layers, aligning with the hierarchical and heterogeneous nature of the large model. This reduces memory requirements and enables token eviction based on the key-value cache budget and accumulated attention scores. This allows for accurate identification of high-value key-value pairs. Furthermore, a weighted moving average method can be used to detect the current similarity between tokens to be evicted in the current time step and tokens in the retained key vector set. Based on the moving average threshold and current similarity, tokens to be evicted can be recalled, thus compensating for information loss during key-value cache eviction and preventing the accidental eviction of important tokens. This maintains the integrity of core contextual information, reducing cache memory while avoiding the loss of important information. It better balances resource consumption and information loss, enabling reliable answers in long-context scenarios such as online consultations or online product promotions. Attached Figure Description

[0018] To more clearly illustrate the technical solutions of the embodiments of the present invention, the drawings used in the description of the embodiments of the present invention will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0019] Figure 1 This is a schematic diagram of an application environment for a key-value caching method according to an embodiment of the present invention;

[0020] Figure 2 This is a flowchart illustrating a key-value caching method according to an embodiment of the present invention;

[0021] Figure 3 yes Figure 2 A schematic diagram of a specific implementation method for step S20;

[0022] Figure 4 yes Figure 2 A schematic diagram of a specific implementation method for step S30;

[0023] Figure 5 yes Figure 2 A schematic diagram of a specific implementation of step S40;

[0024] Figure 6 This is a schematic diagram of a key-value caching device in one embodiment of the present invention;

[0025] Figure 7 This is a schematic diagram of the structure of a computer device according to an embodiment of the present invention;

[0026] Figure 8 This is another structural schematic diagram of a computer device according to one embodiment of the present invention. Detailed Implementation

[0027] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of the present invention. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0028] The key-value caching method provided in this embodiment of the invention can be applied to, for example... Figure 1In this application environment, the client communicates with the server via a network. The key-value caching method of this invention can be configured to execute on the server side, based on data interaction between the client and the server. For example, a large model application can be deployed on the client side. The server can obtain the inference business data of the large model through the client. The server converts the inference business data into token key-value pair data. When performing key-value caching, it first obtains the token key-value cache vector set of each attention layer in the large model. The token key-value cache vector set can be composed of token key-value pair data of each layer. Then, it adjusts the KV cache budget according to the density distribution of attention weights in each attention layer. Based on the KV cache budget and the cumulative attention scores of the tokens marked in the token key-value cache vector set of each attention layer and the remaining tokens, it obtains the token key-value vector set to be evicted and the token key-value vector set to be retained. Furthermore, it uses a weighted moving average method to detect the current similarity between the tokens to be evicted in the token key-value vector set and the tokens in the token key-value vector set at the current time step, so as to recall the current tokens to be evicted based on the moving average threshold and the current similarity. Finally, it merges the recalled tokens to be evicted into the token key-value vector set according to the current similarity. The key-value vector is obtained by quantizing the key-value pairs. In this invention, for large-scale online services such as online consultation services in smart healthcare or business promotion in the fintech field, the key-value cache budget can be adjusted according to the density distribution of attention weights in each attention layer of the large model. This allows the key-value cache management to adapt to the attention weight distribution characteristics of different layers, adapting to the hierarchical and heterogeneous characteristics of the large model, thereby reducing memory requirements. Tokens can be evicted based on the key-value cache budget and accumulated attention scores to accurately identify high-value key-value pairs. Furthermore, a weighted moving average method can be used to detect the current similarity between the tokens to be evicted in the key-value vector set and the tokens in the key-value vector set of the retained tokens at the current time step. Based on the moving average threshold and current similarity, the tokens to be evicted can be recalled, thereby compensating for information loss during key-value cache eviction, preventing the wrong eviction of important tokens, and maintaining the integrity of core context information. While reducing cache memory, it avoids the loss of important information, better balancing resource consumption and information loss. This enables reliable answers in long-context scenarios such as online question-and-answer in smart healthcare and fintech fields. The client can be, but is not limited to, various personal computers, laptops, smartphones, tablets, and portable wearable devices. The server can be implemented using a standalone server or a server cluster consisting of multiple servers. The invention will be described in detail below through specific embodiments.

[0029] Please see Figure 2 As shown, Figure 2A flowchart illustrating the key-value caching method provided in this embodiment of the invention includes the following steps S10-S40:

[0030] S10: Obtain the key-value cache vector set of tokens for each attention layer in the large model.

[0031] In this invention, the large model can be a large model of various dialogue modes, and the token key-value cache vector set includes the token key vector set and the corresponding value vector set.

[0032] In this embodiment, the token key-value cache vector set of each layer can be obtained by accessing the list of key-value pair tuples maintained by each attention layer in the large model, or by traversing the attention layers in the large model and capturing the token key-value cache vector set of each layer through hook functions.

[0033] The key-value caching method of this invention is applicable to large models of various dialogue modes and can be applied in scenarios such as medical consultation or financial promotion. For example, when providing medical consultation services through a large model, the key-value caching method provided in this embodiment can be used to cache the key-value data of inference business data in the large model, thereby reducing memory requirements. While accurately identifying high-value keys, it can also prevent the false expulsion of important tokens, which is beneficial to improving the inference performance of the large model. This allows the large model to process the generation of extended text with higher memory efficiency and high throughput inference, while minimizing the loss of contextual information. This can better assist doctors in disease diagnosis and prediction, and help patients obtain accurate and reliable medical advice. The inference business data refers to the business data involved when the large model performs related inference tasks, which may include user-submitted questions and / or responses from the large model, such as user consultation questions and treatment suggestions.

[0034] S20: Adjust the KV cache budget according to the density distribution of attention weights in each attention layer, and obtain the set of key-value vectors of tokens to be evicted and the set of key-value vectors of tokens to be retained based on the KV cache budget and the cumulative attention scores of the tokens marked in the key-value cache vector set of each attention layer and the remaining tokens.

[0035] In large-scale models, different attention layers typically focus on feature processing at different granularities. In existing LLMs, the lower layers usually focus on syntactic structure, with attention weights exhibiting a globally dense distribution, while the upper layers typically focus on key context fragments, with attention weights exhibiting a stepped sparsity. To adapt to the hierarchical heterogeneous features in large-scale models, this step adjusts the KV cache budget using the difference in cross-layer attention feature density, and determines the set of token key-value vectors to be evicted and the set of token key-value vectors to be retained based on the attention scores. The set of token key-value vectors to be evicted includes the set of token key vectors to be evicted and the corresponding set of value vectors, while the set of token key-value vectors to be retained includes the set of token key vectors to be retained and the corresponding set of value vectors.

[0036] Specifically, such as Figure 3 As shown, step S20 includes the following steps S21-S26:

[0037] S21: Obtain the density distribution of attention weights in each attention layer based on the variance of the attention score matrix in each attention layer.

[0038] In this step, an attention score matrix is ​​used to evaluate the attention density of each layer l. Specifically, according to the formula... Calculate the attention density of each layer in This represents the attention score matrix of layer l in the encoding stage. and These represent the query matrix and the key matrix, respectively, with D being the dimension of the hidden features.

[0039] S22: Determine the memory budget parameters based on the target compression ratio and the density distribution of the attention weights.

[0040] In this step, according to the formula Calculate the memory budget parameter α l Where ρ represents the target compression ratio of the cache, and L is the number of attention layers in the model.

[0041] S23: Adjust the KV cache budget for each layer based on the memory budget parameters.

[0042] In this step, the formula S is used based on the memory budget parameter. l =α l ·S Adjust the KV cache budget for each layer S l Where S is the original buffer capacity, typically the length L of the Prompt during the encoding phase. prompt As can be seen from the above, in this embodiment, the smaller the variance, the denser the attention distribution. A smaller KV cache budget can be allocated to layers with higher variance (e.g., the top layer), while layers with lower variance (e.g., the bottom layer) can obtain a larger KV cache budget.

[0043] S24: Tag the token for the current time step in each attention layer based on the sliding window length and the initial token length.

[0044] In this step, the most recent token at the current time step is marked based on the sliding window length M, and the initial token at the current time step is marked based on the initial tag length T.

[0045] S25: Calculate the cumulative attention score of all tokens in the token key-value cache vector set of each attention layer, excluding the labeled tokens.

[0046] In this invention, the KV cache is also dynamically updated using the accumulated attention score.

[0047] In this step, the cumulative attention score for all tokens except the labeled token is calculated according to the following formula:

[0048]

[0049] S26: Based on the KV cache budget, sliding window length, initial tag length, and the cumulative attention scores of the remaining tokens, sort them to obtain the key-value vector set of tokens to be evicted and the key-value vector set of tokens to be retained.

[0050] Specifically, in this step, based on the difference between the KV cache budget and the sliding window length and the initial label length, the top N tokens with the highest cumulative attention scores are selected from the remaining tokens; the tokens labeled by the sliding window length and the initial label length, along with the selected top N tokens, form a set of retained token key-value vectors, and the unselected tokens from the remaining tokens form a set of token key-value vectors to be evicted; in this embodiment, (Kc, Vc) can be used to represent the key-value pairs retained after eviction, and the set of token key-value vectors to be evicted, Ke = KK. c .

[0051] Understandably, the initial tokens in the retained token key-value vector set typically contain the core context of the hint / question, while the most recent tokens in the sliding window length typically contain the most relevant local context generated at the moment. In this invention, the initial tokens and most recent tokens have the highest retention priority to ensure the retention of key contexts, and the historical importance of tokens is quantified by accumulating attention scores to adaptively identify and retain important tokens.

[0052] Based on steps S21-S26 above, the attention density is measured by variance; the smaller the variance, the denser the attention distribution. A function similar to Softmax is used to dynamically adjust the memory budget parameter α of each layer according to the attention density. l This allows for the adjustment of the KV cache budget based on the density distribution of attention weights. Layers with higher variance (e.g., the top layer) are allocated smaller KV cache budgets, while layers with lower variance (e.g., the bottom layer) receive larger KV cache budgets, thus enabling on-demand and adaptive allocation of computing resources.

[0053] S30: Utilize the weighted moving average method to detect the current similarity between the tokens to be evicted in the token key vector set and the tokens in the token key vector set to be retained at the current time step, so as to recall the tokens to be evicted based on the moving average threshold and the current similarity.

[0054] To avoid losing the integrity of the long context by directly discarding the tokens in the key-value vector set of the tokens to be evicted, this step determines whether the token to be evicted contains key information based on the similarity between the token to be evicted and the retained tokens, and decides whether to recall or discard the token to be evicted.

[0055] Specifically, such as Figure 4 As shown, step S30 includes the following steps S31-S34:

[0056] S31: Calculate the cosine similarity matrix between the set of tokens to be evicted and the set of tokens to be retained at the current time step using the nearest neighbor matching algorithm.

[0057] In this invention, the cosine similarity matrix between the set of key vectors of tokens to be evicted and the set of key vectors of tokens to be retained is calculated using the nearest neighbor matching algorithm. That is, the current similarity between any two tokens is obtained by calculating the cosine similarity between the tokens to be evicted in the set of key vectors of tokens to be evicted and the tokens in the set of key vectors of tokens to be retained. Form a cosine similarity matrix.

[0058] S32: Obtain the current maximum similarity between the tokens to be expelled in the set of token key vectors to be expelled and the tokens in the set of token key vectors to be retained, based on the cosine similarity in the cosine similarity matrix.

[0059] In this step, the most similar token to the expellable token in the retained token key vector set is determined based on the cosine similarity between the two tokens, thus obtaining the current maximum similarity between the expellable token and the retained token key vector set. Understandably, a token in the retained token key vector set can correspond to multiple most similar expellable tokens.

[0060] S33: Smooth the cosine similarity matrix using the exponential moving average method to obtain the moving average threshold for the current time step.

[0061] Using the formula τ t =βMax(U t [:])+(1-β)τ t-1 The moving average threshold τ at the current time step t is calculated based on the cosine similarity matrix. t Among them, U t Let represent the cosine similarity matrix, and β represent the smoothing constant.

[0062] Understandably, in this embodiment, the initial expulsion threshold τ can be set as the mean of the maximum similarity between all tokens to be expelled and the tokens in the key vector set of the retained tokens. Therefore, the moving average threshold τ... t It can be calculated using the following formula:

[0063]

[0064] Among them, L e L represents the length of the key vector set of tokens to be evicted, i.e., the number of tokens to be evicted. c U represents the number of tokens in the retained token key vector set, and the smoothing constant β controls the U. t With historical threshold τ t-1 The weighted combination relationship shows that the higher the β value, the stronger the sensitivity to current similarity fluctuations, while also smoothing out historical similarity fluctuations.

[0065] S34: Compare the moving average threshold with the current maximum similarity. If the current maximum similarity is not lower than the moving average threshold, recall the current token to be expelled.

[0066] In this step, the similarity between the token to be expelled and the retained token is evaluated by a moving average threshold to dynamically determine whether to recall the current token to be expelled based on the current moving average threshold. If the maximum similarity is lower than the moving average threshold, the current token to be expelled is permanently expelled; otherwise, the current token to be expelled is recalled.

[0067] Based on the above steps S31-S34, the similarity between the token to be evicted and the retained token is detected by the exponential moving average method, thereby determining whether the token to be evicted contains key information. If so, the token to be evicted is recalled, thereby retaining valuable information from the token to be evicted and compensating for the loss of information from the KV cache eviction.

[0068] S40: Based on the current similarity, merge the recalled currently expellable tokens into the retained token key-value vector set to obtain the final retained key-value vector.

[0069] In this invention, the currently pending eviction tokens are recalled and then weighted and merged into the retained token key-value vector set.

[0070] Specifically, such as Figure 5 As shown, step S40 includes the following steps S41-S42:

[0071] S41: Calculate the weights of the currently recalled Tokens to be expelled and the similar Tokens in the key-value vector set of the retained Tokens based on the current similarity.

[0072] In this embodiment, the currently recalled token to be expelled is merged with the most similar token in the set of key-value vectors of the retained tokens.

[0073] In this step, the formula is used. Based on the current similarity u ij Calculate the weights. Where k ei The key w represents the currently pending eviction token to be recalled. ei k represents the contribution weight of the currently expellable tokens to be recalled. cj w represents the key of the tokens retained in the key-value vector set. cj Indicates the weight of the retained token, m ij This represents the nearest neighbor mask, if k cj It is k ei The nearest neighbor, i.e., the most similar token, is m. ij =1, if not the nearest neighbor, then m ij =0, e is a constant.

[0074] S42: Based on the weight, the key values ​​of the current Token to be expelled and the key values ​​of similar Tokens in the set of key value vectors of the retained Tokens are weighted and merged to obtain the final retained key value vector.

[0075] In this step, the recalled currently expellable tokens are weighted and merged into the most similar token in the retained token key-value vector set, replacing the original retained token's key value. One retained token can correspond to multiple most similar recalled currently expellable tokens, resulting in the final retained key-value vector. Where, k ei With v ei Let k be the i-th key-value pair in the set of key-value vectors to be evicted from the token set. cj With v cjTo preserve the j-th key-value pair in the Token key-value vector set.

[0076] Based on the above steps S41-S42, the currently recalled tokens to be expelled and the retained tokens can be dynamically weighted and merged according to the current similarity, thus maintaining the integrity of the core context information without the need for fine-tuning.

[0077] As can be seen, in the above solution, for large-scale dialogue-based models providing online services such as online consultations in smart healthcare or business promotion in the fintech field, the key-value cache budget can be adjusted according to the density distribution of attention weights in each attention layer of the large model. This allows key-value cache management to adapt to the attention weight distribution characteristics of different layers, fitting the hierarchical and heterogeneous characteristics of the large model, thereby reducing memory requirements. Furthermore, tokens can be evicted based on the key-value cache budget and accumulated attention scores to accurately identify high-value key-value pairs. Additionally, a weighted moving average method can be used to detect the current similarity between tokens to be evicted in the key vector set and tokens in the key vector set to be retained at the current time step, based on... By recalling the currently eviction token based on the moving average threshold and current similarity, the information loss caused by KV cache eviction is compensated, and important tokens are prevented from being mistakenly evicted, thus maintaining the integrity of core context information. While reducing cache memory, it can avoid the loss of important information, thereby better assisting doctors in disease diagnosis and prediction during medical consultations, or better understanding user intent and recommending more suitable products during financial promotions. It can be seen that the key-value caching method of this invention better balances the resource consumption and information loss of large models, which is conducive to improving the reasoning and generation capabilities in long context scenarios and can achieve reliable answers in long context scenarios. When used, it can improve the efficiency and reliability of medical consultations or financial promotions.

[0078] It should be understood that the sequence number of each step in the above embodiments does not imply the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiments of the present invention.

[0079] In one embodiment, a key-value caching device is provided, which corresponds one-to-one with the key-value caching methods described in the above embodiments. For example... Figure 6 As shown, the key-value caching device includes an acquisition module 110, a cache adjustment module 120, a recall module 130, and a merging module 140. Detailed descriptions of each functional module are as follows:

[0080] Module 110 is used to obtain the Token key-value cache vector set for each attention layer in the large model;

[0081] The cache adjustment module 120 is used to adjust the KV cache budget according to the density distribution of attention weights in each attention layer, and obtain the set of key-value vectors of tokens to be evicted and the set of key-value vectors of tokens to be retained based on the KV cache budget and the cumulative attention scores of the tokens marked in the key-value cache vector set of each attention layer and the remaining tokens.

[0082] The recall module 130 is used to detect the current similarity between the tokens to be expelled in the token key vector set and the tokens in the token key vector set to be retained in the current time step using a weighted moving average method, so as to recall the tokens to be expelled based on the moving average threshold and the current similarity.

[0083] The merging module 140 is used to merge the recalled currently expellable tokens into the retained token key-value vector set according to the current similarity, so as to obtain the final retained key-value vector.

[0084] In one embodiment, the cache adjustment module 120 is specifically used for:

[0085] The density distribution of attention weights in each attention layer is obtained by calculating the variance of the attention score matrix in each attention layer.

[0086] The memory budget parameters are determined based on the target compression ratio and the density distribution of the attention weights;

[0087] Adjust the KV cache budget for each layer based on the aforementioned memory budget parameters.

[0088] In one embodiment, the cache adjustment module 120 is further configured to:

[0089] The token for the current time step in each attention layer is marked based on the sliding window length and the initial token length;

[0090] Calculate the cumulative attention score for all tokens in the token key-value cache vector set of each attention layer, excluding the labeled tokens;

[0091] Based on the KV cache budget, sliding window length, initial tag length, and the cumulative attention scores of the remaining tokens, the key-value vector set of tokens to be evicted and the key-value vector set of tokens to be retained are obtained.

[0092] In one embodiment, the cache adjustment module 120 is further configured to:

[0093] Based on the difference between the KV cache budget and the sliding window length and the initial tag length, select the top N tokens with the highest cumulative attention scores from the remaining tokens;

[0094] The tokens marked by the sliding window length and the initial token length, along with the first N selected tokens, form a set of retained token key-value vectors. The remaining unselected tokens form a set of token key-value vectors to be evicted.

[0095] In one embodiment, the recall module 130 is specifically used for:

[0096] The cosine similarity matrix of the set of tokens to be evicted and the set of tokens to be retained at the current time step is calculated using the nearest neighbor matching algorithm.

[0097] The current maximum similarity between the tokens to be expelled in the set of token key vectors to be expelled and the tokens in the set of token key vectors to be retained is obtained based on the cosine similarity in the cosine similarity matrix.

[0098] The cosine similarity matrix is ​​smoothed using the exponential moving average method to obtain the moving average threshold for the current time step;

[0099] Compare the moving average threshold with the current maximum similarity. If the current maximum similarity is not lower than the moving average threshold, recall the current token to be expelled.

[0100] In one embodiment, the recall module 130 is further specifically used for:

[0101] Using the formula τ t =βMax(U t [:])+(1-β)τ t-1 The moving average threshold τ at the current time step t is calculated based on the cosine similarity matrix. t Among them, U t Let represent the cosine similarity matrix, and β represent the smoothing constant.

[0102] In one embodiment, the merging module 140 is specifically used for:

[0103] The weights of the currently recalled Tokens to be expelled and the similar Tokens in the key-value vector set of the retained Tokens are calculated based on the current similarity.

[0104] The key values ​​of the current Token to be expelled and the key values ​​of similar Tokens in the set of key value vectors of the retained Tokens are weighted and merged according to the weight to obtain the final retained key value vector.

[0105] This invention provides a key-value caching device that can acquire the key-value cache vector set of tokens for each attention layer in a large model; adjust the KV cache budget according to the density distribution of attention weights in each attention layer; and obtain the key-value vector set of tokens to be evicted and the key-value vector set of retained tokens based on the KV cache budget and the cumulative attention scores of the tokens marked and other tokens in the key-value cache vector set of each attention layer; furthermore, it uses a weighted moving average method to detect the current similarity between the tokens to be evicted in the key-value vector set of tokens to be evicted and the tokens in the key-value vector set of retained tokens at the current time step, so as to recall the current tokens to be evicted based on the moving average threshold and the current similarity; finally, it merges the recalled current tokens to be evicted into the key-value vector set of retained tokens according to the current similarity to obtain the final retained key-value vector, i.e., the working key-value vector. It can adjust the KV cache budget according to the density distribution of attention weights in each attention layer of a large model, making key-value cache management adaptive to the attention weight distribution characteristics of different layers and adapting to the hierarchical heterogeneity of large models, thereby reducing memory requirements. It can also evict tokens based on the KV cache budget and accumulated attention scores to accurately identify high-value key-value pairs. Furthermore, it can use a weighted moving average method to detect the current similarity between the tokens to be evicted in the current time step and the tokens in the retained token key vector set, and recall the tokens to be evicted based on the moving average threshold and the current similarity, thereby compensating for the information loss of KV cache eviction, preventing the wrong eviction of important tokens, and maintaining the integrity of core context information. While reducing cache memory, it can avoid the loss of important information, better balance resource consumption and information loss, and achieve reliable answers in long context scenarios.

[0106] For specific limitations regarding the key-value caching device, please refer to the limitations of the key-value caching method above, which will not be repeated here. Each module in the aforementioned key-value caching device can be implemented entirely or partially through software, hardware, or a combination thereof. These modules can be embedded in hardware or independent of the processor in the computer device, or stored in software in the memory of the computer device, so that the processor can call and execute the operations corresponding to each module.

[0107] In one embodiment, a computer device is provided, which may be a server, and its internal structure diagram may be as follows: Figure 7As shown, the computer device includes a processor, memory, network interface, and database connected via a system bus. The processor provides computing and control capabilities. The memory includes non-volatile and / or volatile storage media and internal memory. The non-volatile storage media stores the operating system, computer programs, and database. The internal memory provides an environment for the operation of the operating system and computer programs in the non-volatile storage media. The network interface is used to communicate with external clients via a network connection. When the computer program is executed by the processor, it implements a key-value caching method on the server side, fulfilling its functions or steps.

[0108] In one embodiment, a computer device is provided, which may be a client, and its internal structure diagram may be as follows: Figure 8 As shown, the computer device includes a processor, memory, network interface, display screen, and input devices connected via a system bus. The processor provides computing and control capabilities. The memory includes non-volatile storage media and internal memory. The non-volatile storage media stores the operating system and computer programs. The internal memory provides an environment for the operation of the operating system and computer programs stored in the non-volatile storage media. The network interface is used to communicate with an external server via a network connection. When executed by the processor, the computer program implements client-side functions or steps of a key-value caching method.

[0109] In one embodiment, a computer device is provided, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to perform the following steps:

[0110] Obtain the token key-value cache vector set for each attention layer in the large model;

[0111] The KV cache budget is adjusted according to the density distribution of attention weights in each attention layer, and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer and the remaining tokens are obtained based on the KV cache budget and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer.

[0112] The weighted moving average method is used to detect the current similarity between the tokens to be evicted in the token key vector set and the tokens in the token key vector set to be retained in the current time step, so as to recall the tokens to be evicted based on the moving average threshold and the current similarity.

[0113] Based on the current similarity, the recalled currently expellable tokens will be merged into the reserved tokens.

[0114] The token key-value vector set is used to obtain the final retained key-value vector.

[0115] In one embodiment, a computer-readable storage medium is provided having a computer program stored thereon, the computer program performing the following steps when executed by a processor:

[0116] Obtain the token key-value cache vector set for each attention layer in the large model;

[0117] The KV cache budget is adjusted according to the density distribution of attention weights in each attention layer, and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer and the remaining tokens are obtained based on the KV cache budget and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer.

[0118] The weighted moving average method is used to detect the current similarity between the tokens to be evicted in the token key vector set and the tokens in the token key vector set to be retained in the current time step, so as to recall the tokens to be evicted based on the moving average threshold and the current similarity.

[0119] Based on the current similarity, the recalled currently expellable tokens will be merged into the reserved tokens.

[0120] The token key-value vector set is used to obtain the final retained key-value vector.

[0121] It should be noted that the functions or steps that can be implemented by the computer-readable storage medium or computer device described above can be referred to the relevant descriptions on the server side and client side in the foregoing method embodiments. To avoid repetition, they will not be described one by one here.

[0122] Those skilled in the art will understand that all or part of the processes in the methods of the above embodiments can be implemented by a computer program instructing related hardware. The computer program can be stored in a non-volatile computer-readable storage medium, and when executed, it can include the processes of the embodiments of the above methods. Any references to memory, storage, databases, or other media used in the embodiments provided in this application can include non-volatile and / or volatile memory. Non-volatile memory can include read-only memory (ROM), programmable ROM (PROM), electrically programmable ROM (EPROM), electrically erasable programmable ROM (EEPROM), or flash memory. Volatile memory can include random access memory (RAM) or external cache memory. By way of illustration and not limitation, RAM is available in various forms, such as static RAM (SRAM), dynamic RAM (DRAM), synchronous DRAM (SDRAM), dual data rate SDRAM (DDRSDRAM), enhanced SDRAM (ESDRAM), synchronous link DRAM (SLDRAM), Rambus direct RAM (RDRAM), direct memory bus dynamic RAM (DRDRAM), and memory bus dynamic RAM (RDRAM), etc.

[0123] Those skilled in the art will clearly understand that, for the sake of convenience and brevity, the above-described division of functional units and modules is used as an example. In practical applications, the above functions can be assigned to different functional units and modules as needed, that is, the internal structure of the device can be divided into different functional units or modules to complete all or part of the functions described above.

[0124] The above-described embodiments are merely illustrative of the technical solutions of the present invention and are not intended to limit it. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features. These modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention, and should all be included within the protection scope of the present invention. Furthermore, any software tools or components not belonging to this company that appear in the embodiments of this application are merely illustrative examples and do not represent actual use.

Claims

1. A key-value caching method, characterized in that, include: Obtain the key-value cache vector set of tokens for each attention layer in the large model; The KV cache budget is adjusted according to the density distribution of attention weights in each attention layer, and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer and the remaining tokens are obtained based on the KV cache budget and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer. The weighted moving average method is used to detect the current similarity between the tokens to be evicted in the token key vector set and the tokens in the token key vector set to be retained in the current time step, so as to recall the tokens to be evicted based on the moving average threshold and the current similarity. Based on the current similarity, the recalled currently expellable tokens are merged into the key-value vector set of retained tokens to obtain the final retained key-value vector; The step of adjusting the KV cache budget based on the density distribution of attention weights in each attention layer includes: The density distribution of attention weights in each attention layer is obtained by calculating the variance of the attention score matrix in each attention layer. The memory budget parameters are determined based on the target compression ratio and the density distribution of the attention weights; Adjust the KV cache budget for each layer based on the aforementioned memory budget parameters; The step of using a weighted moving average method to detect the current similarity between the tokens to be evicted in the key vector set and the tokens in the key vector set to be retained at the current time step, and recalling the tokens to be evicted based on the moving average threshold and the current similarity, includes: The cosine similarity matrix of the set of tokens to be evicted and the set of tokens to be retained at the current time step is calculated using the nearest neighbor matching algorithm. The current maximum similarity between the tokens to be expelled in the set of token key vectors to be expelled and the tokens in the set of token key vectors to be retained is obtained based on the cosine similarity in the cosine similarity matrix. The cosine similarity matrix is ​​smoothed using the exponential moving average method to obtain the moving average threshold for the current time step; Compare the moving average threshold with the current maximum similarity. If the current maximum similarity is not lower than the moving average threshold, recall the current token to be expelled.

2. The key-value caching method as described in claim 1, characterized in that, The method, based on the KV cache budget and the cumulative attention scores of the tokens marked in the key-value cache vector set of each attention layer and the remaining tokens, yields the key-value vector set of tokens to be evicted and the key-value vector set to be retained, including: The token for the current time step in each attention layer is marked based on the sliding window length and the initial token length; Calculate the cumulative attention score for all tokens in the token key-value cache vector set of each attention layer, excluding the labeled tokens; Based on the KV cache budget, sliding window length, initial tag length, and the cumulative attention scores of the remaining tokens, the key-value vector set of tokens to be evicted and the key-value vector set of tokens to be retained are obtained.

3. The key-value caching method as described in claim 2, characterized in that, The process of sorting the tokens based on the KV cache budget, sliding window length, initial tag length, and the cumulative attention scores of the remaining tokens to obtain the set of key-value vectors for tokens to be evicted and the set of key-value vectors for tokens to be retained includes: Based on the difference between the KV cache budget and the sliding window length and the initial tag length, from Select the top N tokens with the highest cumulative attention scores from the remaining tokens; The tokens marked by the sliding window length and the initial token length, along with the first N selected tokens, form a set of retained token key-value vectors. The remaining unselected tokens form a set of token key-value vectors to be evicted.

4. The key-value caching method as described in claim 1, characterized in that, The step of smoothing the cosine similarity matrix using the exponential moving average method to obtain the moving average threshold for the current time step includes: Using formula The moving average threshold for the current time step t is calculated based on the cosine similarity matrix. ;in, Represents the cosine similarity matrix. This represents the smoothing constant.

5. The key-value caching method as described in claim 1, characterized in that, The step of merging the recalled currently expellable tokens into the retained token key-value vector set based on the current similarity to obtain the final retained key-value vector includes: The weights of the currently recalled Tokens to be expelled and the similar Tokens in the key-value vector set of the retained Tokens are calculated based on the current similarity. The key values ​​of the current Token to be expelled and the key values ​​of similar Tokens in the set of key value vectors of the retained Tokens are weighted and merged according to the weight to obtain the final retained key value vector.

6. A key-value caching device, characterized in that, include: The acquisition module is used to obtain the key-value cache vector set of tokens for each attention layer in the large model; The cache adjustment module is used to adjust the KV cache budget according to the density distribution of attention weights in each attention layer, and based on the KV cache budget and the cumulative attention scores of the tokens marked in the Token key-value cache vector set of each attention layer and the remaining tokens, obtain the set of token key-value vectors to be evicted and the set of token key-value vectors to be retained. The recall module is used to detect the current similarity between the tokens to be evicted in the token key vector set and the tokens in the retained token key vector set at the current time step using a weighted moving average method, so as to recall the tokens to be evicted based on the moving average threshold and the current similarity; specifically, it is used to calculate the cosine similarity matrix between the token key vector set to be evicted and the retained token key vector set at the current time step using a nearest neighbor matching algorithm; and obtain the current maximum similarity between the tokens to be evicted in the token key vector set and the tokens in the retained token key vector set based on the cosine similarity in the cosine similarity matrix. The cosine similarity matrix is ​​smoothed using the exponential moving average method to obtain the moving average threshold for the current time step; Compare the moving average threshold with the current maximum similarity; if the current maximum similarity is not lower than the moving average threshold, recall the current token to be expelled. The merging module is used to merge the recalled currently expellable tokens into the retained token key-value vector set according to the current similarity, so as to obtain the final retained key-value vector; Specifically, the cache adjustment module is further configured to obtain the density distribution of attention weights in each attention layer based on the variance of the attention score matrix in each attention layer; determine the memory budget parameters based on the target compression ratio and the density distribution of the attention weights; and adjust the KV cache budget for each layer based on the memory budget parameters.

7. A computer device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that, When the processor executes the computer program, it implements the steps of the key-value caching method as described in any one of claims 1 to 5.

8. A computer-readable storage medium storing a computer program, characterized in that, When the computer program is executed by a processor, it implements the steps of the key-value caching method as described in any one of claims 1 to 5.

Images