Large-scale language model reasoning acceleration method based on cache fusion, electronic equipment and storage medium

Abstract

The invention discloses a cache fusion-based large-scale language model reasoning acceleration method, electronic equipment and a storage medium, and aims at realizing minimization of information loss in a compression process through a mode of merging caches and effectively solving problems in a KV cache compression process. The method comprises the following steps of: 1, calculating a cache value of a given large language model based on a Token Merge technology, and performing pruning and merging operation according to the calculated cache value; 2, for a KV cache matrix of an original large language model, generating a binary mask through Bernoulli distribution for Token needing to be abandoned; 3, for the reserved KV cache matrix of the large language model, firstly calculating the merged probability value through the attention weight, and then averagely merging the abandoned Token into the reserved Token by using the binary mask in the step 2; and 4, in the forward propagation process of the large language model, continuously executing the processing processes from the step 1 to the step 3 on the generated KV cache matrix.

Description

Technical Field

[0001] The present invention relates to the field of compression and acceleration technology of artificial neural networks, and in particular to related technologies for analyzing theoretical errors using a Token Merge strategy, and more particularly to a large-scale language model reasoning acceleration method, electronic device, and storage medium based on cache fusion. Background Art

[0002] The sparsity of deep convolutional neural networks (CNNs) refers to the sparsity of neuron activations in the network, that is, only a small fraction 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 computation, reduce storage requirements, and improve the interpretability of the model.

[0003] Currently, the memory load of large language models (LLMs) during inference 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, and has achieved significant results in solving the problems of the number of model parameters and the size of the activation buffer. In LLMs, the memory overhead of its KV cache is related to the hidden dimension of the model and is more dependent on the sequence length and batch size during inference. The memory required for the KV cache increases proportionally with the persistent storage of previous key-value embeddings in the attention to generate new tokens. In scenarios involving extended sequences, such as sequences spanning 2048 tokens, memory consumption usually increases to 2.5 to 5 times the size of the model itself, causing a significant burden.

[0004] Existing technologies for reducing the burden of KV caching in LLM reasoning are generally divided into two categories:

[0005] (1) The first category introduces sparsity to discard unnecessary tags. For example, H2O identifies important tags based on the accumulated attention scores, TOVA proposes to utilize the recent attention patterns, while StreamingLLM and LM-Infinite propose the concept of “convergence tag” (i.e., the beginning of the tag sequence, and it is recommended to remove the KV cache of the intermediate tags). 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 category is to reduce the number of bits required for each KV embedding, which aims to reduce the number of bits required to store KV embeddings. For example, methods such as FlexGen and KIVI propose to quantize keys by channel and values ​​by tag. Newer methods such as milestone tagging or Gist tagging compress the KV cache cost by setting special tags, but they all require additional training and cannot be seamlessly integrated into the existing LLM. SparQ enhances KV cache eviction by averaging the global embedding over all tags. However, averaging the global embedding over all tags will jeopardize the original output of the current tag.

[0007] In summary, the existing technology is not perfect enough in solving the problems related to KV cache compression. Summary of the invention

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

[0009] In order to achieve the above object, the solution of the present invention is:

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

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

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

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

[0014] Step 4: During the forward propagation of the large language model, the processing from step 1 to step 3 is continuously performed on the generated KV cache matrix.

[0015] Furthermore, in step 2, the binary mask generated for:

[0016]

[0017] Among them, Bernoulli(·) represents the Bernoulli function; clamp(·) represents the truncation operation; Indicates taking to The average value operation of i and j k ,k∈[1,m] corresponds to the discarded Token and the retained kth Token in its associated KV cache matrix, and m represents the number of retained Tokens; and k∈[1,m] represents the cumulative attention scores of the discarded Token and the retained k-th Token in the previous iteration step.

[0018] Preferably, the calculation formulas for the cumulative attention scores are:

[0019]

[0020] Where, t represents the time; A i and Represent the attention scores of the discarded and retained Tokens respectively.

[0021] Preferably, in step 3, the calculation formula corresponding to the operation of merging the discarded tokens into the retained tokens on average is:

