Kv cache quantization method based on geometric perception manifold alignment llm

Abstract

A KV cache quantization method based on geometrically aware manifold-aligned LLM, which performs a unified quantization process on the Key vector and Value vector in the KV cache, decomposes the vector into two independent components of norm and direction, and respectively adopts a quantization strategy matching the geometric structure, so that the quantization process is aligned with the hyperspherical manifold of the Key vector and the heavy-tailed hyperbolic geometry of the Value vector, without training, the effective inference performance can be maintained at an extremely low bit rate, and under the condition of 2.28 bits of extreme compression, the effective inference performance can still be maintained, and under the condition of 3.27 bits, the performance level close to full precision can be achieved.

Description

Technical Field

[0001] This invention relates to a technology in the field of Large Language Models (LLMs), specifically a KV cache quantization method based on geometry-aware manifold aligned LLM. Background Technology

[0002] As large language models evolve towards long-context inference, key-value (KV) caching has become a major memory bottleneck in the inference stage. Existing KV caching compression methods suffer from the following shortcomings: 1) Uniform quantization methods use a uniform Euclidean linear grid to quantize the key and value. However, key and value vectors play different functional roles in the attention mechanism and have fundamentally different geometric distribution characteristics. Uniform quantization ignores this geometric heterogeneity, leading to a sharp amplification of quantization errors under low-bit conditions. 2) Sparse eviction methods based on token importance reduce memory by discarding non-key tokens, but this causes irreversible loss of contextual information, limiting its applicability in scenarios requiring retrieval of dense information. 3) Rotation-based methods use Hadamard transformation to smooth outliers, but require offline preprocessing, increasing deployment complexity. 4) Local polar coordinate decomposition methods use the dimension of Rotated Position Encoding (RoPE) to perform local polar coordinate transformation on the structure, but only focus on the local rotation patterns within the key vector, failing to capture the global intrinsic manifold structure of the key and value, thus having limited effectiveness under extreme compression conditions. Summary of the Invention

[0003] This invention addresses the performance degradation problem of existing large language models under sub-4-bit low-bit quantization conditions of KV cache, which is easily caused by ignoring the geometric heterogeneity of key and value vectors. It proposes a KV cache quantization method based on geometry-aware manifold alignment (LLM). By decomposing the vector into two independent components, norm and direction, and adopting quantization strategies that match the geometric structure, the quantization process is aligned with the hyperspherical manifold of the key vector and the heavy-tailed hyperbolic geometry of the value vector. It can maintain effective inference performance at extremely low bit rates without training, and can still maintain effective inference performance under extreme compression conditions of 2.28 bits. Under 3.27 bits, it can achieve a near-full-precision performance level.

[0004] This invention is achieved through the following technical solution:

[0005] This invention relates to a key-value cache quantization method based on geometry-aware manifold alignment (LLM), which performs the following unified quantization process on both key and value vectors in the key-value cache, including:

[0006] Step S1: For each vector in the KV cache during the large language model inference process Perform geometric decomposition and calculate vectors of Norm derives norm components , will vector Divide by norm component Obtain the unit direction vector ;

[0007] Preferably, in calculating the unit direction vector Before, norm components Apply a lower bound truncation, i.e. ,in: These are the preset truncation parameters.

[0008] The truncation parameter is preferably a very small positive number to prevent numerical instability caused by an excessively small norm.

[0009] Step S2: For norm components Logarithmic field quantization is performed by first converting the norm components... The process involves mapping the logarithmic norm to logarithmic space, then uniformly quantizing the logarithmic norm in a block-based manner within the logarithmic space. Finally, the quantized logarithmic norm is mapped back to the original space via exponential operations to obtain the quantization norm. Specifically, this includes:

[0010] Step S21: Convert the norm components Taking the natural logarithm yields the logarithmic norm. .

[0011] Step S22: Calculate the logarithmic norm values ​​of multiple vectors according to a preset quantization block size. It is divided into several quantization blocks, and the minimum and maximum values ​​of the logarithmic norm are calculated within each quantization block.

