Hybrid precision parallel compression method and system for optimizing large model key value cache

Abstract

The invention discloses a mixed precision parallel compression method for optimizing large model key value cache, which combines the advantages of mixed precision key value cache compression with an advanced system optimization technology. Based on the characteristic that key value pairs needing high-precision retention in mixed precision compression are the same as key value pairs used for attention calculation in a prefetching strategy, low-precision KV caches are stored in a GPU memory, and meanwhile predicted high-precision important KV caches are dynamically prefetched from a CPU memory according to needs. The technical problems that an existing multi-head attention mechanism-based method is incompatible with an existing pre-training model and cannot be directly applied to a closed source or a fine-tuned large model, and generalization of the method is reduced can be solved, and the situation that an existing pruning-based method is prone to deleting unimportant marks in the current stage can be solved. And contextual information loss is caused.

Description

Technical Field

[0001] This invention belongs to the field of large model inference optimization technology, and more specifically, relates to a mixed-precision parallel compression method and system for optimizing large model key-value caching. Background Technology

[0002] Large Transformer-based models have achieved significant success in machine translation, document summarization, and code generation, primarily due to their self-attention mechanism. However, the autoregressive nature of token generation during the inference phase necessitates recalculating the key-value vectors of all previously generated tokens, increasing computational complexity. To mitigate this complexity, large models utilize key-value (KV) caching to store previously computed key-value pairs on the GPU, significantly improving deployment efficiency. However, the size of the KV cache grows linearly with the length of the input sequence, easily exceeding memory bottlenecks as the inference context expands. Therefore, key-value cache compression has emerged as a solution.

[0003] Existing key-value cache compression methods mainly include multi-head attention-based methods, pruning-based methods, quantization-based methods, and offloading-based methods. Multi-head attention-based methods include Grouped Query Attention (GQA), which can reduce the number of heads in the key-value cache, but requires modification of the model architecture. Pruning-based methods retain historical tags with high attention weights for the current tag, while discarding or normalizing the key-value caches of other tags. Quantization-based methods retain key tags with high-precision data types, while quantizing less important tags into low-precision data types. Offloading-based methods offload the key-value cache from GPU memory to CPU memory, and in the decoding stage, partially load the key-value cache of important tags from the CPU to perform computation.

[0004] However, the key-value caching compression methods mentioned above all have some drawbacks that cannot be ignored:

[0005] First, methods based on multi-head attention mechanisms are structural optimizations that require modification of the large model architecture during training. This makes them incompatible with existing pre-trained models and cannot be directly applied to closed-source or finely tuned large models, thus reducing the generalization ability of the method.

[0006] Second, pruning-based methods rely on the current attention score to determine the importance of a label, which can easily lead to the deletion of unimportant labels in the current stage, affecting subsequent generation and causing the loss of contextual information.

[0007] Third, in complex generation tasks, quantization-based methods require complex scheduling strategies for mixed precision, and low-precision key-value caching may lead to inaccurate model attention, thereby reducing the reliability of the output.

[0008] Fourth, the offloading-based method requires transferring a large amount of KV cache from CPU memory to GPU memory. However, since the data transfer speed of the PCIe channel is much lower than the access speed of the video memory, the communication time becomes a new performance bottleneck. Summary of the Invention

[0009] To address the aforementioned shortcomings or improvement needs of existing technologies, this invention provides a mixed-precision parallel compression method and system for optimizing key-value caching in large models. Its purpose is to combine the advantages of mixed-precision key-value caching compression with advanced system optimization techniques. Based on the characteristic that the key-value pairs requiring high-precision retention in mixed-precision compression are the same as those used for attention calculation in the prefetching strategy, low-precision key-value caching is stored in GPU memory, while high-precision salient key-value pairs are dynamically prefetched from CPU memory as needed. This solves the technical problems of existing multi-head attention mechanisms being incompatible with existing pre-trained models, unable to be directly applied to closed-source or fine-tuned large models, and reducing the generalization ability of the method. It also addresses the issue that existing pruning-based methods are prone to deleting unimportant labels at the current stage.

[0010] The technical issues include the loss of contextual information, the need for complex scheduling strategies for mixed precision in existing quantization-based methods in complex generation tasks, the potential for inaccurate model attention due to low-precision key-value caching, which reduces the reliability of the output, and the need for existing offloading-based methods to transfer large amounts of KV cache from CPU memory to GPU memory, where communication time becomes a new performance bottleneck because the data transfer speed of the PCIe channel is much lower than the access speed of GPU memory.

