Key-value cache compression and sparse attention computation method and system for large language model inference

Abstract

The present disclosure relates to the technical field of artificial intelligence and natural language processing, in particular to a key-value cache compression and sparse attention calculation method and system for large language model inference. The key-value cache compression and sparse attention calculation method for large language model inference comprises: an offline calibration stage; an online inference stage, comprising: a pre-filling step; an autoregressive generation step, for each newly generated word element: projecting the current query vector Q and the key vector K in the key cache into a low-dimensional space to obtain Q' and K'; calculating the approximate attention score based on Q' and K', and selecting the top k most relevant word element indexes I in descending order; calculating the accurate attention score based on Q and K[I], and calculating the output of the current word element with the value vector V[I]. The above scheme solves the memory and calculation bottleneck of large model inference in the scenario of long text sequence input, and has the advantages of reducing the memory occupation and the calculation complexity at the same time.

Description

TECHNICAL FIELD

[0001] The present disclosure relates to the technical field of artificial intelligence and natural language processing, in particular to a key-value cache compression and sparse attention calculation method and system for large language model inference. BACKGROUND

[0002] The Transformer architecture has become the mainstream structure of current large language models, and the computational complexity of its core self-attention mechanism grows quadratically with the sequence length. In the inference stage, in order to accelerate autoregressive generation, the KV-Cache technology is usually used to cache the key and value vectors of all previous tokens. However, when the context window reaches tens of thousands or even hundreds of thousands of tokens, the memory occupation of KV-Cache rapidly expands, and may even exceed the parameter amount of the model itself. This quadratic expansion is the main bottleneck of inference efficiency, which seriously limits the scalability of LLM (large language model) in practical applications.

[0003] Existing technologies propose two types of improvement schemes: one type only caches the key vector through mathematical equivalent transformation, and calculates the value vector in real time when needed, thereby reducing the context memory by half, but this method still needs to perform full-dimensional attention calculation on all tokens in the generation stage, and the calculation amount is not reduced; the other type projects the key vector into a low-dimensional space for sparse attention calculation, which can reduce the attention calculation time, but this method still needs complete KV-Cache, and the memory occupation is not reduced. These existing technologies cannot solve the two major bottlenecks of high memory occupation and high computational complexity at the same time. SUMMARY

[0004] In order to solve the problems in the related art, the embodiments of the present disclosure provide a key-value cache compression and sparse attention calculation method and system for large language model inference.

[0005] In a first aspect, the embodiments of the present disclosure provide a key-value cache compression and sparse attention calculation method for large language model inference, comprising:

[0006] In the offline calibration stage, principal component analysis is performed on the key vector set obtained from the historical key vector data set, and a dimension reduction projection matrix P is generated and stored;

[0007] In the online inference stage, comprising:

[0008] The pre-filling step receives an input sequence, then parallel processes all tokens of the input sequence, and only caches the corresponding key vector K to the key cache, while the value vector V is transformed by the reversible transformation V = K·W kv in real time generation, wherein W kv is an offline precomputed weight matrix; the input sequence is a long text sequence comprising a plurality of tokens;

[0009] autoregressively generating, for each newly generated token:

[0010] projecting the current query vector Q and the key vectors K in the key cache into a low-dimensional space by the dimensionality reduction projection matrix P to obtain Q' and K', respectively;

[0011] calculating an approximate attention score based on Q' and K' in the low-dimensional space computing unit, and selecting the indexes I of the top-k most relevant tokens in descending order by the hardware sorter based on the approximate attention score;

[0012] generating the value vector V[I] based on the full-dimensional key vector K[I] corresponding to the index I;

[0013] calculating an accurate attention score based on Q and K[I] in the high-dimensional computing unit, and calculating the output of the current token based on the accurate attention score and the value vector V[I].

[0014] In an embodiment of the present disclosure, obtained by offline calculation and pre-stored in the weight storage unit, wherein and are the original projection weight matrices of the keys and values, respectively.

[0015] In an embodiment of the present disclosure, the dimension of the low-dimensional space is 25%-50% of the original dimension.

[0016] In an embodiment of the present disclosure, the value k of the top-k in descending order is set to 12.5%-25% of the current sequence length.

[0017] In an embodiment of the present disclosure, the generating the value vector V[I] based on the full-dimensional key vector K[I] corresponding to the index I includes:

[0018] reading the full-dimensional key vector K[I] from the key cache based on the index I, and performing linear transformation V[I] = K[I]·W kv to generate the value vector V[I].

[0019] In an embodiment of the present disclosure, the approximate attention score calculation and the index screening of the top-k most relevant tokens in descending order are performed in parallel by a GPU, and a sparse matrix multiplication kernel optimization is adopted.