[0012] Step S23: Within each quantization block, the logarithmic norm value is uniformly quantized according to the minimum value, maximum value, and norm quantization bit number of the quantization block to obtain the quantized logarithmic norm value.

[0013] Step S24: Take the exponent of the quantization logarithm norm value to obtain the quantization norm.

[0014] The range of the quantized number of bits is 2 to 8 bits.

[0015] Step S3: For the unit direction vector Channel-by-channel directional quantization is performed, which involves independently calculating the minimum and maximum values ​​for each dimension of the direction vector within each quantization block, and then quantizing using a uniform quantizer to obtain the quantized direction vector. Specifically, this includes:

[0016] Step S31: Divide the unit direction vectors of multiple vectors according to the preset quantization block size. It is divided into several quantization blocks.

[0017] Step S32: Within each quantization block, calculate the minimum and maximum values ​​of all vectors in each dimension of the unit direction vector.

[0018] Step S33: Within each quantization block, based on the minimum and maximum values ​​of each dimension and the number of quantization bits for the direction, perform uniform quantization independently on each dimension of the unit direction vector to obtain the quantized direction vector.

[0019] The number of directional quantization bits ranges from 2 to 4 bits.

[0020] The number of direction quantization bits in step S3 and the number of norm quantization bits in step S2 are set independently.

[0021] Step S4: Perform spherical reprojection on the quantization direction vector, i.e., calculate the quantization direction vector. After the norm, divide the quantization direction vector by it. The norm normalizes the quantized direction vector back onto the unit sphere, resulting in the reprojected direction vector.

[0022] Step S5: Multiply the quantization norm obtained in step S2 with the reprojection direction vector obtained in step S4 to obtain the reconstructed quantization vector.

[0023] This invention relates to a streaming inference scenario application based on the aforementioned KV cache quantization method. During the autoregressive generation process, the KV cache is divided into a sealed storage area and a pending buffer area according to boundary indices; when the number of tokens accumulated in the pending buffer reaches the quantization block size... At that time, the above KV cache quantization is performed on the vectors in the block, and the quantization result is moved into the sealed storage area to ensure that each token is quantized only once.

[0024] The sealed storage area contains a complete quantized block that has been quantized. The data is frozen after quantization and does not require subsequent updates.

[0025] The pending buffer contains tokens that are less than a complete quantization block recently, and is maintained in FP16 high-precision format to accommodate statistical fluctuations in incomplete blocks.

[0026] The KV caching quantization method described above requires no training and no fine-tuning or offline calibration of the large language model; it can be directly applied to the inference process of the pre-trained large language model.

[0027] Technical effect

[0028] This invention leverages the inherent geometrical differences between the key and value vectors in a KV cache. By decomposing the vectors into two independent components—norm and direction—and employing quantization strategies matched to the geometrical structure, it achieves high-fidelity compression under extremely low bit conditions. Logarithmic norm quantization utilizes the proportional compression property of logarithmic transformation, mapping norm values ​​of different orders of magnitude in heavy-tailed distributions to a uniformly distributed logarithmic space, ensuring that quantization accuracy is not biased towards any particular order of magnitude. Spherical reprojection utilizes the orthogonality of norm-direction decomposition, strictly limiting directional quantization errors to the angular domain and preventing error leakage across channels. The dynamic undetermined buffer utilizes the dependence of block quantization on complete statistics, ensuring the stability of the statistics for each block by delaying the quantization of incomplete blocks. Attached Figure Description

[0029] Figure 1 This is a flowchart of the present invention;

[0030] Figure 2 This is a flowchart illustrating the quantization of the norm domain in the embodiments.

[0031] Figure 3 This is a flowchart of the spherical reprojection in the embodiment;

[0032] Figure 4 This is a flowchart illustrating the management process of the dynamic pending buffer in the embodiment.

[0033] Figure 5 The following is a diagram illustrating the overall architecture of the quantitative system as an example. Detailed Implementation

[0034] like Figure 1 As shown, this embodiment illustrates a KV cache quantization method based on geometry-aware manifold alignment LLM, comprising:

[0035] Step S1: For each vector in the KV cache during the large language model inference process Perform geometric decomposition and calculate vectors The L2 norm is used to obtain the norm components. , will vector Divide by norm component Obtain the unit direction vector Specifically: , .

[0036] To prevent numerical instability caused by excessively small norms, a lower bound is first applied to the norm components for truncation. Then, the vector is divided by the truncated norm components to obtain the unit direction vector, specifically: ,in: This is for truncating parameters.

[0037] In this embodiment, the truncation parameter is taken as... .

[0038] The information of the original vector is decomposed into norm components. (Encoding the semantic strength of the token) and unit direction vector (Encoding the directional position of the token in the feature space), both mathematically satisfy... Quantization error can be optimized independently.

[0039] Step S2: As Figure 2 As shown, for norm components Logarithmic field quantization is performed to obtain the quantization norm, specifically including:

[0040] Step S21: Convert the norm components Taking the natural logarithm yields the logarithmic norm. .

[0041] Step S22: Calculate the logarithmic norm values ​​of multiple vectors according to a preset quantization block size. After dividing the quantization into several blocks, the minimum value of the logarithmic norm is calculated within each quantization block. and maximum value .

[0042] For example At that time, each quantization block contains the logarithmic norm values ​​of 128 vectors.

[0043] When there are invalid positions in the quantization block due to padding, these positions are excluded from the statistical range by using a mask to avoid polluting the statistics of valid data.

[0044] Step S23: Within each quantization block, based on the minimum value of the quantization block... Maximum value And the number of bits quantified Uniform quantization is performed on the logarithmic norm values, specifically: quantization series. Quantization step size Quantization index Cut off to Within range. Inverse quantization value For example, when At that time, the quantization level is 16, and the quantization step size is... .

[0045] Step S24: Perform exponential operation on the quantized logarithmic norm value. Obtain the quantization norm .

[0046] Logarithmic field quantization keeps the relative reconstruction error of the norm constant, regardless of the norm size. Small and large norm values ​​achieve the same relative accuracy, effectively adapting to the heavy-tailed distribution of the value vector.

[0047] For each quantization block, storage is required. and Two FP32 metadata.

[0048] Step S3: For the unit direction vector Perform channel-by-channel directional quantization to obtain quantized direction vectors. Specifically, this involves quantizing the unit direction vectors of multiple vectors according to a preset quantization block size. It is divided into several quantization blocks. Within each quantization block, for each dimension of the direction vector... , For each vector along the attention head dimension, calculate its minimum value. and maximum value Based on the minimum value of each dimension. Maximum value and direction quantization bit number Each dimension of the unit direction vector is independently and uniformly quantized.

[0049] For example, when At that time, the quantization level is 4, and the quantization step size is... Quantization index Cut off to Within range. Inverse quantization value Channel-by-channel independent processing can adapt to the heterogeneous distribution of the key vector caused by Rotated Position Encoding (RoPE) on different frequency channels, and at the same time adapt to the variance differences of different feature dimensions of the value vector.

[0050] For each channel of each quantization block, the preferred storage is... and Two FP32 metadata.

[0051] Step S4: As Figure 3 As shown, the quantization direction vector is reprojected onto a spherical surface to obtain the reprojected direction vector. Specifically, after channel-by-channel independent quantization in step S3, the quantization direction vector is... The quantization errors in each dimension are independent of each other, causing them to deviate from the unit sphere, i.e. Spherical reprojection reprojects the direction vector back onto the unit sphere through L2 normalization: .

[0052] The core function of the reprojection back to the unit sphere operation is to restore the norm-orthogonal decomposition relationship established in step S1. Without reprojection, the actual norm of the final reconstructed vector is... This means that the error from directional quantization leaks into the norm channel, causing error coupling between the two channels. After reprojection, the norm of the reconstructed vector is strictly equal to... The directional error is limited to the angular domain.

[0053] Spherical reprojection ensures that the norm of the reconstructed vector is exactly equal to the quantization norm. This prevents directional quantization errors from leaking into the norm channel, ensuring that the quantization errors of the norm and directional channels are independent of each other.

