Key value cache compression and sparse attention calculation method and system for large language model reasoning

Abstract

The invention 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 reasoning, and the method comprises the steps: an offline calibration stage; the online reasoning stage comprises the following steps: a pre-filling step; an autoregression generation step: for each newly generated lexical element, projecting a current query vector Q and a key vector K in a key cache to a low-dimensional space to obtain Q'and K '; calculating an approximate attention score based on Q'and K ', and selecting an index I of the first k most relevant lexical elements which are ranked from high to low; and calculating an accurate attention score based on Q and K [I], and calculating with the value vector V [I] to obtain the output of the current lexical element. According to the scheme, the memory and calculation bottleneck of large model reasoning in a scene of long text sequence input are solved, and the method has the advantages of reducing video memory occupation and calculation complexity at the same time.

Description

Technical Field

[0001] This disclosure relates to the fields of artificial intelligence and natural language processing technology, specifically to a key-value caching compression and sparse attention calculation method and system for reasoning in large-scale language models. Background Technology

[0002] The Transformer architecture has become the mainstream structure for large language models, but the computational complexity of its core self-attention mechanism increases quadratically with sequence length. During the inference phase, to accelerate autoregressive generation, a KV-Cache technique is typically used to cache the key and value vectors of all previous tokens. However, when the context window reaches tens or even hundreds of thousands of tokens, the memory usage of the KV-Cache rapidly expands, potentially exceeding the number of parameters in the model itself. This quadratic expansion is the main bottleneck for inference efficiency, severely limiting the scalability of LLMs (Large Language Models) in practical applications.

[0003] Existing technologies propose two types of improvement schemes: one type caches only the key vectors through mathematical equivalence transformations, calculating the value vectors on demand, thus halving the context memory. However, this method still requires full-dimensional attention calculations for all lexical units during the generation phase, so the computational load is not reduced. The other type projects the key vectors into a low-dimensional space for sparse attention calculations, which can reduce the attention calculation time, but this method still requires a complete KV-Cache, so the GPU memory usage is not reduced. These existing technologies cannot simultaneously solve the two major bottlenecks of high memory consumption and high computational complexity. Summary of the Invention

[0004] To address the problems in the related technologies, this disclosure provides a method and system for key-value caching compression and sparse attention computation for large-scale language model inference.

[0005] In a first aspect, embodiments of this disclosure provide a method for key-value caching compression and sparse attention computation for large-scale language model inference, including: 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].

[0006] In one embodiment of this disclosure, 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.

[0007] In one embodiment of this disclosure, the dimension of the low-dimensional space is 25%-50% of the original dimension.

[0008] In one embodiment of this disclosure, the k value of the first k items sorted from high to low is set to 12.5%-25% of the current sequence length.

[0009] In one embodiment of this disclosure, the instantaneous generation of the 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].

[0010] In one embodiment of this disclosure, the approximate attention score calculation and the indexing and filtering of the top k most relevant terms sorted from high to low are performed in parallel by a GPU and optimized using a sparse matrix multiplication kernel.

[0011] In one embodiment of this 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 to the key cache for expansion.

[0012] Secondly, embodiments of this disclosure provide a system for implementing the method described in any one of the first aspects, comprising: 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.

[0013] Thirdly, embodiments of this disclosure provide an electronic device including 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 as described in any of the first aspects.

[0014] Fourthly, this disclosure provides a computer-readable storage medium having computer instructions stored thereon, which, when executed by a processor, implement the method as described in any of the first aspects.

[0015] The technical effects provided by the embodiments of this disclosure may include the following beneficial effects: According to the technical solution provided in this disclosure, a key-value caching compression and sparse attention calculation method for large-scale language model inference generates a dimensionality-reduced projection matrix in the offline calibration stage, and caches only the key vector and dynamically generates the value vector when processing long text sequences in the online inference stage. By combining low-dimensional space sparse attention filtering and full-dimensional accurate calculation, it achieves dual optimization of memory usage and computational complexity, solves the memory and computational bottlenecks of large model inference in scenarios with long text sequence input, and has the advantage of simultaneously reducing memory usage and computational complexity.