[0020] In an embodiment of the present disclosure, the value vector V is not stored in the pre-filling step, and the GPU video memory space occupied by the value vector V is released, and the released video memory space is dynamically allocated to the key cache for expansion.

[0021] In a second aspect, the disclosure provides a system for implementing the method of any one of the first aspect, comprising:

[0022] an offline calibration module configured to perform principal component analysis on a set of key vectors obtained from a historical key vector dataset, generate and store a dimension reduction projection matrix P;

[0023] a key cache memory configured to store only key vectors K;

[0024] a reversible value generation module configured to generate a value vector V according to a key vector K and a pre-stored weight matrix W kv an instant value vector V;

[0025] a low-dimensional projection module configured to project a query vector Q and a key vector K into a low-dimensional space;

[0026] an index selection module configured to select the most relevant token index according to the approximate attention score;

[0027] a high-dimensional attention calculation module configured to perform full-dimensional accurate attention calculation on the top k tokens sorted from high to low and output the result.

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

[0029] In a fourth aspect, the disclosure provides a computer-readable storage medium having computer instructions stored thereon, wherein the computer instructions are executed by a processor to implement the method of any one of the first aspect.

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

[0031] According to the technical scheme provided by the embodiments of the disclosure, the key-value cache compression and sparse attention calculation method for large language model inference generates a dimension reduction projection matrix through an offline calibration stage, only caches key vectors and dynamically generates value vectors when processing long text sequences in the online inference stage, and combines low-dimensional space sparse attention screening and full-dimensional accurate calculation, thereby realizing the dual optimization of memory occupation and calculation complexity, solving the memory and calculation bottleneck of large model inference in the scenario of long text sequence input, and having the advantages of simultaneously reducing memory occupation and calculation complexity.

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

[0033] Other features, objects, and advantages of the present disclosure will become more apparent from the following detailed description when read in conjunction with the accompanying drawings. The following is a description of the drawings.

[0034] Figure 1 A structural diagram showing attention mechanism optimization according to an embodiment of the present disclosure is shown.

[0035] Figure 2 A flowchart showing a key-value cache compression and sparse attention calculation method for large language model inference according to an embodiment of the present disclosure is shown.

[0036] Figure 3 A structural block diagram showing a key-value cache compression and sparse attention calculation system for large language model inference according to an embodiment of the present disclosure is shown.

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

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

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

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

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

[0042] Definitions of terms involved in the present disclosure:

[0043] LLM: LLM is the abbreviation of "Large Language Model", which is an artificial intelligence model with a huge number of parameters (usually reaching tens of billions or even hundreds of billions). It can generate natural language text through learning and training on large-scale text data, and has the ability of language understanding and generation, and can be applied to various natural language processing tasks such as text generation, machine translation, question and answer systems, etc.

[0044] Transformer: Transformer is a deep learning model architecture based on self-attention mechanism, used for processing sequential data such as natural language processing tasks. It effectively captures long-range dependencies in sequences through parallel computation and dynamic weight distribution, and performs well in tasks such as machine translation and text generation.

[0045] Query, Key, Value: Q represents the query vector, K represents the key vector, and V represents the value vector. These three vectors are obtained by linear transformation from the original input vector (usually word embedding representation), where Q is Query, K is Key, and V is Value.

[0046] KV-Cache: KV-cache (Key-Value Cache) is a technique used to optimize computational efficiency during large language model (LLM) inference. The core idea of KV-cache is to cache the key (Key) and value (Value) vectors calculated when generating previous tokens, and directly reuse these cached vectors when generating subsequent tokens, rather than recalculating, thereby avoiding a large amount of repeated calculations, significantly improving the speed and efficiency of text generation, and reducing computational complexity.

[0047] Token: Token is the basic unit of language, which can be a word, a subword (such as "un-" and "able") or a single character (in some language models), used to divide and represent text into computer-processable form. In multi-turn dialogue, the entire dialogue is also represented as a series of tokens.

[0048] MHA: MHA is the abbreviation of "Multi-Head Attention", which is one of the core mechanisms in the Transformer architecture. It calculates the correlation between elements in the input sequence through multiple attention "heads", capturing diverse dependencies and feature dimensions in the sequence. The output is the concatenation and linear transformation of these multi-head attention results, making the model's understanding of input data more comprehensive and in-depth.

[0049] Attention: Attention is a technique in deep learning that allows the model to dynamically focus on different parts of the input when processing sequential data. It calculates the correlation weight between input elements and the current processing position, allowing the model to selectively focus on important information. In the Transformer architecture, the Attention mechanism is a core component that captures long-range dependencies between elements in the sequence, thereby improving the model's understanding and generation capabilities for sequential data.