[0054] Spherical reprojection is particularly important for the accuracy of attention routing for the key vector, as the attention weights of the key are highly sensitive to directional shifts.

[0055] Step S5: Quantize the norm obtained in step S2 The reprojection direction vector obtained in step S4 Multiply to obtain the quantization vector .

[0056] The norm quantization bit number in step S2 The number of direction quantization bits in step S3 This can be set independently. Preferably, the number of directional quantization bits... Quantitative number of bits At this point, the effective bit rate is approximately 2.28 bits. Another example is the number of bits for directional quantization. Quantitative number of bits At this point, the effective bit rate is approximately 3.27 bits.

[0057] The effective bit rate The first term represents the per-element cost of direction quantization, the second term represents the per-element amortized cost of norm quantization, and the third term represents the per-element amortized cost of FP32 metadata. For long sequences, the metadata cost term approaches zero, and the effective bit rate approaches the number of direction bits. .

[0058] like Figure 5As shown, the KV buffer quantization system based on geometry-aware manifold alignment (LLM) to implement the above method includes: a geometric decomposition module, a logarithmic domain norm quantization module, a channel-by-channel direction quantization module, a spherical reprojection module, a vector reconstruction module, and a dynamic undetermined buffer management module. Specifically: the geometric decomposition module calculates the L2 norm for each vector in the KV buffer to obtain the norm component. After dividing the vector by the norm component to obtain the unit direction vector, the norm component is output to the logarithmic domain norm quantization module, and the unit direction vector is output to the channel-by-channel direction quantization module. The logarithmic domain norm quantization module maps the norm component to the logarithmic space, performs uniform quantization in a block-based manner in the logarithmic space, and then uses exponential... The computation is mapped back to the original space to obtain the quantization norm, which is then output to the vector reconstruction module. The channel-by-channel directional quantization module independently performs uniform quantization on each dimension of the direction vector within each quantization block to obtain the quantized direction vector, which is then output to the spherical reprojection module. The spherical reprojection module calculates the L2 norm of the quantized direction vector, divides the quantized direction vector by its L2 norm, and renormalizes it to the unit sphere to obtain the reprojected direction vector, which is then output to the vector reconstruction module. The vector reconstruction module multiplies the quantization norm by the reprojected direction vector to obtain the quantized vector. The dynamic pending buffer management module manages the division between the sealed storage area and the pending buffer during streaming inference and controls the timing of quantization triggering.

[0059] The logarithmic domain norm quantization module maintains the minimum and maximum logarithmic domain metadata for each quantization block.

[0060] The channel-by-channel directional quantization module maintains metadata for the minimum and maximum values ​​of each channel in each quantization block.

[0061] The aforementioned control over quantization triggering timing refers to the following: when the number of tokens accumulated in the pending buffer reaches the size of the quantization block, the four-stage processing pipeline from the geometric decomposition module to the vector reconstruction module is triggered, and the quantization result is moved into the sealed storage area. During inference, the attention calculation module simultaneously accesses the quantized KV cache in the sealed storage area (which needs to be dequantized first) and the FP16 KV cache in the pending buffer, merges the two parts, and then performs attention weight calculation and value aggregation.

[0062] like Figure 4 As shown in the experiment, when the above quantization method is applied to a streaming inference scenario, a dynamic pending buffer mechanism is used to manage the triggering timing of steps S1 to S5, specifically:

[0063] A) Using quantized block size For example, let the current sequence length be... Calculate the boundary index. The sealed storage area [0, 256) contains two complete quantized blocks formed by the first 256 tokens. The statistics of each block have been determined and frozen, and the data is stored in low-bit format. The pending buffer [256, 350) contains the most recent 94 tokens (350 − 256 = 94), maintaining the FP16 high-precision format.

[0064] B) When the sequence length reaches When the pending buffer accumulates to 128 tokens, reaching the quantization block size, the final statistics of these 128 tokens are calculated, and the quantization processing of steps S1 to S5 is performed. The block is then moved to the sealed storage area, and the new pending buffer is empty.