[0022]

[0023] in, and Respectively represent the cache values ​​of the kth Token before and after the merge; V i Indicates the cached value of the discarded Token.

[0024] An electronic device includes a processor, a memory and an application; the application is stored in the memory and is configured to be executed by the processor to accelerate the large-scale language model reasoning based on cache fusion.

[0025] A computer-readable storage medium stores a computer program; when the computer program is executed in the computer, the computer executes the large-scale language model reasoning acceleration method based on cache fusion.

[0026] After adopting the above technical solution, the present invention provides an effective way to improve the performance of large language models, which is particularly suitable for the application scenario of low latency memory (LLMs) and has the following outstanding advantages:

[0027] (1) Efficient resource utilization: By using the Token Merge-based pruning method, the neural network can utilize resources more effectively, reduce unnecessary computing and storage overhead, and thus improve the overall performance and efficiency.

[0028] (2) Intelligent cache management: By adopting a novel sampling strategy that is influenced by the attention score, it is able to dynamically manage the cache while retaining key tag information to adapt to the needs of different tasks and scenarios; this intelligent cache management can significantly reduce memory usage and help improve the speed and efficiency of the reasoning process.

[0029] (3) Universality and scalability: The present invention has been extensively verified in experiments, covering multiple benchmarks, including LLaMA, OPT, and GPT-NeoX. This shows that the present invention is not only applicable to a specific neural network structure, 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 tag information: Through reasonable merging and discarding operations, the present invention can maintain the integrity and validity of key tag information without maintaining their corresponding caches, thereby simplifying the management and maintenance process of the system and reducing the complexity of implementation and deployment. After verification by a large amount of experimental data, the present invention has been very successful in improving the performance of various current KV cache matrix pruning and compression methods in various benchmark tests. Specifically, the present invention improved the accuracy of StreamingLLM (a leading KV cache matrix compression technology) when testing the optimized LLaMA-7b model on Openbookqa. In addition, at similar sparsity rates, the improved performance even exceeds the accuracy of current advanced KV cache matrix compression. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] Figure 1 are the results of applying the present invention to an extended set of existing baselines;

[0032] Figure 2 The performance comparison results of the present invention and three baseline methods applied to each representative model are shown;

[0033] Figure 3 Results for three downstream tasks;

[0034] Figure 4 is the efficiency data of the present invention;

[0035] Figure 5 A comparison chart of memory overhead and inference performance of dense attention, StreamingLLM and the present invention;

[0036] Figure 6Schematic diagram of the variance of attention weights in the sequence generation process of the present invention;

[0037] Figure 7 It is a comparison chart of the performance changes of full KV cache, StreamingLLM and the KV cache used for compressing mainstream LLM of the present invention. DETAILED DESCRIPTION

[0038] In order to further explain the technical solution of the present invention, the present invention is described in detail below through specific embodiments.

[0039] The present invention discloses a large language model reasoning acceleration method based on cache fusion, comprising:

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

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

[0042] Step 3: For the KV cache matrix of the retained large language model, first calculate its probability of being merged by the attention weight, and then use the binary mask of step 2 to merge the discarded tokens into the retained tokens on average. This ensures that even if some information is discarded, the large language model can still maintain its overall integrity and effectiveness, thereby minimizing performance loss; and this probability value is determined based on the importance of each token and the overall attention distribution, making the merging operation more intelligent and effective.

[0043] Step 4: During the forward propagation of the large language model, the generated KV cache matrix is ​​continuously processed from step 1 to step 3. This means that the present invention dynamically adjusts and optimizes the cache structure of the network in each round of forward propagation to adapt to the needs of the current task and changes in resources, thereby achieving a more efficient and flexible reasoning process.

[0044] Through the above scheme, the present invention provides an effective way to improve the performance of large language models, which is particularly suitable for application scenarios of low latency memories (LLMs) and has the following outstanding advantages:

[0045] (1) Efficient resource utilization: By using the Token Merge-based pruning method, the neural network can utilize resources more effectively, reduce unnecessary computing and storage overhead, and thus improve the overall performance and efficiency.