[0050] In the prior art, the self-attention mechanism of the Transformer architecture has a quadratic growth in computational complexity with the sequence length, and the KV-Cache technology leads to a sharp increase in memory usage in long context scenarios. For example, when processing an input sequence of tens of thousands of tokens, the memory requirement of the key-value cache can exceed the model parameters themselves. Existing improvement solutions have limitations: the method of only caching the key vector reduces memory usage but does not reduce computational complexity; the sparse attention based on low-dimensional projection speeds up the calculation, but still requires complete storage of the key-value cache. This results in a coexistence of memory bottlenecks and computational efficiency problems when deploying long context models on edge devices.

[0051] The inventors found that there is a reversible mapping relationship between key-value vectors, and the key vector has a low-rank characteristic. By analyzing the mathematical structure of the key vector, a solution combining key-value cache compression and sparse attention is proposed: first, a low-dimensional representation of the key vector is established using principal component analysis, and the value vector storage requirement is eliminated through reversible transformation; second, after screening the key tokens in the low-dimensional space, only a small number of high-correlation tokens are executed for full-dimensional calculation. This idea breaks through the original bottleneck from the storage and calculation dimensions by two dimension reduction operations.

[0052] Therefore, the present application proposes a technical solution of attention mechanism optimization including an offline calibration phase and an online inference phase. The offline calibration phase performs principal component analysis on the historical key vector dataset to generate a dimension reduction projection matrix; the online inference phase includes a pre-filling and a self-recurrence generation step, the pre-filling step only caches the key vector and generates the value vector through reversible transformation in real time, and the self-recurrence generation step projects the query vector and the key vector into a low-dimensional space to calculate the approximate attention score, screens the key tokens, and generates the value vector based on the full-dimensional key vector and performs accurate attention calculation.

[0053] Figure 1 A structural schematic diagram of attention mechanism optimization according to an embodiment of the present disclosure is shown.

[0054] As Figure 1 shown, in the offline calibration phase, a calibration dataset can be input, which can be a historical key vector dataset, for generating the key vector K required by the model. The calibration dataset is input into the large language model LLM to generate the original key vector K. The key vector K is generated and saved, and the key vector K output by the model is used as the input for subsequent principal component analysis PCA processing. The saved key vector K is subjected to principal component analysis and projected into a low-dimensional space to generate a reduced key vector. The parameter file of the PCA-processed dimension reduction projection matrix P is saved for fast implementation of low-dimensional projection in the inference phase.

[0055] In the online inference phase, the current query vector Q and all key vectors K in the K-cache are applied to the dimension reduction projection matrix P for processing to obtain the reduced dimension Q' and K' vectors. Then, the approximate attention scores are calculated using Q' and K'. The top-k most relevant tokens are selected according to the approximate attention scores, that is, the tokens ranked from high to low in the first k. The key vectors K corresponding to the top-k tokens are restored from the low-dimensional space to the original dimension. The final attention scores are calculated using the current Q vector and the restored full-dimensional key vectors K. The final attention scores are output for subsequent model inference.

[0056] Figure 2 A flowchart of a key-value cache compression and sparse attention calculation method for large language model inference according to an embodiment of the present disclosure is shown.

[0057] As shown in Figure 2 , the key-value cache compression and sparse attention calculation method for large language model inference specifically includes the following steps:

[0058] In the offline calibration phase, principal component analysis is performed on the key vector set obtained from the historical key vector dataset to generate and store the dimension reduction projection matrix P;

[0059] In the online inference phase, the following steps are included:

[0060] In the pre-filling step, an input sequence is received, and then all tokens of the input sequence are processed in parallel. Only the corresponding key vectors K are cached to the key cache, and the value vectors V are generated on the fly through the reversible transformation V = K·W kv , where W kv is an offline precomputed weight matrix; the input sequence is a long text sequence including multiple tokens;

[0061] In the autoregressive generation step, for each newly generated token:

[0062] The current query vector Q and the key vectors K in the key cache are projected to a low-dimensional space through the dimension reduction projection matrix P respectively to obtain Q' and K';

[0063] Approximate attention scores are calculated based on Q' and K' in the low-dimensional space calculation unit, and the indexes I of the top-k most relevant tokens ranked from high to low are selected by the hardware sorter according to the approximate attention scores;

[0064] The value vectors V[I] are generated on the fly based on the full-dimensional key vectors K[I] corresponding to the indexes I;