[0016] It should be understood that the above general description and the following detailed description are exemplary and explanatory only, and are not intended to limit this disclosure. Attached Figure Description

[0017] Other features, objects, and advantages of this disclosure will become more apparent from the following detailed description of non-limiting embodiments, taken in conjunction with the accompanying drawings. The following is a description of the accompanying drawings.

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

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

[0020] Figure 3A block diagram of a key-value caching compression and sparse attention computing system for large-scale language model inference according to an embodiment of the present disclosure is shown.

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

[0022] Figure 5 A schematic diagram of the structure of a computer system suitable for implementing the method according to embodiments of the present disclosure is shown. Detailed Implementation

[0023] In the following, exemplary embodiments of the present disclosure will be described in detail with reference to the accompanying drawings to enable those skilled in the art to readily implement them. Furthermore, for clarity, portions unrelated to the description of exemplary embodiments have been omitted from the drawings.

[0024] In this disclosure, it should be understood that terms such as “comprising” or “having” are intended to indicate the presence of features, figures, steps, behaviors, components, parts or combinations thereof disclosed in this specification, and are not intended to exclude the possibility of the presence or addition of one or more other features, figures, steps, behaviors, components, parts or combinations thereof.

[0025] It should also be noted that, unless otherwise specified, the embodiments and features described in this disclosure can be combined with each other. This disclosure will now be described in detail with reference to the accompanying drawings and embodiments.

[0026] Definitions of terms used in this disclosure: LLM: LLM is an abbreviation for "Large Language Model". It is an artificial intelligence model with a massive number of parameters (usually reaching billions or even hundreds of billions). Through learning and training on large-scale text data, it can generate natural language text, has the ability to understand and generate language, and can be applied to a variety of natural language processing tasks, such as text generation, machine translation, question answering systems, etc.

[0027] Transformer: Transformer is a deep learning model architecture based on self-attention mechanism, used to process sequential data (such as natural language processing tasks). Through parallel computation and dynamic weight allocation, it can effectively capture long-distance dependencies in sequences and performs well in tasks such as machine translation and text generation.

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

[0029] 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 and value vectors calculated when generating tokens, and then reuse these cached vectors when generating subsequent tokens, instead of recalculating them. This avoids a lot of redundant calculations, significantly improving the speed and efficiency of text generation and reducing computational complexity.

[0030] Tokens: Tokens are the basic units of language. They can be a word, a subword (such as "un-", "able"), or a single character (in some language models) used to segment and represent text in a computer-processable form. In multi-turn dialogues, the entire dialogue is also represented as a series of tokens.

[0031] MHA: MHA is an abbreviation for "Multi-Head Attention," one of the core mechanisms in the Transformer architecture. By simultaneously calculating the correlations between elements in the input sequence on multiple different attention "heads," it can capture diverse dependencies and feature dimensions in the sequence. The output is a concatenation and linear transformation of these multi-head attention results, enabling the model to have a more comprehensive and in-depth understanding of the input data.

[0032] Attention: Attention is a technique in deep learning that allows models to dynamically focus on different parts of the input when processing sequential data. It enables the model to selectively focus on important information by calculating the relevance weights between input elements and the current processing position. In the Transformer architecture, the Attention mechanism is a core component used to capture long-distance dependencies between elements in a sequence, thereby improving the model's understanding and generation capabilities of sequential data.

[0033] In existing technologies, the computational complexity of the self-attention mechanism in the Transformer architecture increases quadratically with sequence length, and the KV-Cache technique leads to a surge in GPU memory usage in long-context scenarios. For example, when processing input sequences with tens of thousands of tokens, the GPU memory requirement for the key-value cache may exceed that of the model parameters themselves. Existing improvements have limitations: methods that only cache key vectors reduce memory usage but do not reduce computational complexity; sparse attention based on low-dimensional projection accelerates computation but still requires complete storage of the key-value cache. This results in both GPU memory bottlenecks and computational efficiency issues when deploying long-context models on edge devices.