[0011] To achieve the above objectives, according to one aspect of the present invention, a mixed-precision parallel compression method for optimizing large model key-value caching is provided, comprising the following steps:

[0012] (1) Obtain the data X input by the user in the pre-filling stage of the large model inference process, and process the data X using the self-attention mechanism to obtain the key K matrix and value V matrix corresponding to the data X. The size of the two matrices is m rows * d columns, where d represents the dimension of the self-attention mechanism hidden layer preset by the large model.

[0013] (2) Perform element-wise parallel 1-bit quantization on the K matrix obtained in step (1) to obtain the corresponding 1-bit quantization symbol matrix K. ±1 =sign(K), where the formula for calculating each element is as follows:

[0014]

[0015] Where i represents the row index and j represents the column index, and i∈[1,m], j∈[1,d];

[0016] (3) In the K matrix obtained in step (1), calculate the sum of the absolute values ​​of the elements in each column in parallel. The calculation formula is as follows:

[0017]

[0018] Sum j Let represent the sum of the absolute values ​​of the elements in the j-th column of matrix K;

[0019] (4) Based on the absolute value of each column element in the K matrix obtained in step (3) and Sum j Calculate the average absolute value (Avg) of the elements in each column in parallel. j ;

[0020] (5) Sort the average absolute values ​​of all columns in the K matrix obtained in step (4) in descending order. Construct an outlier column matrix O with the first k absolute values ​​and their corresponding columns, which has a size of m rows * k columns. Construct an outlier column index set I with the indexes of all columns corresponding to the first k absolute values ​​and their corresponding columns in the K matrix, where k is equal to 10% of the number of columns d in the K matrix, rounded up. I = {i 1， i 2， ..., i k}, where i n Let i represent the index of the nth column, where n ∈ [1, k], 1 ≤ i n ≤d;

[0021] (6) Quantize the abnormal column matrix O obtained in step (5) and perform parallel computation on each column to obtain the scaling factor s of each column. n and Zero Point zp n ;

[0022] (7) Based on the scaling ratio s of each column in the abnormal column matrix O obtained in step (6), n and Zero Point zp n The abnormal column matrix O is then subjected to element-wise parallel 8-bit quantization to obtain the quantized abnormal column matrix QNT(O).

[0023] (8) Take the 1-bit quantization symbol matrix K obtained in step (2) ±1 The 8-bit quantization anomaly column matrix QNT(O) obtained in step (7) is merged to obtain the final quantization matrix K, which is then saved to GPU memory.

[0024] (9) Based on the quantization matrix obtained in step (8) The K matrix obtained in step (1) is subjected to column-parallel residual processing to obtain the residual matrix K. residual And unload it to CPU memory:

[0025]

[0026] (10) Select the maximum value in each column of the V matrix obtained in step (1) as the scaling factor c for that column. j ;

[0027] (11) Based on the scaling factor c of each column in the V matrix obtained in step (10) j The V matrix obtained in step (1) is subjected to column-parallel scaling and quantization to obtain the quantized matrix V, and then saved to the GPU memory.

[0028] (12) Based on the quantization matrix obtained in step (11) The V matrix obtained in step (1) is subjected to column-parallel residual processing to obtain the residual matrix V. residual And unload it to CPU memory:

[0029]

[0030] (13) Obtain the data x input by the user in the decoding stage of the large model inference process, and process the data x using the self-attention mechanism to obtain the query q vector corresponding to the data x, the length of which is d;

[0031] (14) Combine the q vector obtained in step (13) with the quantization matrix obtained in step (8). Each line in and the quantization matrix obtained in step (11) Each line in Similarity calculations are performed separately to obtain the q vector and the quantization matrix. Each line Similarity score between And the q vector and the quantization matrix Each line Similarity score between

[0032] (15) Combine the q vector obtained in step (14) with the quantization matrix Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are used to construct an index set K.index The q vector and quantization matrix obtained in step (14) Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are constructed into an index set V. index , where t equals 10% of the row number d of matrix K, rounded up;

[0033] (16) Based on the index set K obtained in step (15) index The residual matrix K obtained in step (9) residual Perform a search to obtain the index set K. index All row indices in the residual matrix K residual The prefetch matrix K consisting of all the corresponding row vectors in the matrix. prefetch ; and based on the index set V obtained in step (15) index The residual matrix V obtained in step (12) residual Perform a search to obtain the index set V index All row indices in the residual matrix V residual The prefetch matrix V, which consists of all the corresponding row vectors in the matrix. prefetch Both prefetch matrices are of size l rows * d columns, and the prefetch matrix K is... prefetch and V prefetch Send to GPU memory;

