Cache fusion-based large language model inference acceleration method, electronic device and storage medium

Abstract

The application discloses a large language model inference acceleration method based on cache fusion, an electronic device and a storage medium. The method realizes the minimization of information loss in the compression process through the merging of the cache, and effectively solves the problems faced in the KV cache compression process. The method comprises the following steps: 1. For a given large language model, the cache value is calculated based on the Token Merge technology, and pruning and merging operations are performed according to the calculated cache value; 2. For the KV cache matrix of the original large language model, a binary mask is generated through Bernoulli distribution for the Token to be discarded; 3. For the KV cache matrix of the large language model to be retained, firstly, the probability value of being merged is calculated through attention weight, and then the discarded Token is evenly merged into the retained Token by using the binary mask of step 2; 4. In the forward propagation process of the large language model, the generated KV cache matrix is continuously subjected to the processing procedures of steps 1 to 3.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of compression and acceleration of artificial neural networks, in particular to related technologies for analyzing theoretical errors using Token Merge strategy, and particularly relates to a large language model inference acceleration method based on cache fusion, an electronic device and a storage medium. BACKGROUND

[0002] The sparsity of deep convolutional neural networks (CNN) refers to the sparsity of neuron activation in the network, that is, only a small part of neurons are activated under any given input. This sparsity is very common in deep convolutional neural networks and is crucial to the performance and efficiency of the network, which can improve the generalization ability of the model, reduce the amount of calculation, reduce the storage demand, and improve the interpretability of the model.

[0003] Currently, the memory load in the inference process of large language models (LLM) is mainly attributed to three main components: the number of model parameters, the size of the activation buffer, and the size of the KV cache. Model compression techniques mainly include parameter quantization and network pruning. Pruning each individual weight level can ignore the performance degradation of the model even at high sparsity rates, which has achieved significant results in addressing the number of model parameters and the size of the activation buffer. In LLM, the memory overhead of its KV cache is related to the hidden dimension of the model, and it is more dependent on the sequence length and batch size in the inference process. The memory required by the KV cache increases proportionally with the persistent storage of previous key-value embeddings in attention to generate new tokens, and in scenarios involving extended sequences, such as sequences spanning 2048 tokens, the memory consumption usually increases to 2.5 to 5 times the size of the model itself, causing a significant burden.

[0004] The existing technology in the research direction of reducing the KV cache burden in LLM inference is usually divided into two categories:

[0005] (1) The first category is to introduce sparsity to discard unnecessary tokens, for example: H2O identifies important tokens based on accumulated attention scores, TOVA suggests using recent attention patterns, and StreamingLLM and LM-Infinite propose the concept of "confluence tokens" (i.e., the beginning of the token sequence, suggesting removing the KV cache of intermediate tokens); in addition, other methods study the attention importance of different attention heads and layers and propose adaptive KV cache eviction strategies.

[0006] (2) The second type is to reduce the number of bits required for each KV embedding, aiming to reduce the number of bits required to store the KV embedding, for example: FlexGen and KIVI and other methods suggest quantizing the key by channel and quantizing the value by token; newer methods such as milestone tokens or gist tokens compress the KV cache cost by setting special tokens, but they all require additional training and cannot be seamlessly integrated into existing LLMs; SparQ enhances KV cache eviction by using global embeddings that average all tokens, but global embeddings that average all tokens can compromise the original output of the current token.

[0007] In summary, the prior art is not perfect in solving the related problems of KV cache compression. SUMMARY

[0008] The purpose of the present application is to provide a large language model inference acceleration method based on cache fusion, an electronic device and a storage medium, to solve the problems in the prior art, to minimize information loss in the compression process by merging caches, and to effectively solve the problems faced in the KV cache compression process.

[0009] In order to achieve the above purpose, the solution of the present application is:

[0010] A large language model inference acceleration method based on cache fusion, comprising:

[0011] Step 1: For a given large language model, calculate its cache value based on Token Merge technology, and perform pruning and merging operations according to the calculated cache value;

[0012] Step 2: For the original KV cache matrix of the large language model, a binary mask is generated through Bernoulli distribution for the Token to be discarded;

[0013] Step 3: For the retained KV cache matrix of the large language model, first calculate the probability value of being merged through attention weight, and then use the binary mask of step 2 to average the discarded Token into the retained Token;

[0014] Step 4: In the forward propagation process of the large language model, the generated KV cache matrix is constantly processed by steps 1 to 3.

[0015] Further, in the step 2, the generated binary mask is:

[0016]