[0034] The inventors discovered an invertible mapping relationship between key-value vectors and that key vectors possess low-rank properties. By analyzing the mathematical structure of key vectors, they proposed a scheme combining key-value caching compression with sparse attention: First, principal component analysis is used to establish a low-dimensional representation of key vectors, eliminating the need for value vector storage through invertible transformations; second, after filtering key terms in the low-dimensional space, full-dimensional computation is performed only on a small number of highly relevant terms. This approach overcomes the original bottlenecks in both storage and computation through two dimensionality reduction operations.

[0035] Therefore, this application proposes a technical solution for optimizing the attention mechanism, including an offline calibration stage and an online inference stage. The offline calibration stage performs principal component analysis on the historical key vector dataset to generate a dimensionality-reduced projection matrix. The online inference stage includes pre-filling and autoregressive generation steps. The pre-filling step only caches the key vectors and generates value vectors instantly through an invertible transformation. The autoregressive generation step projects the query vector and key vectors into a low-dimensional space to calculate approximate attention scores. After filtering keyword elements, value vectors are generated based on the full-dimensional key vectors, and precise attention calculation is performed.

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

[0037] like Figure 1 As shown, in the offline calibration phase, a calibration dataset is first input, which can be a historical key vector dataset, to generate the key vectors K required by the model. The calibration dataset is then input into the large language model (LLM) to generate the original key vectors K. The key vectors K are generated and saved, and the key vectors K output by the model are used as input for subsequent principal component analysis (PCA). PCA is performed on the saved key vectors K, projecting them into a low-dimensional space to generate dimensionality-reduced key vectors. The parameter file of the dimensionality-reduced projection matrix P obtained from the PCA process is saved for rapid low-dimensional projection during the inference phase.

[0038] During the online inference phase, a dimensionality-reduced projection matrix P is applied to the current query vector Q and all key vectors K in the K-cache to obtain the dimensionality-reduced Q′ and K′ vectors. Then, an approximate attention score is calculated using Q′ and K′. Based on the approximate attention scores, the top-k most relevant terms are selected, i.e., the k terms with the highest approximate attention scores in descending order. The key vectors K corresponding to the top-k terms are then restored from the low-dimensional space to their original dimensions. The final attention score is calculated using the current Q vector and the restored full-dimensional key vector K. The final attention score is output for subsequent model inference.

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

[0040] like Figure 2 As shown, the key-value caching compression and sparse attention calculation method for large-scale language model inference specifically includes the following steps: 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].

[0041] The key-value caching compression and sparse attention computation method for large-scale language model inference disclosed herein generates a dimensionality-reduced projection matrix during the offline calibration stage, and caches only the key vectors and dynamically generates the value vectors when processing long text sequences during the online inference stage. By combining low-dimensional space sparse attention filtering and full-dimensional accurate computation, it achieves dual optimization of memory usage and computational complexity, solves the memory and computational bottlenecks of large model inference in scenarios with long text sequence input, and has the advantage of simultaneously reducing memory usage and computational complexity.

[0042] In this disclosure, a long text sequence refers to an input sequence with a length ≥ 10k words. Principal Component Analysis (PCA) in the offline calibration stage extracts the main feature directions by statistically analyzing the historical key vector distribution. This can be achieved using covariance matrix eigenvalue decomposition, providing a data-driven dimensionality reduction basis for subsequent low-dimensional projection. The reversible transformation in the pre-filling step derives value vectors from key vectors through linear mapping, eliminating the need to store value vectors and saving GPU memory. The low-dimensional space refers to mapping high-dimensional vectors to a subspace that retains the main information. This can be achieved using matrix projection operations, reducing the computational cost of approximate attention calculations. The indexing and filtering of the top k most relevant words, sorted from highest to lowest, selects the most relevant words based on their low-dimensional attention scores. This can be implemented using a sorting algorithm, ensuring that subsequent precise calculations only apply to highly relevant words. Precise attention calculation executes a standard attention mechanism across the entire dimensional space. This can be achieved using dot product operations, ensuring the accuracy of capturing keyword semantics. The low-dimensional computation unit refers to the computational unit that maps high-dimensional vectors to a low-dimensional space. Specifically, it can be implemented using block matrix multiplication and SIMD instruction set optimizations to reduce the dimensionality of subsequent attention calculations. The hardware sorter is a logical unit that filters keyword token indices based on approximate attention scores. Specifically, it can be implemented using parallel sorting algorithms and atomic operations, reducing the number of tokens requiring precise attention calculations. The high-dimensional computation unit refers to dedicated hardware that performs full-dimensional attention calculations on selected tokens. Specifically, it can be implemented using mixed-precision computation and pipelined scheduling techniques to ensure the accuracy of keyword token attention calculations.