[0034] (17) Use the prefetch matrix K obtained in step (16) prefetch For the quantization matrix stored in GPU memory Perform row vector replacement to obtain the final compressed matrix K. final And use the prefetch matrix V obtained in step (16) prefetch Quantization matrix stored in GPU memory Perform row vector replacement processing to obtain the final compressed matrix V in GPU memory. final .

[0035] Preferably, the calculation formula for step (6) is as follows:

[0036]

[0037] zp n =X min,n

[0038] Among them, X max,n and X min,n These are the maximum and minimum values ​​of the elements in the nth column of the abnormal column matrix O obtained in step (5).

[0039] Preferably, in step (7), the element QNT(O) in the i-th row and n-th column of the quantization anomaly column matrix QNT(O) is... i,n The calculation formula for ) is as follows:

[0040]

[0041] The `round` function rounds the result to the nearest integer, and the `clamp` function sets the result to the nearest integer (0, 2). 8 Integers within the range of -1].

[0042] Preferably, the quantization matrix in step (8) The element in the i-th row and n-th column equal:

[0043]

[0044] Preferably, the quantization matrix in step (11) medium elements The calculation formula is as follows:

[0045]

[0046] Where V i,j This represents the element in the i-th row and j-th column of matrix V.

[0047] Preferably, the specific calculation formula in step (14) is as follows:

[0048]

[0049] in Indicates line The transpose of the matrix, Indicates line The transpose of .

[0050] According to another aspect of the present invention, a mixed-precision parallel compression system for optimizing large model key-value caching is provided, comprising:

[0051] The first module is used to obtain the data X input by the user in the pre-filling stage of the large model inference process, and to process the data X using the self-attention mechanism to obtain the key K matrix and value V matrix corresponding to the data X. Both matrices are m rows * d columns, where d represents the dimension of the self-attention mechanism hidden layer preset by the large model.

[0052] The second module performs element-wise parallel 1-bit quantization on the K matrix obtained from the first module to obtain the corresponding 1-bit quantization symbol matrix K. ±1=sign(K), where the formula for calculating each element is as follows:

[0053]

[0054] Where i represents the row index and j represents the column index, and i∈[1,m], j∈[1,d];

[0055] The third module is used to obtain the K matrix obtained by the first module, and to calculate the sum of the absolute values ​​of each column elements in parallel. The calculation formula is as follows:

[0056]

[0057] Sum j Let represent the sum of the absolute values ​​of the elements in the j-th column of matrix K;

[0058] The fourth module is used to calculate the sum of the absolute values ​​of the elements in each column of the K matrix obtained in the third module. j Calculate the average absolute value (Avg) of the elements in each column in parallel. j ;

[0059] The fifth module sorts the average absolute values ​​of all columns in the K matrix obtained in the fourth module in descending order. It then constructs an outlier column matrix O, with size m rows * k columns, from the first k absolute values ​​and their corresponding columns in the sorted result. Finally, it constructs an outlier column index set I, where k equals 10% of the number of columns d in the K matrix, rounded up. I = {i1, i2, ..., i...} k}, where i n Let i represent the index of the nth column, where n ∈ [1, k], 1 ≤ i n ≤d;

[0060] The sixth module is used to quantize the abnormal column matrix O obtained from the fifth module, performing parallel computation on each column to obtain the scaling factor s for each column. n and Zero Point zp n ;

[0061] The seventh module is used to scale each column of the anomaly column matrix O obtained from the sixth module by the scaling factor s. n and Zero Point zp n The abnormal column matrix O is then subjected to element-wise parallel 8-bit quantization to obtain the quantized abnormal column matrix QNT(O).

[0062] The eighth module is used to process the 1-bit quantization symbol matrix K obtained from the second module. ±1 The 8-bit quantization anomaly column matrix QNT(O) obtained from the seventh module is merged to obtain the final quantization matrix. And save it to GPU memory;

[0063] The ninth module is used to quantize the matrix obtained from the eighth module. The K matrix obtained from the first module is subjected to column-parallel residual processing to obtain the residual matrix K. residual And unload it to CPU memory:

[0064]

[0065] The tenth module is used to select the maximum value of each column element from the V matrix obtained from the first module as the scaling factor c for that column. j ;

[0066] The eleventh module is used to calculate the scaling factor c for each column of the V matrix obtained in the tenth module. j The V matrix obtained from the first module is subjected to column-parallel scaling and quantization processing to obtain the quantization matrix. And save it to GPU memory;