[0065] The exact attention scores are calculated based on Q and K[I] in the high-dimensional calculation unit, and the output of the current token is calculated based on the exact attention scores and the value vectors V[I].

[0066] The key-value cache compression and sparse attention calculation method for large language model inference provided by the present disclosure generates a dimension reduction projection matrix through an offline calibration stage, only caches key vectors and dynamically generates value vectors when processing long text sequences in the online inference stage, combines low-dimensional space sparse attention screening and full-dimensional accurate calculation, realizes the dual optimization of memory occupation and calculation complexity, solves the memory and calculation bottleneck of large model inference in the scene of long text sequence input, and has the advantages of reducing memory occupation and calculation complexity at the same time.

[0067] In the present disclosure, the long text sequence refers to an input sequence with a sequence length ≥10k word units. The principal component analysis (PCA) in the offline calibration stage refers to extracting the main feature direction by statistical history key vector distribution, which can be realized by covariance matrix eigenvalue decomposition, providing data-driven dimension reduction basis for subsequent low-dimensional projection. The reversible transformation in the pre-filling step refers to deriving the value vector from the key vector through linear mapping, so that the value vector does not need to be stored, thereby saving memory space. The low-dimensional space refers to mapping the high-dimensional vector to a subspace that retains the main information, which can be realized by matrix projection operation, reducing the calculation amount of approximate attention calculation. The index screening of the top k most relevant word units from high to low refers to selecting the most relevant word units according to the low-dimensional attention score, which can be realized by sorting algorithm. This step ensures that the subsequent accurate calculation only acts on the high-relevance word units. The accurate attention calculation refers to executing the standard attention mechanism in the full-dimensional space, which can be realized by dot product operation, thereby ensuring the capture accuracy of the key word unit semantics. The low-dimensional space calculation unit refers to a calculation unit that maps the high-dimensional vector to the low-dimensional space, which can be realized by block matrix multiplication and SIMD instruction set optimization, used to reduce the dimension of subsequent attention calculation. The hardware sorter refers to a logic unit that selects the key word unit index based on the approximate attention score, which can be realized by parallel sorting algorithm and atomic operation. Its role is to reduce the number of word units that need to be accurately calculated. The high-dimensional calculation unit refers to a special hardware for executing full-dimensional attention calculation on selected word units, which can be realized by mixed precision calculation and pipeline scheduling technology, ensuring the attention calculation accuracy of the key word unit.

[0068] Specifically, the offline stage extracts principal components by analyzing the historical key vector distribution, generates a dimension reduction projection matrix P for storage in advance. During online inference, the pre-filling step processes all the word tokens in the input sequence, only stores the key vectors in the cache, and does not store the value vectors, which can be transformed instantly by the matrix. During autoregressive generation, the current query vector and the cached key vectors are reduced in dimension by the dimension reduction projection matrix P, and the approximate attention scores of all word tokens are quickly calculated in the low-dimensional space. A small number of highly relevant word token indexes are selected. Then, the full-dimensional key vectors are read from the key cache according to the indexes, and the corresponding value vectors are generated instantly. Finally, the selected word tokens are executed in the original high-dimensional space for accurate attention calculation. This process realizes the optimization of memory and calculation through two-dimensional transformations.

[0069] Through the above technical solutions, the application effectively relieves the memory pressure of long context reasoning while maintaining the accuracy of the model. For a dialogue system that needs to process tens of thousands of word tokens, the key-value cache memory usage can be reduced by about 50%, and the computational complexity is reduced from O(N 2 ) to O(Nd+k 2 ), where d is the dimension of the low-dimensional space, and k is the number of selected word tokens. This allows edge devices to deploy large language models while supporting longer context windows and maintaining real-time response speed.

[0070] The technical solutions of the present disclosure will be described below through a specific embodiment:

[0071] Suppose for a large language model, an input sequence of length 8192 needs to be processed.

[0072] The technical solutions of the present disclosure include the following steps:

[0073] K-cache construction: In the pre-filling stage, the model processes all input word tokens in parallel, but only stores the K matrix and does not store the V matrix. In the decoding stage, for each newly generated word token, a new K vector is calculated and appended to the K-cache, while the V vector is generated instantly by V = K·W kv .

[0074] PCA projection: Assuming the original K vector dimension is 128, the first 64 principal components (about 50%) are retained after PCA dimension reduction. The conversion matrix P of 128*64 is calculated offline.

[0075] Low-rank screening: Apply PCA projection to the current Q vector and all K vectors in the K-cache to obtain 128-dimensional Q' and K' vectors.