[0065] This mechanism ensures that each token is quantized only once, avoiding the cumulative drift introduced by repeated token-by-token quantization, while maintaining the memory compression advantages of streaming inference. The upper bound of the memory overhead of the pending buffer is... FP16 vectors, for example Head dimension At this time, the buffer occupies a maximum of 32KB, which is negligible compared to the overall KV cache of long sequences.

[0066] This application was experimentally validated on various large language models, covering tasks such as basic language modeling, mathematical reasoning, code generation, and long context understanding. With a 2.28-bit configuration (2 bits for direction, 4 bits for norm), the Llama-3-8B model achieved a WikiText-2 perplexity of 10.52 (FP16 baseline: 8.28), maintaining effective performance even when other methods generally failed. With a 3.27-bit configuration (3 bits for direction, 3 bits for norm), the Qwen2-7B model achieved a GSM8K mathematical reasoning accuracy of 81.58%, completely consistent with the FP16 baseline, achieving lossless compression. In long context scenarios, the 3.27-bit configuration still achieved a Needle-in-a-Haystack retrieval accuracy of 95% with a 128K token context length.

[0067] The following experiments were conducted on a server equipped with an NVIDIA H20 GPU. The test models included five large language models: Llama-3-8B, Llama-3-70B, Qwen2-7B, Qwen2.5-Coder-7B, and Mistral-7B. The quantization block size B=128, and the attention head dimension d=128. This invention primarily tests two configurations: d2 / n4 configuration (2 bits for direction, 4 bits for norm, effective bit rate approximately 2.28 bits) and d3 / n3 configuration (3 bits for direction, 3 bits for norm, effective bit rate approximately 3.27 bits). Comparison methods included existing methods such as Uniform quantization, KIVI, PolarQuant, and QuaRot.

[0068] Table 1 shows the perplexity (PPL) test results on the WikiText-2 dataset after applying the method of this invention to various large language models. A lower perplexity indicates stronger language modeling ability.

[0069] Table 1

[0070] Experimental results show that under extreme compression conditions of 2.28 bits, uniform 2-bit quantization causes the perplexity to spike to 248.40 on Llama-3-8B and 2000.33 on Llama-3-70B, resulting in severe performance degradation. The PolarQuant method, under the same conditions, exhibits perplexities of 717.43 and 3267.70, respectively, with even more severe degradation. In contrast, this invention, at the same bit rate (2.28 bits), achieves a perplexity of only 10.52 on Llama-3-8B and 7.20 on Llama-3-70B, effectively avoiding performance degradation. With a 3.27-bit configuration, the perplexity degradation of this invention does not exceed 0.3 points across all models, achieving near-lossless compression.

[0071] Table 2 shows the accuracy test results on the GSM8K mathematical inference dataset.

[0072] Table 2

[0073] Experimental results show that, under a 3.27-bit configuration, the Qwen2-7B model achieves a GSM8K mathematical inference accuracy of 81.58%, perfectly consistent with the FP16 full-precision baseline, demonstrating lossless compression. Under extreme compression at 2.28 bits, the invention maintains an accuracy of 58.07% on Llama-3-8B, nearly five times that of the KIVI method (12.21%), while using less memory.

[0074] Table 3 shows the Pass@1 accuracy test results on the HumanEval code generation dataset.

[0075] Table 3

[0076] Experimental results show that the present invention recovers 97% (78.05% vs. 80.49%) of the full-precision baseline performance of the professional code generation model Qwen2.5-Coder-7B under a 3.27-bit configuration, verifying the effectiveness of the spherical reprojection mechanism in maintaining the accuracy of attention routing.

[0077] Table 4 shows the test results (F1 / accuracy score) on the LongBench comprehensive evaluation dataset, including four types of tasks: single-document question answering (S-Doc), multi-document question answering (M-Doc), text summarization (Summ), and few-shot learning.

[0078] Table 4 Llama-3-8B Model

[0079] Table 5 Mistral-7B Model

[0080] Table 6 Qwen2-7B Model