[0017] wherein Bernoulli(·) represents a Bernoulli function; clamp(·) represents a clipping operation; represents an average operation of to ; i and j k ,k∈[1,m] correspond to the discarded Token and the kth reserved Token in the associated KV cache matrix, respectively, and m represents the number of reserved Tokens; and k∈[1,m] represent the accumulated attention scores of the discarded Token and the kth reserved Token in the previous iteration step, respectively.

[0018] Preferably, the calculation formula of the accumulated attention score is respectively:

[0019]

[0020] wherein t represents a time; A i and represent the attention scores of the discarded Token and the reserved Token, respectively.

[0021] Preferably, in step 3, the operation of averaging the discarded Token into the reserved Token corresponds to the calculation formula:

[0022]

[0023] wherein, and represent the cache values of the kth reserved Token before and after the merging, respectively; V i represents the cache value of the discarded Token.

[0024] An electronic device comprising a processor, a memory and an application program; the application program is stored in the memory and is configured to be executed by the processor to perform the cache fusion-based large language model inference acceleration method.

[0025] A computer-readable storage medium, the storage medium stores a computer program; the computer program makes the computer execute the cache fusion-based large language model inference acceleration method when executed in the computer.

[0026] After adopting the above technical solution, the present application provides an effective way for performance improvement of large language models, especially suitable for low-latency memory (LLMs) application scenarios, and has the following outstanding advantages:

[0027] (1) Efficient resource utilization: By using the Token Merge-based pruning method, the neural network can make more efficient use of resources, reducing unnecessary computation and storage overhead, thereby improving overall performance and efficiency.

[0028] (2) Intelligent cache management: By adopting a novel sampling strategy that is influenced by attention scores, the system can dynamically manage the cache while preserving key token information, adapting to different tasks and scenarios; this intelligent cache management can significantly reduce memory usage and help improve the speed and efficiency of the inference process.

[0029] (3) Universality and scalability: The invention has been extensively tested in experiments, covering multiple benchmark tests, including LLaMA, OPT, and GPT-NeoX. This indicates that the invention is not only suitable for specific neural network structures, but also has certain universality and scalability, and can be applied to various deep learning tasks and models.

[0030] (4) No need to maintain key token information: Through reasonable merging and discarding operations, the invention can maintain the integrity and effectiveness of key token information without maintaining their corresponding caches, thereby simplifying the management and maintenance process of the system and reducing the complexity of implementation and deployment. Through a large amount of experimental data verification, the invention is very successful in improving the performance of current pruning compression methods for various KV cache matrices in various benchmark tests. Specifically, when testing the optimized LLaMA-7b model on Openbookqa, the invention improves the accuracy of StreamingLLM (a leading KV cache matrix compression technology) by 1.5%. In addition, under similar sparsity, the improved performance even exceeds the accuracy of the current advanced KV cache matrix compression. BRIEF DESCRIPTION OF DRAWINGS

[0031] Figure 1 Results of the application of the invention to a set of extended existing baselines;

[0032] Figure 2 Performance comparison results of the application of the invention and three baseline methods to each representative model;

[0033] Figure 3 Results for three downstream tasks;

[0034] Figure 4 Efficiency data for the invention;

[0035] Figure 5 Memory usage and inference performance comparison chart of dense attention, StreamingLLM, and the invention;

[0036] Figure 6A variance diagram of attention weights in the sequence generation process of the application;

[0037] Figure 7 A performance change comparison diagram of the full KV cache, Streaming LLM and the KV cache of the application for compressing mainstream LLM. DETAILED DESCRIPTION

[0038] In order to further explain the technical solutions of the application, the application will be described in detail below through specific embodiments.

[0039] The application discloses a large language model inference acceleration method based on cache fusion, comprising:

[0040] Step 1: For a given large language model, calculate its cache value based on Token Merge technology, and perform pruning and merging operations according to the calculated cache value.

[0041] Step 2: For the KV cache matrix of the original large language model, a binary mask is generated for the Token to be discarded through Bernoulli distribution. The generation of this mask matrix is obtained by analyzing and calculating the structure of the neural network, which is convenient for subsequent processing steps.

[0042] Step 3: For the retained KV cache matrix of the large language model, first calculate the probability value of being merged through attention weight, and then use the binary mask of step 2 to average the discarded Token into the retained Token. Thus, it is ensured that even in the case of discarding part of the information, the large language model can still maintain the integrity and effectiveness of the whole, thereby minimizing the performance loss; and this probability value is determined according to the importance of each Token and the overall attention distribution, so that the merging operation is more intelligent and effective.