[0046] (2) Intelligent cache management: By adopting a novel sampling strategy that is influenced by the attention score, it is able to dynamically manage the cache while retaining key tag information to adapt to the needs of different tasks and scenarios; this intelligent cache management can significantly reduce memory usage and help improve the speed and efficiency of the reasoning process.

[0047] (3) Universality and scalability: The present invention has been extensively verified in experiments, covering multiple benchmarks, including LLaMA, OPT, and GPT-NeoX. This shows that the present invention is not only applicable to a specific neural network structure, 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 tag information: Through reasonable merging and discarding operations, the present invention can maintain the integrity and validity of key tag information without maintaining their corresponding caches, thereby simplifying the management and maintenance process of the system and reducing the complexity of implementation and deployment. After verification by a large amount of experimental data, the present invention has been very successful in improving the performance of various current KV cache matrix pruning and compression methods in various benchmark tests. Specifically, the present invention improved the accuracy of StreamingLLM (a leading KV cache matrix compression technology) when testing the optimized LLaMA-7b model on Openbookqa. In addition, at similar sparsity rates, the improved performance even exceeds the accuracy of current advanced KV cache matrix compression.

[0049] Specific embodiments of the present invention are shown below.

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

[0051]

[0052] Among them, Bernoulli(·) represents the Bernoulli function; clamp(·) represents the truncation operation; Indicates taking to The average value operation of i and j k ,k∈[1,m] corresponds to the discarded Token and the retained kth Token in its associated KV cache matrix, respectively. m represents the number of retained Tokens, corresponding to the tag span when considering the average attention of the tag range; and k∈[1,m] represents the cumulative attention scores of the discarded Token and the retained k-th Token in the previous iteration step, replacing the traditional single Token attention score in a single step, which is more robust in predicting future attention.

[0053] Furthermore, the calculation formulas for the above cumulative attention scores are:

[0054]

[0055] Where, t represents the time; A i and Represent the attention scores of the discarded and retained Tokens respectively.

[0056] At the same time, in the above step 3, the calculation formula for merging the discarded tokens into the retained tokens on average is:

[0057]

[0058] in, and Respectively represent the cache values ​​of the kth Token before and after the merge; V i Indicates the cached value of the discarded Token.

[0059] Above, the present invention introduces a 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 memories (LLMs) to retain key tag information during reasoning, even without maintaining the corresponding cache. Extensive experiments have verified the effectiveness of this method in improving memory efficiency performance, covering multiple benchmarks including LLaMA, OPT and GPT-NeoX. The following experimental data demonstrates the technical effect of the present invention:

[0060] See also Figure 1, illustrates the results of applying the present invention to an extended set of existing baselines, including H2O and Scissorhands. These methods selectively retain the previous cache of a token by looking at its attention score. Since these strategies essentially retain local tokens, they are consistent with the application strategy of the present invention. Empirical evidence shows that the present invention always improves the reasoning performance of all methods across a range of tasks. For example, the present invention enhances the performance of H2O, eliminates the memory budget of LLaMA-7b, and improves by 4.8%, 3.3%, and 9.0% on the OpenBookQA, PiQA, and COPA datasets, respectively. These results show that the present invention is a general and scalable method that complements current techniques aimed at improving output fluctuations caused by cache eviction.

[0061] See also Figure 2 , evaluate multiple representative LLM-based model series, including OPT, LLaMA, and GPT-NeoX, with model sizes ranging from 7 billion parameters to 65 billion parameters, and select three baseline methods for comparison, including StreamingLLM, H2O, and Scissorhands. All experiments are performed on NVIDIA Tesla A800 GPU, using five different random seeds, from which we report the mean and variance of the results. Because the present invention evenly merges the value cache to be discarded into all local tokens of all converter layers. Considering that all existing KV cache compression methods are stored in local tokens, the present invention can achieve universality.