[0067] Module 12 is used to quantize the matrix obtained from Module 11. The V matrix obtained from the first module is subjected to column-parallel residual processing to obtain the residual matrix V. residual And unload it to CPU memory:

[0068]

[0069] The thirteenth module is used to obtain the data x input by the user during the decoding stage of the large model inference process, and to process the data x using a self-attention mechanism to obtain the query q vector corresponding to the data x, which has a length of d.

[0070] The fourteenth module is used to combine the q-vector obtained from the thirteenth module with the quantization matrix obtained from the eighth module. Each line in and the quantization matrix obtained from module eleven Each line in Similarity calculations are performed separately to obtain the q vector and the quantization matrix. Each line Similarity score between And the q vector and the quantization matrix Each line Similarity score between

[0071] Module 15 is used to combine the q vector obtained from Module 14 with the quantization matrix. Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are used to construct an index set K. index The q vector and quantization matrix obtained in the fourteenth module Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are constructed into an index set V. index , where t equals 10% of the row number d of matrix K, rounded up;

[0072] Module sixteen is used to determine the index set K obtained from module fifteen. index The residual matrix K obtained from the ninth module residual Perform a search to obtain the index set K. index All row indices in the residual matrix K residual The prefetch matrix K consisting of all the corresponding row vectors in the matrix. prefetch ; and based on the index set V obtained from module fifteen index The residual matrix V obtained from the twelfth module residual Perform a search to obtain the index set V index All row indices in the residual matrix V residual The prefetch matrix V, which consists of all the corresponding row vectors in the matrix. prefetch Both prefetch matrices are of size l rows * d columns, and the prefetch matrix K is... prefetch and V prefetch Send to GPU memory;

[0073] Module 17 is used to apply the prefetch matrix K obtained from Module 16. prefetch For the quantization matrix stored in GPU memory Perform row vector replacement to obtain the final compressed matrix K. final And using the prefetch matrix V obtained from the sixteenth module prefetch Quantization matrix stored in GPU memory Perform row vector replacement processing to obtain the final compressed matrix V in GPU memory. final .

[0074] In summary, compared with the prior art, the above-described technical solutions conceived by this invention can achieve the following beneficial effects:

[0075] (1) Since the present invention adopts steps (1) to (12) to optimize the storage and calculation of key value cache without modifying the original architecture of the model, it can be applied to any existing pre-trained model and thus can solve the generalization problem of existing methods based on multi-head attention mechanism.

[0076] (2) Since the present invention adopts steps (1) to (12), instead of directly deleting the marker, it represents the low importance marker as a combination of low precision compression and residual terms, and still retains its information. Therefore, it can solve the technical problem that existing pruning-based methods are prone to losing context information.

[0077] (3) Since the present invention uses steps (1) to (12) to compress at the semantic level, the semantic information retained is better than the element-by-element compression of traditional quantization. Therefore, it can solve the technical problem that the output reliability of existing quantization-based methods is reduced due to inaccurate model attention.

[0078] (4) Since the present invention adopts steps (16) to (17), it is based on the cognition that the important key value cache required for adjacent generated tags is highly similar. It prefetches the key value cache by calculating the similarity score, thereby achieving the overlap of communication time and calculation time. Therefore, it can solve the technical problem that the communication time becomes a technical bottleneck in the existing unloading-based methods.

[0079] (5) This invention allows for flexible adjustment between accuracy and inference speed by setting high-precision quantization ratio, residual retention ratio, and prefetch information ratio according to different task requirements, thus adapting to scenarios with different resource budgets;

[0080] (6) Since the present invention contains multiple steps that can be executed independently and in parallel, it can significantly improve throughput and reduce inference latency, is compatible with efficient inference frameworks, and is more suitable for solving the distributed training problem of existing large model inference optimization methods in the future direction of hardware acceleration. Attached Figure Description

[0081] Figure 1 This is a flowchart of the mixed-precision parallel compression method for optimizing large model key-value caching according to the present invention. Detailed Implementation

[0082] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the invention. Furthermore, the technical features involved in the various embodiments of this invention described below can be combined with each other as long as they do not conflict with each other.

[0083] The basic idea of ​​this invention is to propose a compression framework that uses mixed-precision quantization to solve the problem of corrupted inference context in large models. It combines the advantages of mixed-precision key-value cache compression with advanced system optimization techniques. Based on the fact that the key-value pairs that need to be preserved with high precision in mixed-precision compression are the same as the key-value pairs used for attention calculation in the prefetching strategy, the low-precision KV cache is stored in GPU memory, while the high-precision important KV cache predicted is dynamically prefetched from CPU memory as needed.