[0043] Step 4: In the forward propagation process of the large language model, the processing procedures of steps 1 to 3 are constantly performed on the generated KV cache matrix. This means that the application dynamically adjusts and optimizes the cache structure of the network in each round of forward propagation to adapt to the changes in current task requirements and resources, thereby realizing a more efficient and flexible inference process.

[0044] Through the above scheme, the application provides an effective way for performance improvement of large language models, especially suitable for low latency memory (LLMs) application scenarios, and has the following outstanding advantages:

[0045] (1) Efficient resource utilization: By using the pruning method based on Token Merge, the neural network can make more efficient use of resources, reduce unnecessary computation and storage overhead, and improve overall performance and efficiency.

[0046] (2) Intelligent cache management: By adopting a novel sampling strategy that is influenced by attention scores, the system can dynamically manage the cache while preserving key token information, adapting to different tasks and scenarios. This intelligent cache management can significantly reduce memory usage and help improve the speed and efficiency of the inference process.

[0047] (3) Universality and scalability: The invention has been extensively tested in experiments, covering multiple benchmark tests including LLaMA, OPT, and GPT-NeoX. This indicates that the invention is not only suitable for specific neural network structures, but also has certain universality and scalability, and can be applied to various deep learning tasks and models.

[0048] (4) No need to maintain key token information: Through reasonable merging and discarding operations, the invention can maintain the integrity and effectiveness of key token information without maintaining their corresponding cache, thereby simplifying the management and maintenance process of the system and reducing the complexity of implementation and deployment. Through a large amount of experimental data verification, the invention is very successful in improving the performance of current pruning compression methods for various KV cache matrices in various benchmark tests. Specifically, when testing the optimized LLaMA-7b model on Openbookqa, the invention improves the accuracy of StreamingLLM, a leading KV cache matrix compression technology. In addition, under similar sparsity, the improved performance even exceeds the accuracy of current advanced KV cache matrix compression.

[0049] The following shows a specific embodiment of the invention.

[0050] In the above step 2, the generated binary mask is:

[0051]

[0052] where Bernoulli(·) represents the Bernoulli function; clamp(·) represents the truncation operation; represents taking the average value operation from to ; i and j k , k ∈ [1, m] correspond to the discarded Token and the retained kth Token in the associated KV cache matrix, respectively, m represents the token span when considering the average attention of the token range; and k∈[1,m] respectively represent the accumulated attention scores of the discarded Token and the retained kth Token in the previous iteration step, replacing the traditional single-step Token attention score value, which is more robust in predicting future attention, etc.

[0053] Further, the calculation formula of the above-mentioned accumulated attention score is respectively:

[0054]

[0055] Where t represents the time; A i and respectively represent the attention scores of the discarded Token and the retained Token.

[0056] Meanwhile, in the above-mentioned step 3, the operation of averaging the discarded Token into the retained Token corresponds to the calculation formula:

[0057]

[0058] Where, and respectively represent the cache values of the retained kth Token before and after merging; V i represents the cache value of the discarded Token.

[0059] Above, the present application introduces a brand-new cache adaptive merging method, which adopts a unique sampling strategy that is affected by the attention score of the discarded position. This method enables efficient low-latency memory (LLMs) to retain key token information during inference, even without maintaining the corresponding cache. Extensive experiments verify the effectiveness of this method in improving the performance of memory, covering multiple benchmark tests, including LLaMA, OPT and GPT-NeoX. The following experimental data proves the technical effect of the present application:

[0060] Reference Figure 1This paper elucidates the results of applying the present invention to an extended set of existing baselines, including H2O and Scissorhands. These methods selectively retain their previously cached tokens by examining the token's attention score. Since these strategies inherently preserve local tokens, they align with the application strategy of the present invention. Empirical evidence suggests that the present invention consistently improves the inference performance of all methods across a range of tasks. For example, the present invention enhances the performance of H2O, eliminating the memory budget of LLaMA-7b, and achieving improvements of 4.8%, 3.3%, and 9.0% on the OpenBookQA, PiQA, and COPA datasets, respectively. These results demonstrate that the present invention is a general and scalable approach that complements current techniques aimed at improving output variability caused by cache eviction.

[0061] See Figure 2 This study evaluates several representative LLM-based model families, including OPT, LLaMA, and GPT-NeoX, ranging in model size from 7 billion to 65 billion parameters, and selects three baseline methods for comparison: StreamingLLM, H2O, and Scis sorhands. All experiments were performed on NVIDIA Tesla A800 GPUs using five different random seeds, and we report the mean and variance of the results. Because this invention caches and uniformly merges the values ​​to be discarded into all local tokens across all transformer layers, it achieves generality, considering that all existing KV caching compression methods store values ​​in local tokens.