[0081] Experimental results show that in long-context understanding scenarios, the performance of this invention closely follows the FP16 baseline in a 3.27-bit configuration, with degradation on both Qwen2-7B and Mistral-7B not exceeding one point. Under extreme compression of 2.28 bits, this invention achieves an average score of 25.22 on Llama-3-8B, significantly outperforming the KIVI 2-bit method (19.79) which uses more memory. In few-shot learning tasks, Llama-3-8B scores 40.58 in a 3.27-bit configuration, even slightly exceeding the FP16 baseline (38.96), indicating that geometrically perceptual quantization may provide a slight regularization effect.

[0082] Table 7 shows the accuracy test results on the Needle-in-a-Haystack precision retrieval benchmark.

[0083] Table 7

[0084] Experimental results show that, with a 3.27-bit configuration, this invention maintains a 95% retrieval accuracy on Qwen2-7B with a 128K token context length. Under 2.28-bit conditions, the retrieval accuracy of both uniform quantization and PolarQuant methods drops to 0%, while this invention maintains 85% accuracy in an Llama-3-8B 8K context.

[0085] Table 8 shows the dequantization latency results (in milliseconds) measured on an NVIDIA H20 GPU using a fused Triton kernel.

[0086] Table 8

[0087] With a context length of 131K, the dequantization operation of this invention takes only 0.085 milliseconds, which is only 0.8% of the attention computation time of Flash-Attention-2 (approximately 10.0 milliseconds), and the additional overhead is negligible.

[0088] Table 9 shows the memory bandwidth saving effect (taking 131K context and 128 head dimension as an example).

[0089] Table 9

[0090] This invention achieves a 7.1x memory compression ratio under a d2 / n4 configuration, which can be directly translated into an increase in batch processing scale.

[0091] To verify the necessity of each component of the present invention, ablation experiments were conducted by removing each component one by one on the Qwen2-7B model. The results are shown in Table 10: The effective bit rate of the ablation experiment was calculated based on the evaluation sequence length (2048 tokens), which included the bit overhead of the FP16 tokens in the dynamic pending buffer. Therefore, the d2 / n4 configuration was reported as 2.71 bits here. The 2.28 bits reported in the main experiment is the asymptotic effective bit rate for long sequences (such as 131K tokens), in which case the buffer overhead can be ignored.

[0092] Table 10

[0093] Experimental Analysis: Geometric decomposition is the most critical component; its removal worsened the perplexity from 8.32 to 11.81, and the GSM8K accuracy plummeted from 69.5% to 38.5%. The dynamic undetermined buffer is crucial for streaming inference; its removal caused the GSM8K accuracy to collapse from 69.5% to 7.0%. Spherical reprojection contributed 0.55 points to the perplexity improvement. Logarithmic norm quantification contributed a 4.5% accuracy improvement on GSM8K.

[0094] Table 11 shows the experimental results of testing different combinations of directional bit number and norm bit number on the Qwen2-7B model.

[0095] Table 11

[0096] Experimental analysis: The number of direction bits plays a decisive role in performance. When the number of direction bits is only 1, even if the norm bit number is increased to 6, the perplexity still exceeds 1800, resulting in a catastrophic collapse. However, when the number of direction bits is 3, even if the norm bit number is only 2, the perplexity is only 7.05.

[0097] Table 12 shows the experimental results of testing different quantization block sizes under the d2 / n4 configuration of the Qwen2-7B model.

[0098] Table 12

[0099] Experimental Analysis: There is a trade-off between precision and compression ratio in quantization block size. Smaller blocks provide finer statistics but increase metadata overhead, while larger blocks reduce metadata but lack sufficient statistical granularity. This invention defaults to B=128, achieving a perplexity degradation of 1.59 at an effective bit rate of 2.71 bits, thus balancing compression ratio and precision.