[0043] Specifically, in the offline stage, principal components are extracted by analyzing the distribution of historical key vectors, and a dimensionality-reduced projection matrix P is generated and stored for later use. During online inference, the pre-filling step processes all words in the input sequence, storing only the key vectors in the cache; the value vectors are not stored and can be instantaneously transformed using matrix transformation. During autoregressive generation, the current query vector and the cached key vectors are dimensionality-reduced using the projection matrix P. Approximate attention scores for all words are then quickly calculated in the low-dimensional space, and a small number of highly relevant word indices are selected. Subsequently, the full-dimensional key vectors are read from the key cache based on the indices, and the corresponding value vectors are generated instantaneously. Finally, precise attention calculation is performed on the selected words in the original high-dimensional space. This process achieves dual optimization of memory and computation through two dimensionality transformations.

[0044] Through the above technical solution, this application effectively alleviates the GPU memory pressure of long context reasoning while maintaining model accuracy. For dialogue systems that need to process tens of thousands of words, it can reduce the GPU memory usage of key-value cache by about 50%, while reducing the computational complexity from O(N) to O(N). 2 Reduced to O(Nd+k) 2), where d is the dimension of the low-dimensional space and k is the number of selected lexical units. This enables the deployment of large language models on edge devices to support longer context windows while maintaining real-time response speed.

[0045] The technical solution of this disclosure is illustrated below through a specific embodiment: Suppose that a large language model needs to process an input sequence of length 8192.

[0046] The technical solution disclosed herein includes the following steps: K-cache construction: In the pre-filling stage, the model processes all input tokens in parallel, but only stores the K matrix, not the V matrix. In the decoding stage, for each newly generated token, a new K vector is calculated and appended to the K-cache, while the V vector is obtained through V = K·W. kv Generate instantly.

[0047] PCA Projection: Assuming the original K vector has a dimension of 128, PCA dimensionality reduction retains the first 64 principal components (approximately 50%). An offline calculation yields a 128*64 transformation matrix P.

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

[0049] Calculate the approximate attention score: Suppose that the top 2048 (25%) terms are selected as the most relevant terms, and the index I of the top 2048 terms is retained.

[0050] High-dimensional attention score calculation: For the selected 2048 lexical units, obtain the original 128-dimensional full-dimensional key vector K[I] based on index I. Calculate the precise attention score: .

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

[0052] This process reduces the storage space required from 8192*128*2 (one for K and one for V) to 8192*128 (K only), achieving a 50% reduction in memory usage. Simultaneously, the attention computation is reduced from 8192 lexical units to 2048, significantly improving computational efficiency. This demonstrates a substantial improvement in the efficiency of large language models when processing long sequences, while maintaining overall model performance.

[0053] In one embodiment of this disclosure, 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.

[0054] In this disclosure method, It refers to the inverse of the original projected weight matrix of the key vectors. This refers to the original projected weight matrix of the value vector, whose dimensions are... Maintain consistency to achieve matrix multiplication operations.

[0055] During the offline computing phase, pre-generated The matrix is ​​then stored in the model parameters. During the online inference phase, when a value vector V needs to be generated, the K vector is read directly from the key cache, and K·W is executed. kv The linear transformation yields a V vector that is mathematically equivalent to the V vector generated by the original projection method. This process avoids the high latency problem of online computation of the inverse matrix, while ensuring the numerical stability of the transformation process through offline pre-computation.

[0056] Through the above technical solution, this application significantly reduces the computational complexity of the online inference stage while ensuring the accuracy of value vector generation. This allows the real-time generation process of value vectors to perform only a single matrix multiplication operation, thereby reducing GPU memory usage and improving inference efficiency, while avoiding latency fluctuations caused by online matrix inversion.

[0057] In one embodiment of this disclosure, the dimension of the low-dimensional space is 25%-50% of the original dimension.