[0062] Since the merged mask is a key factor in the present invention, the performance impact of different position intervals is further studied. Figure 3 Results for three downstream tasks are shown in . Intuitively, increasing the lower limit of the interval will lead to an increase in the number of merged tokens, and vice versa. It can be seen that adjusting the upper limit of the merge range can improve the performance of some tasks. However, setting a higher lower limit will inevitably lead to performance degradation. These results show that it is crucial to selectively merge movable caches. In the present invention, a consistent truncation interval [0,1] is adopted to ensure that the present invention remains a flexible and easy-to-use method. Undoubtedly, these findings highlight that the design of merged word tokens is a very promising avenue for future research.

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

[0064] See also Figure 5 , where part (a) is an attention graph using dense attention; part (b) shows that StreamingLLM reduces memory overhead by discarding unnecessary token caches determined by the attention scores of some tokens; part (c) shows that the present invention merges the cache of to-be-eviction tokens with higher attention scores (purple boxes) into subsequent tokens. Therefore, from the comparison in part (d), it can be seen that the output disturbance caused by cache discard in the present invention is smaller, and more excellent inference performance is achieved (LLaMA-7B has a 20% memory budget).

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

[0066] See also Figure 7 , which shows the performance comparison of full KV cache, Streaming LLM and the present invention on multi-language tasks.

[0067] The present invention also discloses an electronic device, comprising a processor, a memory and an application program; the application program is stored in the memory and is configured so that the processor executes the above method.

[0068] The present invention also discloses a computer-readable storage medium, which stores a computer program; when the computer program is executed in a computer, the computer executes the above method. The storage medium includes various media that can store program codes, such as a USB flash drive, a mobile hard disk, a read-only memory (ROM), a random access memory (RAM), a magnetic disk, or an optical disk.

[0069] The above embodiments and drawings do not limit the product form and style of the present invention. Any appropriate changes or modifications made by ordinary technicians in the relevant technical field should be deemed to be within the patent scope of the present invention.

Claims

1. A large language model reasoning acceleration method based on cache fusion, characterized in that include: Step 1: For a given large language model, calculate its cache value based on the Token Merge technology, and perform pruning and merging operations based on the calculated cache value; Step 2: For the KV cache matrix of the original large language model, generate a binary mask through Bernoulli distribution for the tokens that need to be discarded; Step 3: For the KV cache matrix of the retained large language model, first calculate its probability of being merged through the attention weight, and then use the binary mask in step 2 to evenly merge the discarded tokens into the retained tokens; Step 4: During the forward propagation of the large language model, the processing from step 1 to step 3 is continuously performed on the generated KV cache matrix.

2. The large-scale language model reasoning acceleration method based on cache fusion according to claim 1, characterized in that: In step 2, the binary mask generated for: Among them, Bernoulli(·) represents the Bernoulli function; clamp(·) represents the truncation operation; Indicates taking to The average value operation of i and j k ,k∈[1,m] corresponds to the discarded Token and the retained kth Token in its associated KV cache matrix, and m represents the number of retained Tokens; and They respectively represent the cumulative attention scores of the discarded Token and the retained k-th Token in the previous iteration step.

3. The large-scale language model reasoning acceleration method based on cache fusion as claimed in claim 2, characterized in that: The calculation formulas for the cumulative attention scores are: Where, t represents the time; A i and A jk Represent the attention scores of the discarded and retained Tokens respectively.

4. The large-scale language model reasoning acceleration method based on cache fusion as claimed in claim 2, characterized in that: In step 3, the calculation formula corresponding to the operation of merging the discarded tokens into the retained tokens on average is: Among them, V jk and Respectively represent the cache values ​​of the kth Token before and after the merge; V i Indicates the cached value of the discarded Token.

5. An electronic device, characterized in that It includes a processor, a memory and an application; the application is stored in the memory and is configured to be executed by the processor as described in any one of claims 1 to 4 for the large-scale language model reasoning acceleration method based on cache fusion.

6. A computer-readable storage medium, characterized in that The storage medium stores a computer program; when the computer program is executed in the computer, the computer executes the large-scale language model reasoning acceleration method based on cache fusion as described in any one of claims 1 to 4.

Images