[0076] Calculate the approximate attention score: . Assuming that the Top-2048 (25%) word tokens are selected as the most relevant word tokens, and the index I of the Top-2048 word tokens is retained.

[0077] High-dimensional attention score calculation: For the selected 2048 word units, the original 128-dimensional full-dimensional key vector K[I] is obtained based on the index I.

[0078] Calculate the accurate attention score: .

[0079] Finally, apply Softmax to the score, and then perform weighted summation with the corresponding V[I] to obtain the final output.

[0080] Through this process, the storage space originally required for 8192*128*2 (one K and one V) can be reduced to 8192*128 (only K), achieving a 50% memory reduction. At the same time, the attention calculation is reduced from the original 8192 word units to 2048 word units, greatly improving the calculation efficiency. The performance of the model is maintained while significantly improving the efficiency of large language models in processing long sequences.

[0081] In an embodiment of the present disclosure, obtained by offline calculation and pre-stored in the weight storage unit, wherein and are the original projection weight matrices of the key and value respectively.

[0082] In the present disclosure, refers to the inverse matrix of the original projection weight matrix of the key vector, refers to the original projection weight matrix of the value vector, which has the same dimension as to realize matrix multiplication operation.

[0083] In the offline calculation stage, the matrix is pre-generated and stored in the model parameters. In the online inference stage, when the value vector V needs to be generated, the K vector is directly read from the key cache, and the linear transformation K·W kv is performed, and the V vector obtained in this way is mathematically equivalent to the V vector generated by the original projection method. This process avoids the high delay problem of online calculation of inverse matrix, and at the same time ensures the numerical stability of the transformation process through offline pre-calculation.

[0084] Through the above technical solution, the present application significantly reduces the calculation complexity in the online inference stage while ensuring the accuracy of the value vector generation, so that the instant generation process of the value vector only needs to perform a single matrix multiplication operation, thereby reducing the GPU memory occupation and improving the inference efficiency, and at the same time avoiding the delay fluctuation problem caused by online matrix inversion.

[0085] In an embodiment of the present disclosure, the dimension of the low-dimensional space is 25%-50% of the original dimension.

[0086] In the disclosure, the low-dimensional space refers to a compressed space to which the key vector is mapped from the original high-dimensional space through principal component analysis, which can be implemented by a dimension reduction projection matrix P, which is obtained through offline training of a historical key vector data set. The original dimension refers to the dimension of the key vector without dimension reduction processing, which can be embodied as the hidden layer dimension of the attention layer in the neural network.

[0087] In the autoregressive generation step, the current query vector and the key vector in the key cache are projected into the low-dimensional space at the same time. The dimension of the low-dimensional space is set to 25%-50% of the original dimension, for example, when the original dimension is 1024, the low-dimensional space can be an integer value between 256 and 512. The selection of this range of proportions is based on the analysis of the distribution of historical key vector data, which can retain key feature information while significantly reducing the amount of data for subsequent approximate attention calculation.

[0088] Through the above technical solution, the application keeps the effectiveness of approximate attention score calculation while controlling the operation amount of the low-dimensional space within an acceptable range. The dimension setting enables the GPU to efficiently process compressed matrix operations, especially when processing long sequences, which not only reduces the memory bandwidth pressure in the sparse attention screening stage, but also maintains the integrity of the features required for subsequent full-dimensional accurate attention calculation.

[0089] In an embodiment of the disclosure, the k value of the top k from high to low is set to 12.5%-25% of the current sequence length.

[0090] In the disclosure, the top k from high to low refers to the number of most relevant word units selected after approximate attention score calculation, which can be implemented by a sorting algorithm, for example, using a quick selection algorithm or a GPU-based parallel sorting kernel function to filter out the most influential historical word units for the current generated word unit.

[0091] 12.5%-25% of the current sequence length refers to determining the number of candidate word units in proportion to the dynamic length of the input sequence, which can be implemented by calculating the sequence length in real time and multiplying a predetermined proportionality coefficient, for example, when the sequence length is 1024, the k value can be set to an integer between 128 and 256. This proportion range can balance the calculation efficiency and attention accuracy.

[0092] Specifically, in the autoregressive generation step, when the query vector of the newly generated token is completed low-dimensional projection, the system dynamically adjusts the value of k according to the total length of the currently processed sequence. For example, in the initial stage of the sequence, the value of k is proportionally lowered to reduce the computational overhead; as the sequence length increases, the value of k is proportionally increased to maintain the attention coverage. By dynamically associating the value of k with the sequence length, the sparse attention calculation amount can be adaptively controlled under different context lengths, avoiding the problem of oversampling or undersampling caused by a fixed threshold.