[0058] In this disclosure, the low-dimensional space refers to the compressed space to which the key vectors are mapped from the original high-dimensional space through principal component analysis. Specifically, this can be achieved using a dimension-reduced projection matrix P, which is obtained through offline training on a historical key vector dataset. The original dimension refers to the dimension of the key vectors before dimension reduction, which can be specifically represented as the hidden layer dimension of the attention layer in a neural network.

[0059] In the autoregressive generation step, the current query vector and the key vectors in the key cache are simultaneously projected into a low-dimensional space. The dimension of this 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. This range is chosen based on an analysis of the distribution of historical key vector data, preserving key feature information while significantly reducing the amount of data required for subsequent approximate attention calculations.

[0060] Through the above technical solution, this application maintains the effectiveness of approximate attention score calculation while keeping the computational load in low-dimensional space within an acceptable range. This dimensional setting enables the GPU to efficiently process compressed matrix operations, especially when processing long sequences. It reduces the memory bandwidth pressure during the sparse attention selection stage while maintaining the feature integrity required for subsequent full-dimensional accurate attention calculation.

[0061] In one embodiment of this disclosure, the k value of the first k items sorted from high to low is set to 12.5%-25% of the current sequence length.

[0062] In this disclosed method, the top k from high to low refers to the number of most relevant lexical units selected after the approximate attention score calculation. Specifically, it can be implemented through a sorting algorithm, such as using a fast selection algorithm or a GPU-based parallel sorting kernel function, to filter out the historical lexical units that have the greatest impact on the currently generated lexical units.

[0063] The 12.5%-25% of the current sequence length refers to determining the number of candidate lexical units proportionally based on the dynamic length of the input sequence. Specifically, this can be achieved by calculating the sequence length in real time and multiplying it by a preset proportional coefficient. For example, when the sequence length is 1024, the k value can be set to an integer between 128 and 256. This proportional range can balance computational efficiency and attention accuracy.

[0064] Specifically, in the autoregressive generation step, after the query vector of the newly generated tokens has completed low-dimensional projection, the system dynamically adjusts the k value based on the total length of the currently processed sequence. For example, in the initial stage of the sequence, the lower limit of the k value is set proportionally to reduce computational overhead; as the sequence length increases, the upper limit of the k value is set proportionally to maintain the attention coverage. By dynamically associating the k value with the sequence length, the computational cost of sparse attention can be adaptively controlled under different context lengths, avoiding oversampling or undersampling problems caused by a fixed threshold.

[0065] Through the above technical solution, this application solves the adaptability defects of the fixed sparsity strategy in variable length sequence scenarios, so that the attention computational complexity can grow linearly with the sequence length rather than quadratically, while avoiding the loss of semantic coherence caused by excessive sparsity.

[0066] In one embodiment of this disclosure, the instantaneous generation of the 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].

[0067] In this disclosed method, index I refers to the position identifier of the most relevant lexical unit selected through approximate attention scores. Specifically, it can be stored in the form of an integer array for fast location of target data in the key cache. The full-dimensional key vector refers to the raw key vector data without dimensionality reduction processing, which can be directly read through GPU memory to preserve complete semantic information. Specifically, in the autoregressive generation stage, after obtaining the index selected by approximate attention scores, the full-dimensional key vector at the corresponding position is directly extracted from the key cache. Then, the pre-stored weight matrix W... kv Matrix multiplication is performed to dynamically generate the corresponding value vectors. This process avoids the memory occupation of pre-stored complete value vectors and performs transformation operations only on a portion of the selected key vectors, significantly reducing the computational load. 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 vectors.

[0068] Through the above technical solution, this application effectively solves the problem of balancing memory usage and computational efficiency in long-context reasoning scenarios. By employing a selective generation mechanism, high-precision calculations are performed on keyword elements only when needed, achieving a dynamic balance between memory and computational resources while maintaining model performance.

[0069] In one embodiment of this disclosure, the approximate attention score calculation and the indexing and filtering of the top k most relevant terms sorted from high to low are performed in parallel by a GPU and optimized using a sparse matrix multiplication kernel.