[0100] Compared with existing technologies, this invention uses a fused Triton kernel to implement quantization and dequantization operations, integrating logarithmic domain dequantization, direction reconstruction, and spherical reprojection into a single GPU kernel, eliminating the memory write overhead of intermediate results. Direction vectors and norm values ​​are stored in compact uint8 tensors, employing custom bit unpacking logic: 4 values ​​are stored per byte for 2 bits, extracted through bit shifting and masking; for 3 bits, cross-byte reading is performed, extracted through double-byte concatenation; for 4 bits, 2 values ​​are stored per byte, extracted through bit shifting and masking. Hardware-accelerated fast reciprocal square root operations replace division for normalization, hardware-accelerated exponentiation operations are used for logarithmic domain transformation, and Triton's automatic tuner selects the optimal thread block parameters for different sequence lengths, achieving a 6 to 8x speedup in long-context scenarios.

[0101] The aforementioned technical means enable this invention to maintain effective inference performance under extreme compression conditions of 2.28 bits (Llama-3-8B perplexity 10.52, while existing uniform quantization methods experience a perplexity spike to 248.40 at similar bit rates, resulting in complete performance degradation); achieve near-lossless compression under 3.27-bit configuration (Qwen2-7B mathematical inference accuracy 81.58%, perfectly consistent with the FP16 full-precision baseline); achieve 95% accuracy in precise retrieval in scenarios with 128K tokens and long contexts; the dequantization operation latency is only 0.085 milliseconds, accounting for less than 1% of the attention computation time; achieve a 7.1x memory compression ratio; and eliminate the need for training, fine-tuning, or offline calibration of large language models, allowing direct application to the inference process of any pre-trained model.

[0102] In summary, this invention solves the performance crash problem under sub-4-bit conditions, eliminates the need for fine-tuning or offline calibration of large language models, and can be directly applied to the inference process of any pre-trained large language model without introducing additional training costs or data dependencies. It boasts broad applicability and plug-and-play deployment advantages. Furthermore, this invention seamlessly adapts to autoregressive streaming inference scenarios, quantizing each token only once to avoid the cumulative error introduced by repeated quantization, while the buffer memory overhead is negligible.

[0103] The above-described specific implementations can be partially adjusted by those skilled in the art in different ways without departing from the principles and purpose of the present invention. The scope of protection of the present invention is defined by the claims and is not limited to the above-described specific implementations. All implementation schemes within the scope of the claims are bound by the present invention.

Claims

1. A KV buffer quantization method based on geometry-aware manifold aligned LLM, characterized in that, The following unified quantization process is performed on both the key vectors and value vectors in the KV cache, including: Step S1: For each vector in the KV cache during the large language model inference process Perform geometric decomposition and calculate vectors of Norm derives norm components , will vector Divide by norm component Obtain the unit direction vector ; Step S2: For norm components Logarithmic field quantization is performed by first converting the norm components... The logarithmic norm is obtained by mapping to the logarithmic space. Then, the logarithmic norm is uniformly quantized in the logarithmic space in a block manner. Finally, the quantized logarithmic norm is mapped back to the original space through exponential operation to obtain the quantization norm. Step S3: For the unit direction vector Channel-by-channel directional quantization is performed, which means that after independently calculating the minimum and maximum values ​​of each dimension of the directional vector within each quantization block, a uniform quantizer is used to quantize the vector to obtain the quantized directional vector. Step S4: Perform spherical reprojection on the quantization direction vector, i.e., calculate the quantization direction vector. After the norm, divide the quantization direction vector by it. The norm normalizes the quantized direction vector back onto the unit sphere, resulting in the reprojected direction vector. Step S5: Multiply the quantization norm obtained in step S2 with the reprojection direction vector obtained in step S4 to obtain the reconstructed quantization vector; The number of directional quantization bits in step S3 and the number of norm quantization bits in step S2 are set independently.

2. The KV cache quantization method based on geometry-aware manifold aligned LLM according to claim 1, characterized in that, in Calculate the unit direction vector Before, norm components Apply a lower bound truncation, i.e. ,in: These are the preset truncation parameters.

3. The KV buffer quantization method based on geometry-aware manifold aligned LLM according to claim 1 or 2, characterized in that, Step S2 specifically includes: Step S21: Convert the norm components Taking the natural logarithm yields the logarithmic norm. ; Step S22: Calculate the logarithmic norm values ​​of multiple vectors according to a preset quantization block size. Divide into several quantization blocks, and calculate the minimum and maximum values ​​of the logarithmic norm within each quantization block; Step S23: Within each quantization block, the logarithmic norm value is uniformly quantized according to the minimum value, maximum value and norm quantization bit number of the quantization block to obtain the quantized logarithmic norm value. Step S24: Take the exponent of the quantization logarithm norm value to obtain the quantization norm.