[0084] like Figure 1 As shown, this invention provides a mixed-precision parallel compression method for optimizing large model key-value caching, comprising the following steps:

[0085] (1) Obtain the data X input by the user in the pre-filling stage of the large model inference process, and process the data X using the self-attention mechanism to obtain the key (K) matrix and value (V) matrix corresponding to the data X. The size of the two matrices is m rows * d columns, where d represents the dimension of the self-attention mechanism hidden layer preset by the large model.

[0086] (2) Perform element-wise parallel 1-bit quantization on the K matrix obtained in step (1) to obtain the corresponding 1-bit quantization symbol matrix K. ±1 =sign(K), the formula for calculating each element in the sign matrix is ​​as follows:

[0087]

[0088] Where i represents the row index and j represents the column index, and i∈[1, m], j∈[1, d].

[0089] (3) In the K matrix obtained in step (1), calculate the sum of the absolute values ​​of the elements in each column in parallel. The calculation formula is as follows:

[0090]

[0091] Sum j Let represent the sum of the absolute values ​​of the elements in the j-th column of matrix K.

[0092] (4) Based on the absolute value of each column element in the K matrix obtained in step (3) and Sum j Calculate the average absolute value (Avg) of the elements in each column in parallel. j The purpose is to reduce the error introduced by 1-bit quantization, and the calculation formula is as follows:

[0093]

[0094] (5) Sort the average absolute values ​​of all columns in the K matrix obtained in step (4) in descending order. Construct an outlier column matrix O with the first k absolute values ​​and their corresponding columns, which has a size of m rows * k columns. Construct an outlier column index set I with the indexes of all columns corresponding to the first k absolute values ​​and their corresponding columns in the K matrix, where k is equal to 10% of the number of columns d in the K matrix, rounded up. I = {i1, i2, ..., i...} k}, where i n Let i represent the index of the nth column, where n ∈ [1, k], 1 ≤ i n ≤d.

[0095] (6) Quantize the abnormal column matrix O obtained in step (5) and perform parallel computation on each column to obtain the scaling factor s of each column. n and Zero Point zp n The calculation formula is as follows:

[0096]

[0097] Among them, X max,n and X min,n These are the maximum and minimum values ​​of the elements in the nth column of the abnormal column matrix O obtained in step (5);

[0098] (7) Based on the scaling ratio s of each column in the abnormal column matrix O obtained in step (6), n and Zero Point zp n The outlier column matrix O is then subjected to element-wise parallel 8-bit quantization to obtain the quantized outlier column matrix QNT(O). The elements in the i-th row and n-th column of this quantized outlier column matrix QNT(O) are... i,n The calculation formula for ) is as follows:

[0099]

[0100] The `round` function rounds the result to the nearest integer, and the `clamp` function sets the result to the nearest integer (0, 2). 8 Integers within the range of -1].

[0101] The advantage of steps (5) to (7) above is that the proportion of important markers represented by high-precision data types can be dynamically determined by adjusting the value of k.

[0102] (8) Take the 1-bit quantization symbol matrix K obtained in step (2) ±1 The 8-bit quantization anomaly column matrix QNT(O) obtained in step (7) is merged to obtain the final quantization matrix. And save it to GPU memory;

[0103] Specifically, the quantization matrix in this step The element in the i-th row and n-th column equal:

[0104]

[0105] (9) Based on the quantization matrix obtained in step (8) The K matrix obtained in step (1) is subjected to column-parallel residual processing to obtain the residual matrix K. residual And it is unloaded to CPU memory, the specific calculation formula is as follows:

[0106]

[0107] The advantage of the above steps is that representing the KV cache as a combination of low-precision compression and residual terms preserves all the information of the tags.

[0108] (10) Select the maximum value in each column of the V matrix obtained in step (1) as the scaling factor c for that column. j (Reduce quantization error by scaling the matrix using a scaling factor);

[0109] (11) Based on the scaling factor c of each column in the V matrix obtained in step (10) j The V matrix obtained in step (1) is subjected to column-parallel scaling and quantization to obtain the quantization matrix. And save it to GPU memory, this quantization matrix medium elements The calculation formula is as follows:

[0110]

[0111] Where V i,j This represents the element in the i-th row and j-th column of matrix V.