[0070] In this disclosure, GPU parallel execution refers to distributing the attention score calculation and index filtering tasks to multiple computing units on the GPU for synchronous processing. Specifically, this can be achieved using thread block partitioning and shared memory optimization in the CUDA programming model. By breaking down large-scale matrix operations into multiple subtasks for parallel processing, computational throughput is improved. Sparse matrix multiplication kernel optimization refers to a dedicated computing kernel designed for the sparse distribution of non-zero elements in the attention score matrix. Specifically, this can be implemented using the sparse matrix-vector multiplication interface in the NVIDIA cuSPARSE library. By skipping invalid calculations of zero-value elements, the total number of floating-point operations is reduced.

[0071] Specifically, in the autoregressive generation stage, when it is necessary to select the top k most relevant terms from the approximate attention score matrix after low-dimensional projection, sorted from highest to lowest, 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 local maxima and generates a candidate index list. Subsequently, the candidate lists of all sub-blocks are merged through a global reduction operation, and finally, the indices of the top k most relevant terms globally sorted from highest to lowest are selected. In this process, the sparse matrix multiplication kernel avoids invalid calculations on all-zero regions by performing multiplication and addition operations only on non-zero elements, thereby reducing computational latency.

[0072] Through the above technical solution, this application can significantly reduce the latency of the autoregressive generation stage in long context reasoning scenarios, so that when deploying large language models in resource-constrained environments such as edge computing devices, it can maintain high-precision attention calculation and meet real-time requirements.

[0073] In one embodiment of this 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 to the key cache for expansion.

[0074] In this disclosed method, not storing the value vector V means that only the cached data of the key vector K is retained during sequence processing. The corresponding value vector V is generated in real time when needed through a reversible transformation relationship. This technique directly reduces video memory usage by eliminating redundant storage. Dynamically allocating and releasing video memory space means reallocating the memory resources originally used to store the value vector to the key vector cache system, so that the key cache capacity can dynamically expand with the sequence length, effectively breaking through the limitation of fixed memory allocation on the context length.

[0075] Specifically, during the model's processing of the input sequence, the pre-filling stage generates key vectors K for all words through parallel computation and stores them in the key cache, while utilizing the weight matrix pre-trained offline. Perform a linear transformation on K to generate a value vector V. Since the value vector V can be transformed through K and... The system performs matrix multiplication and reconstruction instantly. After completing the current calculation, it immediately releases the video memory space occupied by the value vector V and uses the reclaimed memory resources to store subsequently generated key vectors. This dynamic memory management mechanism increases the key cache capacity without changing the total video memory, thus supporting longer context windows.

[0076] Through the above technical solution, this application effectively alleviates the memory pressure in long sequence reasoning scenarios, significantly increasing the length of context that can be processed under the same hardware conditions, while maintaining the real-time performance and accuracy of the computation process. This method is particularly suitable for edge computing device deployment scenarios, enabling the support of long text processing tasks with tens of thousands of tokens under limited memory resources.

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

[0078] like Figure 3 As shown, the key-value caching compression and sparse attention computing system 300 for large-scale language model inference includes: The offline calibration module 310 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; Key cache memory 320 is used to store only the key vector K; The invertible value generation module 330 is used to generate invertible 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 340 is used to project the query vector Q and the key vector K into a low-dimensional space. The index selection module 350 is used to select the most relevant lexical index based on the approximate attention score; The high-dimensional attention calculation module 360 ​​is used to perform full-dimensional precise attention calculation on the top k lexical units sorted from high to low and output the results.

[0079] According to the technical solution provided in this disclosure, a dimension-reduced projection matrix is ​​generated through an offline calibration stage, and only the key vector is cached and the value vector is dynamically generated when processing long text sequences in the online inference stage. By combining low-dimensional space sparse attention filtering and full-dimensional accurate calculation, dual optimization of memory usage and computational complexity is achieved, which solves the memory and computational bottleneck of large model inference in the scenario of long text sequence input, and has the advantage of reducing memory usage and computational complexity at the same time.

[0080] In one embodiment of this disclosure, 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.

[0081] In one embodiment of this disclosure, the dimension of the low-dimensional space is 25%-50% of the original dimension.

[0082] In one embodiment of this disclosure, the k value of the first k items sorted from high to low is set to 12.5%-25% of the current sequence length.