[0093] Through the above technical solutions, the present application solves the adaptability defects of the fixed sparsity strategy in the variable-length sequence scenario, so that the attention calculation complexity can linearly grow with the sequence length instead of quadratically, while avoiding the loss of semantic coherence caused by excessive sparsification.

[0094] In an embodiment of the present disclosure, the full-dimensional key vector K[I] corresponding to the index I generates a value vector V[I] in real time, including:

[0095] The full-dimensional key vector K[I] is read from the key cache based on the index I, and a linear transformation V[I] = K[I] · W is performed kv to generate a value vector V[I].

[0096] In the present disclosure, the index I refers to the position identifier of the most relevant token selected by the approximate attention score, which can be stored in the form of an integer array for quick positioning of the target data in the key cache. The full-dimensional key vector refers to the original key vector data without dimension reduction processing, which can be directly read through the GPU memory, and is used to retain complete semantic information.

[0097] Specifically, in the autoregressive generation stage, after obtaining the index selected by the approximate attention score, the full-dimensional key vector corresponding to the position is directly extracted from the key cache. Then, through the pre-stored weight matrix W kv matrix multiplication operation is performed to dynamically generate the corresponding value vector. This process avoids the occupation of the memory by pre-storing the complete value vector, and only performs conversion operations on the selected part of the key vector, significantly reducing the calculation amount. For example, when the sequence length is 1024 and the selection ratio is 25%, only 256 key vectors need to be linearly transformed to generate the corresponding value vector.

[0098] Through the above technical solutions, the present application effectively solves the problem that the memory occupation and calculation efficiency are difficult to balance in the long context reasoning scenario. Through the selective generation mechanism, high-precision calculation is only performed on the key tokens when needed, while maintaining the model performance and realizing the dynamic balance of memory resources and calculation resources.

[0099] In an embodiment of the present disclosure, the approximate attention score calculation and the index screening of the top k most relevant tokens in descending order are performed in parallel by a GPU, and a sparse matrix multiplication kernel optimization is adopted.

[0100] In the present disclosure, GPU parallel execution refers to synchronously processing the attention score calculation and the index screening task by multiple computing units of the GPU, which can be specifically implemented by using thread block division and shared memory optimization in the CUDA programming model. By decomposing large-scale matrix operations into multiple subtasks for parallel processing, the computing throughput is improved. The sparse matrix multiplication kernel optimization refers to a special computing kernel designed for the sparse distribution characteristics of non-zero elements in the attention score matrix. Specifically, it can be implemented by using the sparse matrix-vector multiplication interface in the NVIDIA cuSPARSE library. By skipping the invalid calculation of zero-value elements, the total amount of floating-point operations is reduced.

[0101] Specifically, in the autoregressive generation phase, when the top k most relevant tokens in descending order need to be screened from the approximate attention score matrix after low-dimensional projection, the attention score matrix is divided into multiple sub-blocks and distributed to different streaming multiprocessors of the GPU for parallel processing. Each sub-block independently calculates the local maximum value and generates a candidate index list. Subsequently, all candidate lists of the sub-blocks are merged by a global reduction operation, and finally the indexes of the top k most relevant tokens in descending order are screened. In this process, the sparse matrix multiplication kernel avoids invalid calculation of the all-zero region by performing multiplication and addition operations only on non-zero elements, thereby reducing the calculation delay.

[0102] Through the above technical solutions, the present application can significantly reduce the delay in the autoregressive generation phase in the long context reasoning scenario, so that when deploying a large language model in a resource-constrained environment such as an edge computing device, both high-precision attention calculation and real-time requirements can be met.

[0103] In an embodiment of the present disclosure, the value vector V is not stored in the pre-filling step, and the GPU memory space occupied by the value vector V is released, and the released memory space is dynamically allocated for key cache expansion.

[0104] In the present disclosure, not storing the value vector V means that only the cache data of the key vector K is retained during sequence processing, and the corresponding value vector V is generated in real time when needed through the reversible transformation relationship. This technical means directly reduces the memory occupation by eliminating redundant storage. Releasing the memory space for dynamic allocation means that the memory resources originally used to store the value vector are reallocated to the key vector cache system, so that the key cache capacity can be dynamically expanded with the sequence length, effectively breaking the limitation of fixed memory allocation on the context length.

[0105] Specifically, when the model processes an input sequence, the pre-filling stage generates key vectors K of all word pieces in parallel and stores them in the key cache, while using the offline pre-trained weight matrix A linear transformation is performed on K to generate a value vector V. Since the value vector V can be reconstructed instantaneously by matrix multiplication of K and , the system releases the memory space occupied by the value vector V immediately after completing the current calculation, and uses the recovered memory resources to store the key vectors generated subsequently. This dynamic memory management mechanism enables the key cache capacity to be improved under the condition that the total memory remains unchanged, thereby supporting longer context windows.