[0112] (12) Based on the quantization matrix obtained in step (11) The V matrix obtained in step (1) is subjected to column-parallel residual processing to obtain the residual matrix V. residual And it is unloaded to CPU memory, the specific calculation formula is as follows:

[0113]

[0114] The advantage of the above steps (1) to (12) is that they distinguish between important and unimportant tags in the KV cache. By using quantization with different precision, they achieve both memory minimization and maximum retention of usable information.

[0115] (13) Obtain the data x input by the user during the decoding stage of the large model inference process, and process the data x using the self-attention mechanism to obtain the query (q) vector corresponding to the data x, with a length of d;

[0116] (14) Combine the q vector obtained in step (13) with the quantization matrix obtained in step (8). Each line in and the quantization matrix obtained in step (11) Each line in Similarity calculations are performed separately to obtain the q vector and the quantization matrix. Each line Similarity score between And the q vector and the quantization matrix Each line Similarity score between The specific calculation formula is as follows:

[0117]

[0118] in Indicates line The transpose of the matrix, Indicates line The transpose of .

[0119] (15) Combine the q vector obtained in step (14) with the quantization matrix Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are used to construct an index set K. index The q vector and quantization matrix obtained in step (14) Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are constructed into an index set V. index , where t equals 10% of the row number d of matrix K, rounded up.

[0120] The advantage of steps (13) to (15) above is that the proportion of important markers in the final compressed matrix can be dynamically determined by adjusting the value of t.

[0121] (16) Based on the index set K obtained in step (15) index The residual matrix K obtained in step (9) residual Perform a search to obtain the index set K. index All row indices in the residual matrix Kresidual The prefetch matrix K consisting of all the corresponding row vectors in the matrix. prefetch ; and based on the index set V obtained in step (15) index The residual matrix V obtained in step (12) residual Perform a search to obtain the index set V index All row indices in the residual matrix V residual The prefetch matrix V, which consists of all the corresponding row vectors in the matrix. prefetch Both prefetch matrices are of size l rows * d columns, and the prefetch matrix K is... prefetch and V prefetch Send to GPU memory;

[0122] (17) Use the prefetch matrix K obtained in step (16) prefetch For the quantization matrix stored in GPU memory Perform row vector replacement to obtain the final compressed matrix K. final And use the prefetch matrix V obtained in step (16) prefetch Quantization matrix stored in GPU memory Perform row vector replacement processing to obtain the final compressed matrix V in GPU memory. final .

[0123] The advantage of steps (16) to (17) above is that the prefetched communication time overlaps with the computation time by predicting the key KV cache for the next step.