[0083] In one embodiment of this disclosure, the approximate attention score calculation and the indexing and filtering of the top k most relevant terms sorted from high to low are performed in parallel by a GPU and optimized using a sparse matrix multiplication kernel.

[0084] This 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.

[0085] like Figure 4 As 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 embodiments of the present disclosure.

[0086] The key-value caching compression and sparse attention computation method for large-scale language model inference includes: 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].

[0087] In one embodiment of this disclosure, the 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.

[0088] In one embodiment of this disclosure, the dimension of the low-dimensional space is 25%-50% of the original dimension.

[0089] In one embodiment of this disclosure, the k value of the first k items sorted from high to low is set to 12.5%-25% of the current sequence length.

[0090] In one embodiment of this disclosure, the instantaneous generation of the 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].

[0091] In one embodiment of this disclosure, the approximate attention score calculation and the indexing and filtering of the top k most relevant terms sorted from high to low are performed in parallel by a GPU and optimized using a sparse matrix multiplication kernel.

[0092] In one embodiment of this 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 to the key cache for expansion.

[0093] Figure 5 A schematic diagram of the structure of a computer system suitable for implementing the method according to embodiments of the present disclosure is shown.

[0094] like Figure 5 As shown, the computer system includes a processing unit that can execute various methods described above based on a program stored in a read-only memory (ROM) or a program loaded from a storage portion into a random access memory (RAM). The RAM also stores various programs and data required for the operation of the computer system. The processing unit, ROM, and RAM are interconnected via a bus. Input / output (I / O) interfaces are also connected to the bus.

[0095] The following components are connected to the I / O interface: input sections including keyboards, mice, etc.; output sections including cathode ray tubes (CRTs), liquid crystal displays (LCDs), and speakers; storage sections including hard disks; and communication sections including network interface cards such as LAN cards and modems. The communication section performs communication processes via a network such as the Internet. Drives are also connected to the I / O interface as needed. Removable media, such as disks, optical disks, magneto-optical disks, semiconductor memories, etc., are installed on the drive as needed so that computer programs read from them can be installed into the storage section as required. The processing unit can be implemented as a CPU, GPU, TPU, FPGA, NPU, etc.

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

[0097] The flowcharts and block diagrams in the accompanying drawings illustrate the architecture, functionality, and operation of possible implementations of systems, methods, and computer program products according to various embodiments of this disclosure. In this regard, each block in a flowchart or block diagram may represent a module, segment, or portion of code containing one or more executable instructions for implementing a specified logical function. It should also be noted that in some alternative implementations, the functions indicated in the blocks may occur in a different order than those indicated in the drawings. For example, two consecutively indicated blocks may actually be executed substantially in parallel, and they may sometimes be executed in reverse order, depending on the functions involved. It should also be noted that each block in the block diagrams and / or flowcharts, and combinations of blocks in the block diagrams and / or flowcharts, can be implemented using a dedicated hardware-based system that performs the specified function or operation, or using a combination of dedicated hardware and computer instructions.

[0098] The units or modules described in the embodiments of this disclosure can be implemented in software or programmable hardware. The described units or modules can also be located in a processor, and the names of these units or modules do not necessarily constitute a limitation on the unit or module itself.

[0099] In another aspect, this disclosure also provides a computer-readable storage medium, which may be a computer-readable storage medium included in the electronic device or computer system described above; or it may be a standalone computer-readable storage medium not assembled into a device. The computer-readable storage medium stores one or more programs, which are used by one or more processors to perform the methods described in this disclosure.

[0100] The above description is merely a preferred embodiment of this disclosure and an explanation of the technical principles employed. Those skilled in the art should understand that the scope of the invention involved in this disclosure is not limited to technical solutions formed by specific combinations of the above-described technical features, but should also cover other technical solutions formed by arbitrary combinations of the above-described technical features or their equivalents without departing from the inventive concept. For example, technical solutions formed by substituting the above-described features with (but not limited to) technical features disclosed in this disclosure that have 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 of the top k most relevant terms sorted from high to low are executed in parallel by the GPU and optimized using sparse matrix multiplication kernels.

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, and the released memory space is 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