[0106] Through the above technical solutions, the application effectively alleviates the memory pressure in the long sequence reasoning scenario, enabling the context length that can be processed under the same hardware conditions to be significantly increased, while maintaining the real-time and accuracy of the calculation process. This method is particularly suitable for edge computing device deployment scenarios and can support long text processing tasks of the order of tens of thousands of word pieces under limited memory resources.

[0107] Figure 3 A structural block diagram of a key-value cache compression and sparse attention calculation system for large language model inference according to an embodiment of the present disclosure is shown. The device can be realized as part or all of an electronic device through software, hardware, or a combination of both.

[0108] As shown in Figure 3 , the key-value cache compression and sparse attention calculation system 300 for large language model inference includes:

[0109] An offline calibration module 310 is configured to perform principal component analysis on a set of key vectors obtained from a historical key vector dataset, generate and store a dimensionality reduction projection matrix P;

[0110] A key cache memory 320 is configured to store only key vectors K;

[0111] A reversible value generation module 330 is configured to generate value vectors V according to key vectors K and a pre-stored weight matrix W kv and calculate value vectors V instantaneously;

[0112] A low-dimensional projection module 340 is configured to project query vectors Q and key vectors K into a low-dimensional space;

[0113] An index selection module 350 is configured to select the most relevant word piece index according to the approximate attention score;

[0114] A high-dimensional attention calculation module 360 is configured to perform full-dimensional accurate attention calculation on the top k word pieces sorted from high to low and output the result.

[0115] According to the technical scheme provided by the embodiment of the present disclosure, the dimension reduction projection matrix is generated in the offline calibration stage, and only the key vector is cached and the value vector is dynamically generated when processing the long text sequence in the online inference stage, and the low-dimensional space sparse attention screening and full-dimensional accurate calculation are combined, so that the memory occupation and the calculation complexity are optimized, the memory and calculation bottlenecks of the large model inference in the scene of the long text sequence input are solved, and the advantages of reducing the memory occupation and the calculation complexity at the same time are achieved.

[0116] In an embodiment of the present disclosure, obtained by offline calculation and pre-stored in the weight storage unit, wherein and are the original projection weight matrices of the keys and values respectively.

[0117] In an embodiment of the present disclosure, the dimension of the low-dimensional space is 25%-50% of the original dimension.

[0118] In an embodiment of the present disclosure, the k value of the top k from high to low is set to 12.5%-25% of the current sequence length.

[0119] In an embodiment of the present disclosure, the approximate attention score calculation and the index screening of the top k most relevant tokens from high to low are executed in parallel by a GPU, and a sparse matrix multiplication kernel optimization is adopted.

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

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

[0122] The key-value cache compression and sparse attention calculation method for large language model inference includes:

[0123] In the offline calibration stage, principal component analysis is performed on the key vector set obtained from the historical key vector data set, and a dimension reduction projection matrix P is generated and stored;

[0124] In the online inference stage, it includes:

[0125] The pre-filling step receives an input sequence, then processes all tokens of the input sequence in parallel, only caches the corresponding key vector K to the key cache, and the value vector V is generated in real time through the reversible transformation V=K·W kv , wherein W kva weight matrix pre-computed offline; the input sequence is a long text sequence comprising a plurality of tokens;

[0126] an autoregressive generation step, for each newly generated token:

[0127] projecting the current query vector Q and the key vectors K in the key cache into a low-dimensional space by the dimensionality reduction projection matrix P, to obtain Q' and K';

[0128] calculating an approximate attention score based on Q' and K' in a low-dimensional space calculation unit, and selecting the top-k most relevant token indices I in descending order based on the approximate attention score by a hardware sorter;

[0129] generating a value vector V[I] based on the full-dimensional key vector K[I] corresponding to the index I;

[0130] calculating an exact attention score based on Q and K[I] in a high-dimensional calculation unit, and calculating the output of the current token based on the exact attention score and the value vector V[I].

[0131] In an embodiment of the present disclosure, the obtained by offline calculation and pre-stored in a weight storage unit, wherein and are the original projection weight matrices of the keys and values, respectively.

[0132] In an embodiment of the present disclosure, the dimension of the low-dimensional space is 25%-50% of the original dimension.

[0133] In an embodiment of the present disclosure, the value k of the top-k in descending order is set to 12.5%-25% of the current sequence length.