[0062] Since the merging mask is a key factor in this invention, the performance impact of different positional intervals was further investigated. Figure 3 The results for three downstream tasks are presented. Intuitively, increasing the lower bound of the interval leads to an increase in the number of merged tokens, and vice versa. It can be seen that adjusting the upper bound of the merge range can improve the performance of some tasks. However, setting a high lower bound inevitably leads to a performance degradation. These results indicate that selectively merging movable caches is crucial. In this invention, a consistent truncation interval [0,1] is used to ensure that the invention remains a flexible and easy-to-use method. Undoubtedly, these findings highlight that the design of merged tokens is a very promising avenue for future research.

[0063] See Figure 4 Applying this invention in StreamingLLM does not significantly reduce generation speed or increase memory usage. This is intuitive, given that the invention does not introduce any complex matrix operations or iterative loops. Combined with previous experimental results, this invention is a plug-and-play method that effectively enhances the performance of existing methods while achieving memory-efficient LLM.

[0064] See Figure 5 Part (a) shows the attention graph using dense attention; Part (b) shows how StreamingLLM mitigates memory overhead by discarding unnecessary token caches determined by the attention scores of some tokens; Part (c) shows how the present invention merges the cache of tokens to be evicted with higher attention scores (purple boxes) into subsequent tokens. Therefore, as can be seen from the comparison in Part (d), the present invention has less output perturbation caused by cache discarding and achieves superior inference performance (LLaMA-7B with a 20% memory budget).

[0065] See Figure 6 This illustrates the variance of attention weights during sequence generation (Llama-7b on the first sample of the XSUM test set). This invention visualizes tokens divided into multiple consecutive groups, each containing 64 tokens. Taking x-axis 0 as an example, the blue dots represent the average attention variance of the 64 consecutive tokens within the index [0, 64], and the green dots plot the attention variance of these individual tokens. It can be seen that the variance of the average attention weights for consecutive tokens is much smaller.

[0066] See Figure 7 The paper presents a performance comparison of full key-value caching, Streaming LLM, and the present invention on multilingual tasks.

[0067] The present invention also discloses an electronic device, including a processor, a memory, and an application program; the application program is stored in the memory and configured to be executed by the processor using the methods described above.

[0068] The present invention also discloses a computer-readable storage medium storing a computer program; when the computer program is executed in a computer, it causes the computer to perform the above-described method. The storage medium may include various media capable of storing program code, such as a USB flash drive, portable hard drive, read-only memory (ROM), random access memory (RAM), magnetic disk, or optical disk.

[0069] The above embodiments and figures are not intended to limit the product form and style of the present invention. Any appropriate changes or modifications made by those skilled in the art should be considered as not departing from the patent scope of the present invention.

Claims

1. A method for accelerating inference of large-scale language models based on cache fusion, characterized in that... include: Step 1: For a given large language model, calculate its cache value based on Token Merge technology, and perform pruning and merging operations based on the calculated cache value; Step 2: For the existing large language model's key-value cache matrix, generate a binary mask for the tokens that need to be discarded using Bernoulli distribution. Step 3: For the KV cache matrix of the retained large language model, first calculate its probability value of being merged through attention weights, and then use the binary mask in step 2 to merge the discarded tokens into the retained tokens on an average basis. Step 4: During the forward propagation of the large language model, continuously perform the processing steps 1 to 3 on the generated KV cache matrix; In step 2, the generated binary mask for: ； in, Represents the Bernoulli function; Indicates a truncation operation; Indicates taking to Average operation; and The corresponding key-value cache matrix contains the discarded tokens and the retained tokens, respectively. Each Token Indicates the number of tokens that are retained; and These represent the discarded tokens and the retained tokens, respectively. The cumulative attention score of each token in the previous iteration step; The formulas for calculating the cumulative attention score are as follows: ， ； in, Indicates time; and These represent the attention scores for the discarded and retained tokens, respectively.

2. The method for accelerating large-scale language model inference based on cache fusion as described in claim 1, characterized in that: In step 3, the calculation formula corresponding to the operation of merging the discarded tokens into the retained tokens is as follows: ； in, and These represent the reserved first and second digits. The cached values ​​of each token before and after the merge; This represents the cached value of the discarded token.

3. An electronic device, characterized in that... It includes a processor, memory, and an application; the application is stored in the memory and configured to be executed by the processor as described in claim 1 or 2, using a cache-fusion-based method for accelerating large language model inference.

4. A computer-readable storage medium, characterized in that... The storage medium stores a computer program; when the computer program is executed in the computer, it causes the computer to perform the large-scale language model inference acceleration method based on cache fusion as described in claim 1 or 2.

Images