4. The KV buffer quantization method based on geometry-aware manifold aligned LLM according to claim 1 or 2, characterized in that, Step S3 specifically includes: Step S31: Divide the unit direction vectors of multiple vectors according to the preset quantization block size. Divided into several quantization blocks; Step S32: Within each quantization block, calculate the minimum and maximum values ​​of all vectors in each dimension of the unit direction vector; Step S33: Within each quantization block, based on the minimum and maximum values ​​of each dimension and the number of quantization bits for the direction, perform uniform quantization independently on each dimension of the unit direction vector to obtain the quantized direction vector.

5. A streaming inference scenario application based on any one of the KV cache quantization methods described in claims 1-4, characterized in that, During the autoregressive generation process, the KV cache is divided into a sealed storage area and a pending buffer based on boundary indices; when the number of tokens accumulated in the pending buffer reaches the quantization block size... At that time, the above KV cache quantization is performed on the vectors in the block, and the quantization result is moved into the sealed storage area to ensure that each token is quantized only once.

6. The application of streaming inference in the scenario according to claim 5, characterized in that, The sealed storage area contains complete quantized blocks that have been quantized. The data is frozen after quantization and does not require subsequent updates. The pending buffer contains tokens that are less than a complete quantization block recently, and is maintained in FP16 high-precision format to accommodate statistical fluctuations of incomplete blocks; The KV caching quantization method described above requires no training and no fine-tuning or offline calibration of the large language model; it can be directly applied to the inference process of the pre-trained large language model.

7. A KV cache quantization system for implementing the method of any one of claims 1-4, characterized in that, include: The system comprises a geometric decomposition module, a logarithmic norm quantization module, a channel-by-channel directional quantization module, a spherical reprojection module, a vector reconstruction module, and a dynamic pending buffer management module. Specifically: the geometric decomposition module calculates the L2 norm for each vector in the KV buffer to obtain the norm component; after dividing the vector by the norm component to obtain the unit direction vector, it outputs the norm component to the logarithmic norm quantization module and the unit direction vector to the channel-by-channel directional quantization module, respectively; the logarithmic norm quantization module maps the norm component to logarithmic space, performs uniform quantization in a block-based manner in logarithmic space, and then maps it back to the original space through exponential operations to obtain the quantized norm, which is then output to the vector reconstruction module; the channel-by-channel directional quantization module independently performs uniform quantization on each dimension of the direction vector within each quantization block to obtain the quantized direction vector, which is then output to the spherical reprojection module. The spherical reprojection module calculates the L2 norm of the quantization direction vector, divides the quantization direction vector by its L2 norm to renormalize it onto the unit sphere, and obtains the reprojection direction vector, which is then output to the vector reconstruction module. The vector reconstruction module multiplies the quantization norm with the reprojection direction vector to obtain the quantization vector. The dynamic pending buffer management module manages the division between the sealed storage area and the pending buffer during the streaming inference process and controls the timing of quantization triggering.

8. The KV cache quantization system according to claim 7, characterized in that, The logarithmic domain norm quantization module maintains the logarithmic domain minimum and maximum metadata for each quantization block; The channel-by-channel directional quantization module maintains metadata for the minimum and maximum values ​​of each channel in each quantization block.

9. The KV cache quantization system according to claim 7, characterized in that, The aforementioned control of quantization trigger timing refers to the following: when the number of tokens accumulated in the pending buffer reaches the size of the quantization block, the four-stage processing pipeline from the geometric decomposition module to the vector reconstruction module is triggered, and the quantization result is moved into the sealed storage area. During the inference process, the attention calculation module simultaneously accesses the dequantized quantized KV cache in the sealed storage area and the FP16 KV cache in the pending buffer, merges the two parts, and performs attention weight calculation and value aggregation.

Images