[0124] Those skilled in the art will readily understand that the above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. A mixed-precision parallel compression method for optimizing key-value caching of large models, characterized in that, Includes the following steps: (1) Obtain the data X input by the user in the pre-filling stage of the large model inference process, and process the data X using the self-attention mechanism to obtain the key K matrix and value V matrix corresponding to the data X. The size of the two matrices is m rows * d columns, where d represents the dimension of the self-attention mechanism hidden layer preset by the large model. (2) Perform element-wise parallel 1-bit quantization on the K matrix obtained in step (1) to obtain the corresponding 1-bit quantization symbol matrix K. ±1 =sign(K), where the formula for calculating each element is as follows: Where i represents the row index and j represents the column index, and i∈[1,m], j∈[1,d]; (3) In the K matrix obtained in step (1), calculate the sum of the absolute values ​​of the elements in each column in parallel. The calculation formula is as follows: Sum j Let represent the sum of the absolute values ​​of the elements in the j-th column of matrix K; (4) Based on the absolute value of each column element in the K matrix obtained in step (3) and Sum j Calculate the average absolute value (Avg) of the elements in each column in parallel. j ; (5) Sort the average absolute values ​​of all columns in the K matrix obtained in step (4) in descending order. Construct an outlier column matrix O with the first k absolute values ​​and their corresponding columns, which has a size of m rows * k columns. Construct an outlier column index set I with the indexes of all columns corresponding to the first k absolute values ​​and their corresponding columns in the K matrix, where k is equal to 10% of the number of columns d in the K matrix, rounded up. I = {i1, i2, ..., i...} k }, where i n Let i represent the index of the nth column, where n ∈ [1, k], 1 ≤ i n ≤d; (6) Quantize the abnormal column matrix O obtained in step (5) and perform parallel computation on each column to obtain the scaling factor s of each column. n and Zero Point zp n ; (7) Based on the scaling ratio s of each column in the abnormal column matrix O obtained in step (6), n and Zero Point zp n The abnormal column matrix O is then subjected to element-wise parallel 8-bit quantization to obtain the quantized abnormal column matrix QNT(O). (8) Take the 1-bit quantization symbol matrix K obtained in step (2) ±1 The 8-bit quantization anomaly column matrix QNT(O) obtained in step (7) is merged to obtain the final quantization matrix. And save it to GPU memory; (9) Based on the quantization matrix obtained in step (8) The K matrix obtained in step (1) is subjected to column-parallel residual processing to obtain the residual matrix K. residual And unload it to CPU memory: (10) Select the maximum value in each column of the V matrix obtained in step (1) as the scaling factor c for that column. j ; (11) Based on the scaling factor c of each column in the V matrix obtained in step (10) j The V matrix obtained in step (1) is subjected to column-parallel scaling and quantization to obtain the quantization matrix. And save it to GPU memory; (12) Based on the quantization matrix obtained in step (11) The V matrix obtained in step (1) is subjected to column-parallel residual processing to obtain the residual matrix V. residual And unload it to CPU memory: (13) Obtain the data x input by the user in the decoding stage of the large model inference process, and process the data x using the self-attention mechanism to obtain the query q vector corresponding to the data x, the length of which is d; (14) Combine the q vector obtained in step (13) with the quantization matrix obtained in step (8). Each line in and the quantization matrix obtained in step (11) Each line in Similarity calculations are performed separately to obtain the q vector and the quantization matrix. Each line Similarity score between And the q vector and the quantization matrix Each line Similarity score between (15) Combine the q vector obtained in step (14) with the quantization matrix Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are used to construct an index set K. index The q vector and quantization matrix obtained in step (14) Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are constructed into an index set V. index , where t equals 10% of the row number d of matrix K, rounded up; (16) Based on the index set K obtained in step (15) index The residual matrix K obtained in step (9) residual Perform a search to obtain the index set K. index All row indices in the residual matrix K residual The prefetch matrix K consisting of all the corresponding row vectors in the matrix. prefetch ; and based on the index set V obtained in step (15) index The residual matrix V obtained in step (12) residual Perform a search to obtain the index set V index All row indices in the residual matrix V residual The prefetch matrix V, which consists of all the corresponding row vectors in the matrix. prefetch Both prefetch matrices are of size l rows * d columns, and the prefetch matrix K is... prefetch and V prefetch Send to GPU memory; (17) Use the prefetch matrix K obtained in step (16) prefetch For the quantization matrix stored in GPU memory Perform row vector replacement to obtain the final compressed matrix K. final And use the prefetch matrix V obtained in step (16) prefetch Quantization matrix stored in GPU memory Perform row vector replacement processing to obtain the final compressed matrix V in GPU memory. final .

2. The mixed-precision parallel compression method for optimizing large model key-value caching according to claim 1, characterized in that, The calculation formula for step (6) is as follows: zp n =X min,n Among them, X max,n and X min,n These are the maximum and minimum values ​​of the elements in the nth column of the abnormal column matrix O obtained in step (5).

3. The mixed-precision parallel compression method for optimizing large model key-value caching according to claim 1 or 2, characterized in that, In step (7), the elements in the i-th row and n-th column of the quantized anomaly column matrix QNT(O) are... i,n The calculation formula for ) is as follows: The `round` function rounds the result to the nearest integer, and the `Clamp` function sets the result to the nearest integer (0, 2). 8 Integers within the range of [-1].

4. The mixed-precision parallel compression method for optimizing large model key-value caching according to any one of claims 1 to 3, characterized in that, Quantization matrix in step (8) The element in the i-th row and n-th column equal:

5. The mixed-precision parallel compression method for optimizing large model key-value caching according to claim 4, characterized in that, Quantization matrix in step (11) Middle elements The calculation formula is as follows: Where V i,j This represents the element in the i-th row and j-th column of matrix V.

6. The mixed-precision parallel compression method for optimizing large model key-value caching according to claim 5, characterized in that, The specific calculation formula in step (14) is as follows: in Indicates line The transpose of the matrix, Indicates line The transpose of .