[0134] In an embodiment of the present disclosure, the generation of the value vector V[I] based on the full-dimensional key vector K[I] corresponding to the index I comprises:

[0135] reading the full-dimensional key vector K[I] from the key cache based on the index I, and performing linear transformation V[I] = K[I]·W kv to generate the value vector V[I].

[0136] In an embodiment of the present disclosure, the approximate attention score calculation and the index screening of the top-k most relevant tokens in descending order are performed in parallel by a GPU, and a sparse matrix multiplication kernel optimization is adopted.

[0137] In an embodiment of the present disclosure, the value vector V is not stored in the pre-filling step, and the GPU video memory space occupied by the value vector V is released, and the released video memory space is dynamically allocated for the expansion of the key cache.

[0138] Figure 5 A structural diagram showing a computer system suitable for use in implementing the method according to embodiments of the present disclosure is shown.

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

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

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

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

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

[0144] As another aspect, the present disclosure also provides a computer readable storage medium, which can be the computer readable storage medium included in the electronic device or the computer system in the above embodiments, or can exist separately from the device and not be assembled into the device. The computer readable storage medium stores one or more programs for execution by one or more processors to perform the methods described in the present disclosure.

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

Claims

1. A method for key-value caching compression and sparse attention computation for large-scale language model inference, characterized in that, include: During the offline calibration phase, principal component analysis is performed on the set of key vectors obtained from the historical key vector dataset to generate and store the dimension-reduced projection matrix P; The online reasoning phase includes: The pre-filling step receives the input sequence and processes all words in the input sequence in parallel, caching only the corresponding key vector K to the key buffer, while the value vector V undergoes an invertible transformation V = K·W. kv Generates instantly, where W kv The weight matrix is ​​pre-computed offline; the input sequence is a long text sequence including multiple words. The autoregressive generation step, for each newly generated lexical: The current query vector Q and the key vector K in the key cache are projected to the low-dimensional space through the dimensionality reduction projection matrix P to obtain Q′ and K′ respectively. Approximate attention scores are calculated in a low-dimensional space computing unit based on Q' and K', and the index I of the top k most relevant words in descending order is selected by a hardware sorter based on the approximate attention scores. The value vector V[I] is generated instantly based on the full-dimensional key vector K[I] corresponding to index I; The precise attention score is calculated in the high-dimensional computing unit based on Q and K[I], and the output of the current word is calculated based on the precise attention score and the value vector V[I].

2. The method according to claim 1, characterized in that, Obtained through offline calculation and pre-stored in the weight storage unit, where and These are the original projected weight matrices for the keys and values, respectively.

3. The method according to claim 1, characterized in that, The dimension of the low-dimensional space is 25%-50% of the original dimension.

4. The method according to claim 1, characterized in that, The k value of the first k items sorted from high to low is set to 12.5%-25% of the current sequence length.

5. The method according to claim 1, characterized in that, The instantaneous generation of value vector V[I] based on the full-dimensional key vector K[I] corresponding to index I includes: Based on index I, read the full-dimensional key vector K[I] from the key cache and perform a linear transformation V[I] = K[I]·W. kv Generate a value vector V[I].

6. The method according to claim 1, characterized in that, The approximate attention score calculation and the indexing and filtering of the top k most relevant terms sorted from high to low are executed in parallel by GPU and optimized using sparse matrix multiplication kernel.

7. The method according to claim 1, characterized in that, In the pre-filling step, the value vector V is not stored, and the GPU memory space occupied by the value vector V is released. The released memory space is then dynamically allocated to the key cache for expansion.

8. A system for implementing the method according to any one of claims 1 to 7, characterized in that, include: The offline calibration module is used to perform principal component analysis on the set of key vectors obtained from the historical key vector dataset, and generate and store the dimension-reduced projection matrix P; A key cache memory is used to store only the key vector K; The invertible value generation module is used to generate values ​​based on the key vector K and the pre-stored weight matrix W. kv Instantaneous calculation of the value vector V; The low-dimensional projection module is used to project the query vector Q and the key vector K into a low-dimensional space; The index selection module is used to select the most relevant lexical indexes based on the approximate attention score. The high-dimensional attention calculation module is used to perform full-dimensional precise attention calculation on the top k lexical units sorted from high to low and output the results.

9. An electronic device, characterized in that, The method includes a memory and a processor; wherein the memory is used to store one or more computer instructions, wherein the one or more computer instructions are executed by the processor to implement the method according to any one of claims 1-7.

10. A computer-readable storage medium storing computer instructions thereon, characterized in that, When executed by a processor, the computer instructions implement the method described in any one of claims 1-7.

Images