7. A hybrid precision parallel compression system for optimizing key-value caching of large models, characterized in that, include: The first module is used to obtain the data X input by the user in the pre-filling stage of the large model inference process, and to process the data X using the self-attention mechanism to obtain the key K matrix and value V matrix corresponding to the data X. Both matrices are m rows * d columns, where d represents the dimension of the self-attention mechanism hidden layer preset by the large model. The second module performs element-wise parallel 1-bit quantization on the K matrix obtained from the first module to obtain the corresponding 1-bit quantization symbol matrix K. ±1 =sign(K), where the formula for calculating each element is as follows: Where i represents the row index and j represents the column index, and i∈[1,m], j∈[1,d]; The third module is used to obtain the K matrix obtained by the first module, and to calculate the sum of the absolute values ​​of each column elements in parallel. The calculation formula is as follows: Sum j Let represent the sum of the absolute values ​​of the elements in the j-th column of matrix K; The fourth module is used to calculate the sum of the absolute values ​​of the elements in each column of the K matrix obtained in the third module. j Calculate the average absolute value (Avg) of the elements in each column in parallel. j ; The fifth module sorts the average absolute values ​​of all columns in the K matrix obtained in the fourth module in descending order. It then constructs an outlier column matrix O, with size m rows * k columns, from the first k absolute values ​​and their corresponding columns in the sorted result. Finally, it constructs an outlier column index set I, where k equals 10% of the number of columns d in the K matrix, rounded up. I = {i1, i2, ..., i...} k }, where i n Let i represent the index of the nth column, where n ∈ [1, k], 1 ≤ i n ≤d; The sixth module is used to quantize the abnormal column matrix O obtained from the fifth module, performing parallel computation on each column to obtain the scaling factor s for each column. n and Zero Point zp n ; The seventh module is used to scale each column of the anomaly column matrix O obtained from the sixth module by the scaling factor s. n and Zero Point zp n The abnormal column matrix O is then subjected to element-wise parallel 8-bit quantization to obtain the quantized abnormal column matrix QNT(O). The eighth module is used to process the 1-bit quantization symbol matrix K obtained from the second module. ±1 The 8-bit quantization anomaly column matrix QNT(O) obtained from the seventh module is merged to obtain the final quantization matrix. And save it to GPU memory; The ninth module is used to quantize the matrix obtained from the eighth module. The K matrix obtained from the first module is subjected to column-parallel residual processing to obtain the residual matrix K. residual And unload it to CPU memory: The tenth module is used to select the maximum value of each column element from the V matrix obtained from the first module as the scaling factor c for that column. j ; The eleventh module is used to calculate the scaling factor c for each column of the V matrix obtained in the tenth module. j The V matrix obtained from the first module is subjected to column-parallel scaling and quantization processing to obtain the quantization matrix. And save it to GPU memory; Module 12 is used to quantize the matrix obtained from Module 11. The V matrix obtained from the first module is subjected to column-parallel residual processing to obtain the residual matrix V. residual And unload it to CPU memory: The thirteenth module is used to obtain the data x input by the user during the decoding stage of the large model inference process, and to process the data x using a self-attention mechanism to obtain the query q vector corresponding to the data x, which has a length of d. The fourteenth module is used to combine the q-vector obtained from the thirteenth module with the quantization matrix obtained from the eighth module. Each line in and the quantization matrix obtained from module eleven Each line in Similarity calculations are performed separately to obtain the q vector and the quantization matrix. Each line Similarity score between And the q vector and the quantization matrix Each line Similarity score between Module 15 is used to combine the q vector obtained from Module 14 with the quantization matrix. Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are used to construct an index set K. index The q vector and quantization matrix obtained in the fourteenth module Sort all similarity scores between all rows in descending order, and then quantize the top t similarity scores in the sorted result in the quantization matrix. The row indexes of the rows contained in the index are constructed into an index set V. index , where t equals 10% of the row number d of matrix K, rounded up; Module sixteen is used to determine the index set K obtained from module fifteen. index The residual matrix K obtained from the ninth module residual Perform a search to obtain the index set K. index All row indices in the residual matrix K residual The prefetch matrix K consisting of all the corresponding row vectors in the matrix. prefetch ; and based on the index set V obtained from module fifteen index The residual matrix V obtained from the twelfth module residual Perform a search to obtain the index set V index All row indices in the residual matrix V residual The prefetch matrix V, which consists of all the corresponding row vectors in the matrix. prefetch Both prefetch matrices are of size l rows * d columns, and the prefetch matrix K is... prefetch and V prefetch Send to GPU memory; Module 17 is used to apply the prefetch matrix K obtained from Module 16. prefetch For the quantization matrix stored in GPU memory Perform row vector replacement to obtain the final compressed matrix K. final And using the prefetch matrix V obtained from the sixteenth module prefetch Quantization matrix stored in GPU memory Perform row vector replacement processing to obtain the final compressed matrix V in GPU memory. final .

Images