Gradient-based layer importance perception KV Cache mixing precision quantification method and system

Abstract

The invention discloses a gradient-based layer importance perception KV Cache mixing precision quantification method and system, and belongs to the technical field of artificial intelligence, and the method comprises the steps: calculating the L2 gradient norm of each layer Key and Value projection weight relative to a model loss function, and evaluating the KV contribution difference between layers; allocating a layer-specific bit width based on the difference; aiming at KV distribution characteristics, asymmetrically quantizing each channel of Key and each token of Value, and compressing to a low-bit space; self-adaptively reserving a recent key token full-precision KV pair in combination with contribution difference, and compressing an old KV pair by multiplexing a quantization rule; and designing a CUDA kernel and optimizing overhead through operator fusion. According to the method, the problems of rigid static distribution, high dynamic overhead and poor long context adaptation of an existing quantification method are solved, dynamic balance of memory occupation, reasoning precision and throughput is realized, the KV Cache memory bottleneck is remarkably reduced, and the deployment and reasoning efficiency of LLM in a resource-constrained environment is improved.

Description

Technical Field

[0001] This invention belongs to the field of artificial intelligence technology, specifically relating to a gradient-based layer importance-aware KV cache mixed precision quantization method and system. Background Technology

[0002] Large language models such as GPT, Llama, and their derivatives have significantly advanced the field of natural language processing. These models excel in various tasks, including text generation, question answering, and machine translation. Key-value (KV) caches play a crucial role in the autoregressive decoding process of LLMs. KV caches drastically reduce redundant computations in attention mechanisms by storing the KV states from previous time steps for subsequent token generation. However, the memory footprint of KV caches increases linearly with sequence length, posing a significant challenge to hardware resources. For example, a 70-byte model may require over 50GB of memory to maintain a KV cache for a 20k token sequence, exceeding the capacity of a typical GPU. In scenarios involving multiple concurrent requests, the KV cache for each request cannot be shared due to its reliance on unique pre-hints. Although model parameters can be reused, memory quickly becomes saturated due to KV cache demands. Once memory is exhausted, data is unloaded to system memory or even disk, leading to frequent high-bandwidth memory (HBM) swapping with system memory. This process causes an exponential spike in latency, resulting in catastrophic performance degradation.

[0003] The aforementioned characteristics of the KV cache severely limit the deployment and inference efficiency of LLMs in resource-constrained environments, highlighting the urgent need for efficient memory optimization. Recent research addressing this issue has primarily focused on reducing the memory overhead of the KV cache through quantization and sparsity techniques. Quantization methods, especially in industry, have been widely adopted, significantly contributing to the scalability and accessibility of large-scale models. Quantizing the KV cache can significantly reduce memory usage. Existing quantization methods have demonstrated impressive model performance even with very low bit widths. For example, KIVI introduces 2-bit asymmetric quantization, KVQuant employs a non-uniform quantization strategy to handle outliers in the key separately, QAQ implements dynamic mixed-precision quantization, Atom explores 4 / 8-bit activation and weighted mixed quantization, QJL introduces 1-bit key quantization, and KVTuner treats mixed KV quantization as a search optimization problem. While these methods are effective, they rely on static, one-size-fits-all precision allocation or introduce high computational overhead, failing to dynamically prioritize key KVs, especially in long-context tasks, forcing a trade-off between memory, precision, and throughput. Furthermore, studies addressing the attention sink phenomenon, such as StreamingLLM and PyramidInfer, optimize long sequences by preserving key contexts but lack guidance on layer importance. Existing methods are inefficient in analyzing the differences in key-value contributions between layers and cannot efficiently allocate bit widths. These methods either rely on static, one-size-fits-all precision allocation schemes, lacking flexibility and performance-aware adaptability, or introduce high computational costs during dynamic quantization, while failing to adaptively prioritize key-values ​​in long-context tasks. Therefore, they are forced to make suboptimal trade-offs between memory usage, model accuracy, and computational throughput.

[0004] To address these issues, this invention proposes a novel hybrid precision quantization method for KV Cache. Compared to existing hybrid quantization methods, this invention analyzes the importance differences between different model layers at a very low cost, thereby allowing for flexible modification of the quantization configuration based on model performance requirements. This flexibility maximizes the compression ratio of the KV Cache and the throughput of the model while maintaining controllable precision. Summary of the Invention

[0005] The technical problem to be solved by this invention is to address the shortcomings of the prior art by providing a gradient-based layer importance-aware KV Cache mixed-precision quantization method and system. This method analyzes the importance differences of different model layers at a very low cost, thereby allowing flexible modification of the quantization configuration based on model performance requirements. This flexibility can maximize the compression ratio of the KV Cache and the throughput of the model while maintaining controllable accuracy. This addresses the shortcomings of existing KV Cache quantization methods, namely, the lack of flexibility in static allocation, the high computational overhead of dynamic methods, and the inability to prioritize long-context key KVs, thus achieving a dynamic balance between accuracy and efficiency.

[0006] The present invention adopts the following technical solution: A gradient-based method for mixed-precision quantization of layer importance-aware KV cache includes the following steps: S1. Calculate the L2 gradient norm of the Key and Value projected weights of each layer relative to the model loss function, and evaluate the difference in KV contribution between layers; S2. Based on the inter-layer KV contribution difference results obtained in step S1, allocate layer-specific bit widths to the Key and Value of different layers; S3. Based on the distribution characteristics of Key and Value, and according to the bit width allocation scheme obtained in step S2, the Key is quantized per channel and the Value is quantized per token. At the same time, the Key and Value are compressed into a low bit space to obtain the quantized KV cache data. S4. Combining the inter-layer KV contribution difference results obtained in step S1, adaptively retain the full-precision KV pairs of the most recent key tokens, and perform low-bit compression on the old KV pairs using the quantization rules of step S3 to obtain a mixed-precision KV cache set. S5. For the bit width allocation scheme in step S2, the quantization data format in step S3, and the mixed precision KV cache set in step S4, design a CUDA kernel with low bit quantization, dequantization, and matrix-vector multiplication. Optimize computational and memory access overhead through operator fusion to adapt to the inference computation requirements of the mixed precision KV cache.

[0007] Preferably, step S1 specifically includes: S101. Load the full-precision model, sample cues from the target dataset and convert them into input tensors, shift the input tensors one bit to the left to obtain the label tensors, and determine the loss value for each input through forward propagation; S102. For each layer of the model, independently calculate the gradient of the loss relative to the Key projection weights and the Value projection weights, calculate the L2 gradient norm of the gradients respectively, and obtain the importance scores of Key and Value for each layer. S103. Take the average gradient norm across multiple sampling prompts to generate the average importance score of Key and Value for each layer, as the result of the difference in KV contribution between layers.

[0008] Preferably, in step S102, the L2 gradient norm of the i-th layer Key projection weights is expressed as: The L2 gradient norm of the projected weights of Value is expressed as: ,in, Let the Key projection weights be for the i-th layer. Let the Value be the projection weight of the i-th layer. The model loss function; Importance score of each Key level and the importance score of each layer of Value They are respectively:

[0009] in, As a prompt word, The number of prompts, Let L2 be the gradient norm of the i-th layer Key projection weight corresponding to the p-th prompt word. Let L2 be the gradient norm of the i-th layer Value projection weight corresponding to the p-th prompt word.

[0010] Preferably, in step S2, the specific method for allocating the layer-specific bit width is as follows: The layers with the highest importance scores (top 20%-30%) are designated as high-importance layers and assigned 3-bit or 4-bit quantization precision; the remaining layers are designated as low-importance layers and assigned 2-bit quantization precision. The proportion of high-importance layers can be dynamically adjusted.

[0011] Preferably, step S3 specifically includes: S301. Group the Key by channel dimension and reshape the tensor as follows: The grouping format is as follows: each group contains all tokens for a single channel; the Value is grouped by the token dimension, with each group containing all channels for a single token, where B is the batch size. For the number of attention heads, The length of the token sequence. For head dimension; S302. For each group, calculate the scaling factor based on the minimum / maximum value of the elements within the group. ; S303, through Quantize the original elements and convert the quantized values. Cut in Within the range; S304. Store the quantized value in an int32 using bit operations. The dequantization process is performed via... accomplish, The original KV element values ​​within the group. For the original element value The quantized value obtained after quantization It is the minimum value among all the original KV element values ​​within the group.

[0012] Preferably, when the bit width allocation scheme is 3 bits, a specific packing strategy is adopted: The quantized elements are organized into blocks of 11 elements, each stored in a 32-bit integer. The first 10 elements are quantized to 3 bits, and the 11th element is quantized to 2 bits. The clipping range is adjusted according to the element index i within the block. .

[0013] Preferably, step S4 specifically includes: S401. Based on the importance scores of the Key and Value at each level, assign the most recent key context selection ratio to each level. ; S402, Through Calculate the number of RPCs, where, It is the sum of the number of new KV states and the number of historical RPCs at the current time step; S403. Retain the full-precision KV pairs corresponding to RPC, and perform low-bit compression on the old KV pairs of non-RPC using the quantization rules of step S3 to form a mixed-precision KV cache set.

[0014] Preferably, in step S5, optimizing computational and memory access overhead through operator fusion specifically includes: merging the quantization operation of the KV state with the operation of connecting it to the historical KV cache into a CUDA kernel; and merging the dequantization operation of the KV cache with the matrix-vector multiplication operation in the attention computation into another CUDA kernel.

[0015] Preferably, it further includes: designing corresponding CUDA kernels for quantization, dequantization, and matrix-vector multiplication for different quantization bit width configurations.

[0016] Secondly, embodiments of the present invention provide a gradient-based layer importance-aware KV cache mixed-precision quantization system, comprising: The analysis module is used to calculate the L2 gradient norm of the Key and Value projected weights of each layer relative to the model loss function, and to evaluate the differences in KV contributions between layers. A quantization module, connected to the analysis module, is used to allocate layer-specific bit widths to the Key and Value of different layers based on the inter-layer KV contribution difference results. The compression module, connected to the quantization module, is used to compress the KV values ​​to a low bit space by using per-channel quantization for the Key and per-token quantization for the Value based on the distribution characteristics of the Key and Value values ​​and the layer-specific bit width allocation scheme, thereby obtaining quantized KV cache data. The selection module is connected to the analysis module and the compression module respectively. It is used to combine the inter-layer KV contribution difference results, adaptively retain the full-precision KV pairs of the most recent key tokens, and perform low-bit compression on the old KV pairs using the reuse quantization rules to obtain a mixed-precision KV cache set. The optimization module, connected to the quantization module, compression module, and selection module respectively, is used to design a CUDA kernel for low-bit quantization, dequantization, and matrix-vector multiplication based on the bit width allocation scheme, quantization data format, and mixed-precision KV cache set. It optimizes computational and memory access overhead through operator fusion to adapt to inference computing requirements.

[0017] Thirdly, a computer device includes a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor, when executing the computer program, implements the steps of the gradient-based layer importance-aware KV Cache mixed precision quantization method described above.

[0018] Fourthly, embodiments of the present invention provide a computer-readable storage medium including a computer program, which, when executed by a processor, implements the steps of the above-described gradient-based layer importance-aware KV Cache mixed-precision quantization method.

[0019] Fifthly, a chip includes a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor, when executing the computer program, implements the steps of the gradient-based layer importance-aware KV Cache mixed-precision quantization method described above.

[0020] In a sixth aspect, embodiments of the present invention provide an electronic device, including a computer program, which, when executed by the electronic device, implements the steps of the above-described gradient-based layer importance-aware KV Cache mixed precision quantization method.

[0021] Compared with the prior art, the present invention has at least the following beneficial effects: This invention presents a gradient-based layer importance-aware KV cache mixed-precision quantization method. For the first time, it organically integrates four key technical aspects: gradient-based layer importance analysis, asymmetric mixed-precision quantization, dynamic key context selection, and dedicated computational kernel optimization, forming a closed-loop optimization system. Compared to existing static or single dynamic quantization methods, this invention accurately identifies the differences in contribution of different layers to the model output at extremely low cost through gradient analysis, providing a scientific basis for differentiated bit allocation. Based on this, asymmetric quantization is performed by combining the distribution characteristics of Key (significant channel outliers) and Value (critical to attention output), maximizing compression efficiency and controlling error. During inference, the full-precision state of key contexts is dynamically preserved based on importance, achieving an adaptive balance between memory and precision for long sequences. By designing a dedicated CUDA kernel that integrates quantization / dequantization and computation processes, the additional memory access and computational overhead introduced by quantization is fundamentally eliminated, ensuring improved actual inference throughput under extremely low bit quantization. This solution systematically solves the industry problems of inflexible static allocation, high overhead of dynamic methods, and the inability to prioritize key information in long contexts.

[0022] Furthermore, this paper presents an efficient, quantifiable, and theoretically sound method for evaluating layer importance. It directly calculates the weight gradients and measures their magnitude using the L2 norm, accurately reflecting the sensitivity of small perturbations to the final loss function. This method avoids the one-sidedness of relying solely on the magnitude of weights or activation values, evaluating from the fundamental loss function of the model output. This directly links importance assessment to the actual performance of the model, providing a highly reliable and clearly targeted guidance signal for subsequent bit allocation, ensuring that resources are accurately allocated to the layers that have the greatest impact on model accuracy.

[0023] Furthermore, by quantifying layer importance through explicit mathematical expressions, ambiguous descriptions are avoided. The cross-cue average score calculation method further enhances the stability of importance ranking, making the division of high and low importance layers more precise. This provides solid theoretical support for layer-specific bit allocation, ensuring that quantization resources are tilted towards key layers, minimizing accuracy loss.

[0024] Furthermore, the 20%-30% high-importance layer ratio, based on experimental verification, achieves an optimal trade-off between accuracy and memory usage. Allocating 3 bits / 4 bits to high-importance layers ensures accuracy, while allocating 2 bits to low-importance layers saves memory, specifically addressing the issue of varying contributions from different layers to the model output. The dynamically adjustable ratio design allows users to customize the configuration according to their specific scenarios, adapting to diverse deployment needs. This achieves precise resource allocation, resulting in lower accuracy loss and lower memory consumption for the same memory compression rate.

[0025] Furthermore, grouping keys by channel isolates the impact of outliers, while grouping values ​​by token preserves the integrity of key tokens, effectively reducing quantization errors. The scaling factor is calculated based on intra-group extrema, adapting to different data distributions and improving quantization accuracy. The design of storing quantized values ​​in int32 improves compression density, and the dequantization formula is simple and efficient, reducing computational overhead. This solves the accuracy loss problem caused by existing quantization methods not considering KV distribution differences, maintaining high model performance even at low bit depths.

[0026] Furthermore, by using 11-element block packing, the compression density is improved by 10% compared to traditional 3-bit quantization, further reducing memory usage; the clipping range is adjusted according to the element index within the block to adapt to the quantization requirements of different elements and reduce the precision loss caused by clipping; it seamlessly integrates with the bit allocation scheme, ensuring both accuracy and maximizing memory savings in high-importance 3-bit quantization scenarios.

[0027] Furthermore, the dynamic key context selection mechanism enables intelligent and adaptive management of the number of full-precision key-value pairs during long sequence inference, effectively alleviating the attention sink problem and linking it with layer importance.

[0028] Furthermore, the fusion of quantization and connection integrates two independent operations into a single CUDA kernel, streams elements and directly appends them to the history cache, reducing the number of memory accesses; the fusion of dequantization and matrix-vector multiplication enables instant dequantization and accumulation of results, avoiding increased memory usage caused by premature dequantization; the design of CUDA thread block parallel processing and shared memory cache enhances data locality and improves computational efficiency.

[0029] Furthermore, the data packing format, bit operation logic, and computational intensity differ for different bit widths. Using a general-purpose or single kernel to handle all cases often leads to performance loss. This invention tailors multiple kernels for its mixed-precision scheme, enabling each kernel to perform optimized thread scheduling, memory access, and computation instruction generation for a specific bit width. This refined kernel design ensures that every quantization configuration, from high-bit to low-bit levels, can be executed on the GPU with near-hardware-limit efficiency, thus fully releasing the performance advantages of the entire mixed-precision scheme and avoiding overall performance degradation due to kernel efficiency bottlenecks.

[0030] It is understood that the beneficial effects of the second to sixth aspects mentioned above can be found in the relevant descriptions in the first aspect mentioned above, and will not be repeated here.

[0031] In summary, this invention achieves low-cost and accurate guidance through gradient analysis; achieves optimal memory-accuracy balance through asymmetric quantization and dynamic selection; and eliminates quantization overhead with the help of a dedicated fusion kernel. Ultimately, it significantly reduces memory usage while greatly improving inference speed with minimal loss of accuracy.

[0032] The technical solution of the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. Attached Figure Description

[0033] Figure 1 This is a flowchart of the gradient-based layer importance-aware KV Cache mixed-precision quantization method in the invention. Figure 2 This is a diagram illustrating the impact of using 2-bit quantization on model accuracy at different layers in this invention. Figure 3 This is a heatmap of the weights of the KV projection matrices at different layers in the Llama 2-7B model of this invention. Figure 4 This is an overall flowchart of the KVmix profiler framework proposed in this invention; Figure 5 This is an illustration of the efficient 3-bit quantization method in this invention; Figure 6 This is an illustration of the dynamic key context selection strategy in this invention; Figure 7 A graph showing the changes in model precision, memory usage, and accuracy under high and low bit quantization configurations using different proportions. Figure 8 This is a graph showing the quantization bit configuration results of different models analyzed using the KVmix profiler in this invention; Figure 9 This is a graph showing the memory compression ratio of the model using different quantization methods in this invention; Figure 10 This is a graph showing the inference throughput of the model using different quantization methods in this invention. Figure 11 A schematic diagram of a computer device provided in an embodiment of the present invention; Figure 12 This is a block diagram of a chip provided according to an embodiment of the present invention.

[0034] Among them, 60. Computer equipment; 61. Processor; 62. Memory; 63. Computer program; 600. Electronic device; 610. Processing unit; 620. Storage unit; 6201. Random access memory unit; 6202. Cache memory unit; 6203. Read-only memory unit; 6204. Program / utility; 6205. Program module; 630. Bus; 640. Display unit; 650. Input / output interface; 660. Network adapter; 700. External device. Detailed Implementation

[0035] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of the present invention. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0036] In the description of this invention, it should be understood that the terms "comprising" and "including" indicate the presence of the described features, integrals, steps, operations, elements and / or components, but do not exclude the presence or addition of one or more other features, integrals, steps, operations, elements, components and / or collections thereof.

[0037] It should also be understood that the terminology used in this specification is for the purpose of describing particular embodiments only and is not intended to limit the invention. As used in this specification and the appended claims, the singular forms “a,” “an,” and “the” are intended to include the plural forms unless the context clearly indicates otherwise.

[0038] It should also be further understood that the term "and / or" as used in this specification and the appended claims refers to any combination and all possible combinations of one or more of the associated listed items, and includes such combinations. For example, A and / or B can represent three cases: A alone, A and B simultaneously, and B alone. Additionally, the character " / " in this invention generally indicates that the preceding and following objects have an "or" relationship.

[0039] It should be understood that although terms such as first, second, third, etc., may be used in the embodiments of the present invention to describe the preset range, these preset ranges should not be limited to these terms. These terms are only used to distinguish the preset ranges from one another. For example, without departing from the scope of the embodiments of the present invention, the first preset range may also be referred to as the second preset range, and similarly, the second preset range may also be referred to as the first preset range.

[0040] Depending on the context, the word "if" as used here can be interpreted as "when," "when," "in response to determination," or "in response to detection." Similarly, depending on the context, the phrase "if determination" or "if detection (of the stated condition or event)" can be interpreted as "when determination," "in response to determination," "when detection (of the stated condition or event)," or "in response to detection (of the stated condition or event)."

[0041] The accompanying drawings illustrate various structural schematic diagrams according to embodiments disclosed in this invention. These drawings are not to scale, and some details have been enlarged for clarity, and some details may have been omitted. The shapes of the various regions and layers shown in the drawings, as well as their relative sizes and positional relationships, are merely exemplary and may deviate from reality due to manufacturing tolerances or technical limitations. Furthermore, those skilled in the art can design regions / layers with different shapes, sizes, and relative positions as needed.

[0042] This invention provides a gradient-based layer importance-aware KV cache mixed-precision quantization method. First, for the Key-Value (KV) cache of a large language model LLM, layer importance analysis is performed using the gradient of KV weights relative to the model loss to evaluate the impact of the Key and Value projection matrices of each layer on the model loss, achieving layer-specific bit width allocation. Based on the distribution characteristics of Key and Value values, asymmetric low-bit quantization is applied to Key and Value. A dynamic key context selection strategy is used, combined with the KV importance analysis results, to adaptively retain the most recent full-precision KV pairs of key tokens while aggressively compressing older KV pairs with low bits. An efficient low-bit quantization method and CUDA kernel are designed to optimize computational overhead. This invention effectively solves the memory bottleneck problem of KV cache in resource-constrained environments, providing users with a flexible quantization strategy that balances memory, precision, and throughput, thereby improving the inference performance of the model.

[0043] Please see Figure 1 This invention discloses a gradient-based layer importance-aware KV cache mixed-precision quantization method, comprising the following steps: S1, KV Importance Analysis: The difference in KV contribution between layers is evaluated by calculating the L2 gradient norm of the Key and Value projected weights of each layer relative to the model loss function; Current KV cache quantization methods rely on a fixed quantization strategy, ignoring the differences in the contribution of KV values ​​to the final output at different layers. To demonstrate that quantizing keys or values ​​at different layers has different effects on the model, 2-bit quantization was selectively applied to keys or values ​​at different layers, and the impact on model accuracy was evaluated. The results are presented in... Figure 2 The results show that quantizing keys or values ​​at different layers has varying impacts on the generated model quality. However, effectively analyzing the contribution differences between model layers to allocate different quantization bits to keys or values ​​remains a key challenge that needs to be addressed.

[0044] In each layer of the KV Cache, the calculation process is as follows: At time step t, the i-th layer receives the hidden state from the output of the previous layer. These hidden states are used to calculate the K and V of the current token as follows: and ,in, and These are the K and V projection weights of the i-th layer. After calculation, the calculated... and By concatenating with the previously stored K and V data, a complete K and V sequence up to time step t is generated: and .

[0045] This calculation process shows that, and This determines how the Key and Value are extracted from the hidden state, directly affecting the quality of the KV pairs generated by the attention mechanism and the contribution of this layer to the model output. Figure 3 The Llama 2-7B model is provided. and Heatmap. The heatmap reveals two key points: ① The KV weight values ​​are significantly different between different layers.

[0046] ② The KV weight distribution patterns are different within the same layer.

[0047] For key-value pairs (KVs), their values ​​dynamically adapt to changes in the input; however... and The weights are learned during the model's training phase and remain static during inference. Therefore, these weights can be used to evaluate the differences in importance between Keys and Values ​​across layers. In the attention mechanism, Key-Value pairs are computed by applying a linear transformation of the Key-Value projection weight matrix to the hidden state of the previous layer, and then combining it with the Query vector to compute the attention output. Therefore, the magnitude of the Key-Value projection weights alone is insufficient to measure the importance of Key-Value pairs across layers; a more precise evaluation metric is needed. This metric should quantify the sensitivity of each layer's Keys and Values ​​to the model's loss function L. Using the chain rule, we obtain:

[0048] here, and They represent L relative to and The gradient. and Let L represent the partial derivatives of L with respect to the Key and Value matrices. To quantize the perturbation, assume that the quantization operation introduces a small perturbation into the Key and Value matrices, i.e. and ,in K and V represents the quantization error.

[0049] The change in L caused by quantization is This is the difference between the original loss and the quantized loss. (For approximation...) L, a first-order Taylor expansion around Key:

[0050] Therefore, the change in loss due to quantization is:

[0051] and It was calculated previously. and ,so ,in This represents the inner product. For a fixed... K, larger (Taking the L2 norm) will amplify L. Therefore, the weight gradient norm reflects the influence of Key on L, because approximate (Scaled by the hidden state H). The principle is the same for Value and Key.

[0052] Based on the above analysis, we propose a gradient-based method, KVmix profiler, to quantify the contribution of each layer's Key and Value to the model output, thereby achieving a smart mixed-precision quantization strategy.

[0053] Specifically, the L2 norm of the model loss function L is calculated relative to the Key and Value projection weight matrices of each model layer. and KVmixprofiler captures the dynamic sensitivity of these parameters during model inference and provides layer-specific importance metrics to support effective mixture precision quantification in subsequent inference stages.

[0054] The implementation of the KVmix profiler includes the following key steps: ① Data Preparation and Forward Propagation. Load the full-precision model and randomly sample multiple cues from the target dataset as input data. These cues are converted into input tensors using a tokenizer. Leveraging the autoregressive properties of LLMs, each input tensor is shifted left by one bit to serve as the corresponding label tensor for calculating the model's loss function. Subsequently, the loss value for each input is determined through forward propagation.

[0055] ② Gradient calculation and importance assessment. For layer i of the model, the loss is calculated independently relative to the Key ( ) and Value ( The gradient of the projected weights.

[0056] The process begins by calculating the gradient through backpropagation, i.e. and Then, the magnitude of these gradients is evaluated using the L2 norm, i.e. and .

[0057] The importance scores for Keys and Values ​​at each level are represented as follows: Larger or This indicates that the Key or Value in layer i has a greater impact on the model output.

[0058] To enhance the reliability of the assessment, the average gradient norm can be taken across multiple cue words (p) to generate an average importance score for the Key and Value at each layer (where P is the number of cue words):

[0059] S2. Mixed precision quantization: Based on the specific bit width of the importance allocation layer, higher precision is used for critical layers and lower precision is used for non-critical layers; Importance scores are used to categorize the importance of Key and Value components across all model layers. and The first 20% of layers use high-bit quantization (e.g., 3-bit or 4-bit), while the remaining 80% of layers use more aggressive low-bit quantization (e.g., 2-bit). This 20%-80% division is not fixed and can be dynamically adjusted according to needs to balance model accuracy and memory usage. Increasing the proportion of low-bit quantization layers can further reduce the memory usage of the KV cache, but may sacrifice some inference accuracy. ③ Model Configuration and Inference. The KV quantization configuration results obtained from the above steps are incorporated into the model configuration, making the quantized model usable for inference. Detailed workflow is as follows... Figure 4 As shown. This process is performed offline, so it does not affect inference efficiency. Moreover, the KVmix profiler analysis is executed once, allowing the model to reuse the analysis results for subsequent inference tasks.

[0060] S3. Asymmetric Quantization and Low-Bit Compression: Based on the distribution characteristics of Key and Value, per-channel quantization is used for Key and per-token quantization for Value. Simultaneously, their values ​​are compressed into a low-bit space as much as possible, with 2-bit and 3-bit values ​​mixed and compressed into int32 to improve compression density. The key and value are grouped and quantized using per-channel and per-token methods respectively. The shape of the KV cache is as follows: Where B is the batch size, nh is the number of attention heads, T is the token sequence length, and D is the head dimension. When the key is quantized by channel (D), the tensor is reshaped into... Each group contains all tokens for a single channel. This approach is inspired by the distribution characteristics of the Key Cache, where significant outliers exist in the channel dimension, with some channels exhibiting large outliers. Channel-wise quantization isolates errors within each channel and prevents outliers from affecting other channels. When Value is quantized by token, the shape of the tensor remains unchanged, and each group contains all channels for a single token. Unlike the Key Cache, the Value Cache does not have obvious outliers, but it plays a crucial role in computing attention output. Token-wise quantization confines errors to a single token, preserving the integrity of other important tokens. This asymmetric quantization strategy effectively reduces the errors introduced during KV Cache quantization. Additionally, grouped low-bit quantization minimizes KV Cache memory usage. The process includes: S301, Calculate the scaling factor s For each group (key by channel, value by token), calculate using the group's minimum / maximum value. ,in This represents the maximum quantization value.

[0061] S302, Element Quantization use Quantize the element, where x represents the original element value and q is the quantized value.

[0062] S303, Cutting Restrict q to .

[0063] S304, Storage and Inverse Quantization The quantized value is stored in a 32-bit integer (int32) using bit operations.

[0064] For 4-bit, 2-bit, and 1-bit, the number of elements in each int32 is: Inverse quantization via implement.

[0065] For 3-bit quantization, a new packing strategy is introduced to maximize the memory compression ratio.

[0066] The quantized elements are organized into blocks of 11 elements, each block is stored in a 32-bit integer, the first 10 elements are quantized to 3 bits, and the 11th element is quantized to 2 bits.

[0067] The cropping range is adjusted based on the element index: Where i is the index of the element within the block, such as Figure 5 As shown, this strategy improves the compression packing density by 10% compared to uniform 3-bit quantization (where each int32 can only hold 10 elements).

[0068] S4. Dynamic Key Context Selection: Adaptively retains the most recent key token's full-precision key-value pair and compresses older key-value pairs; Based on the importance analysis results provided by the KVmix profiler, a dynamic nearest critical context (RPC) selection strategy is proposed.

[0069] Specifically, for layer i, based on and The score is assigned an RPC selection ratio r, the higher the ratio. and This leads to a larger r. The number of RPCs is through calculate. It is the sum of the number of new KV states and the number of historical RPCs at the current time step. Based on Select an appropriate number of key-value pairs as the RPC. Retain full precision for the RPC, while performing mixed quantization on less critical and older key-value pairs, such as... Figure 6 As shown. This strategy ensures a dynamic reduction in the number of full-precision RPCs during long context inference, thus avoiding excessive memory pressure caused by retaining a large number of full-precision key-value pairs, while maintaining high-quality sequence generation. Furthermore, since the importance of keys and values ​​may differ within the same layer, the RPC selection ratio for keys and values ​​will differ accordingly within that layer. The RPC selection ratio can be adjusted to balance accuracy and memory: increasing it will improve accuracy but require more memory.

[0070] S5 and CUDA kernel optimization: Design efficient 2, 3, and 4-bit quantization and gemv kernels to reduce computational overhead, while using operator fusion computation to reduce memory access overhead.

[0071] During model inference, quantization of the KV cache introduces additional overhead due to quantization and dequantization operations. To improve inference efficiency, an efficient CUDA kernel is designed for quantization, dequantization, and matrix-vector multiplication.

[0072] S501, the fusion of quantization and connectivity During the decoding phase, the current layer's key-value (KV) state is joined with the historical KV cache. Joining the previous KV cache with the current KV would incur additional memory access overhead. Quantization and this join are integrated into a single CUDA kernel, processing each element in a streaming manner. KV elements are quantized and directly appended to the historical KV cache, reducing memory accesses. CUDA thread blocks process tokens in parallel, using shared memory to cache intermediate results to enhance data locality.

[0073] S502, Fusion of Dequantization and Matrix-Vector Multiplication In attention computation, quantized key-value pairs (KV) require dequantization before matrix-vector multiplication. Pre-dequantizing the complete KV increases memory usage. Integrating dequantization with multiplication—dequantizing each element on the fly and immediately multiplying it by its corresponding element and summing the results—minimizes memory overhead.

[0074] S503, multi-bit quantization configuration To support various quantization bit widths in KVmix, CUDA kernels for 1, 2, 3, and 4-bit quantization were developed, along with matrix-vector multiplication kernels tailored to each configuration, ensuring compatibility across bit width quantization.

[0075] In another embodiment of the present invention, a gradient-based layer importance-aware KV Cache mixed-precision quantization system is provided. This system can be used to implement the above-mentioned gradient-based layer importance-aware KV Cache mixed-precision quantization method. Specifically, the gradient-based layer importance-aware KV Cache mixed-precision quantization system includes an analysis module, a quantization module, a compression module, a selection module, and an optimization module.

[0076] The analysis module is used to calculate the L2 gradient norm of the Key and Value projection weights of each layer relative to the model loss function, and to evaluate the differences in KV contributions between layers. A quantization module, connected to the analysis module, is used to allocate layer-specific bit widths to the Key and Value of different layers based on the inter-layer KV contribution difference results. The compression module, connected to the quantization module, is used to compress the KV values ​​to a low bit space by using per-channel quantization for the Key and per-token quantization for the Value based on the distribution characteristics of the Key and Value values ​​and the layer-specific bit width allocation scheme, thereby obtaining quantized KV cache data. The selection module is connected to the analysis module and the compression module respectively. It is used to combine the inter-layer KV contribution difference results, adaptively retain the full-precision KV pairs of the most recent key tokens, and perform low-bit compression on the old KV pairs using the reuse quantization rules to obtain a mixed-precision KV cache set. The optimization module, connected to the quantization module, compression module, and selection module respectively, is used to design a CUDA kernel for low-bit quantization, dequantization, and matrix-vector multiplication based on the bit width allocation scheme, quantization data format, and mixed-precision KV cache set. It optimizes computational and memory access overhead through operator fusion to adapt to inference computing requirements.

[0077] This invention provides a terminal device comprising a processor and a memory. The memory stores a computer program, which includes program instructions. The processor executes the program instructions stored in the computer storage medium. The processor may be a Central Processing Unit (CPU), or other general-purpose processors, graphics processing units (GPUs), tensor processing units (TPUs), digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. It is the computing and control core of the terminal, suitable for implementing one or more instructions, specifically suitable for loading and executing one or more instructions to achieve a corresponding method flow or function. The processor described in this embodiment can be used for the operation of a gradient-based layer importance-aware KV cache mixed-precision quantization method, including: The algorithm calculates the L2 gradient norm of the projected weights of Key and Value in each layer relative to the model loss function to evaluate the differences in KV contributions between layers. Based on these differences, it allocates layer-specific bit widths to Key and Value in different layers. Based on the distribution characteristics of Key and Value, and according to the bit width allocation scheme, it performs per-channel quantization on Key and per-token quantization on Value, while compressing Key and Value into a low-bit space to obtain quantized KV cache data. Combining the inter-layer KV contribution differences, it adaptively retains the full-precision KV pairs of the most recent key tokens and performs low-bit compression on older KV pairs using the reuse quantization rules, resulting in a mixed-precision KV cache set. For the bit width allocation scheme, quantized data format, and mixed-precision KV cache set, it designs a CUDA kernel for low-bit quantization, dequantization, and matrix-vector multiplication. Operator fusion optimizes computational and memory access overhead to meet the inference computation requirements of the mixed-precision KV cache.

[0078] Please see Figure 11The terminal device is a computer device. In this embodiment, the computer device 60 includes a processor 61, a memory 62, and a computer program 63 stored in the memory 62 and executable on the processor 61. When executed by the processor 61, the computer program 63 implements the gradient-based layer importance-aware KV Cache mixed-precision quantization method described in this embodiment. To avoid repetition, details are omitted here. Alternatively, when executed by the processor 61, the computer program 63 implements the functions of each model / unit in the gradient-based layer importance-aware KV Cache mixed-precision quantization system described in this embodiment. To avoid repetition, details are omitted here.

[0079] Computer device 60 can be a desktop computer, laptop, handheld computer, cloud server, or other computing device. Computer device 60 may include, but is not limited to, a processor 61 and a memory 62. Those skilled in the art will understand that... Figure 11 This is merely an example of computer device 60 and does not constitute a limitation on computer device 60. It may include more or fewer components than shown, or combine certain components, or different components. For example, computer device may also include input / output devices, network access devices, buses, etc.

[0080] The processor 61 may be a Central Processing Unit (CPU), or other general-purpose processors, graphics processing units (GPUs), tensor processing units (TPUs), digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. A general-purpose processor may be a microprocessor or any conventional processor.

[0081] The memory 62 can be an internal storage unit of the computer device 60, such as a hard disk or memory of the computer device 60. The memory 62 can also be an external storage device of the computer device 60, such as a plug-in hard disk, smart media card (SMC), secure digital (SD) card, flash card, etc. equipped on the computer device 60.

[0082] Furthermore, the memory 62 may include both internal storage units of the computer device 60 and external storage devices. The memory 62 is used to store computer programs and other programs and data required by the computer device. The memory 62 can also be used to temporarily store data that has been output or will be output.

[0083] Please see Figure 12 The terminal device is an electronic device 600, which is manifested in the form of a general-purpose computing device. The components of the electronic device may include, but are not limited to: at least one processing unit 610, at least one storage unit 620, a bus 630 connecting different platform components (including storage unit 620 and processing unit 610), a display unit 640, etc.

[0084] The storage unit stores program code, which can be executed by the processing unit 610 to perform the steps described in the method section of this specification according to various exemplary embodiments of the present invention. For example, the processing unit 610 can perform actions such as... Figure 1 The steps are shown in the figure.

[0085] Storage unit 620 may include a readable medium in the form of a volatile storage unit, such as random access memory (RAM) 6201 and / or cache memory 6202, and may further include a read-only memory (ROM) 6203.

[0086] Storage unit 620 may also include a program / utility 6204 having a set (at least one) program module 6205, such program module 6205 including but not limited to: operating system, one or more application programs, other program modules and program data, each or some combination of these examples may include an implementation of a network environment.

[0087] Bus 630 can represent one or more of several types of bus structures, including a memory cell bus or memory cell controller, a peripheral bus, a graphics acceleration port, a processing unit, or a local bus using any of the multiple bus structures.

[0088] Electronic device 600 can also communicate with one or more external devices 700 (e.g., keyboard, pointing device, Bluetooth device, etc.), and with one or more devices that enable a user to interact with electronic device 600, and / or with any device that enables electronic device 600 to communicate with one or more other computing devices (e.g., router, modem). This communication can be performed via input / output interface 650. Furthermore, electronic device 600 can also communicate with one or more networks (e.g., local area network, wide area network, and / or public network, such as the Internet) via network adapter 660. Network adapter 660 can communicate with other modules of electronic device 600 via bus 630. It should be understood that, although not shown in the figures, other hardware and / or software modules can be used in conjunction with electronic device 600, including but not limited to: microcode, device drivers, redundant processing units, external disk drive arrays, RAID systems, tape drives, and data backup storage platforms.

[0089] Example 4 This invention also provides a storage medium, specifically a computer-readable storage medium, which is a memory device in a terminal device for storing programs and data. It is understood that the computer-readable storage medium here can include both built-in storage media in the terminal device and extended storage media supported by the terminal device; it can be any tangible medium containing or storing a program that can be used by or in conjunction with an instruction execution system, apparatus, or device. The computer-readable storage medium provides storage space that stores the terminal's operating system. Furthermore, the storage space also stores one or more instructions suitable for loading and execution by a processor, which can be one or more computer programs (including program code). More specific examples of the computer-readable storage medium include: an electrical connection with one or more wires, a portable disk, a hard disk, random access memory, read-only memory, erasable programmable read-only memory, optical fiber, portable compact disk read-only memory, optical storage device, magnetic storage device, or any suitable combination thereof.

[0090] Computer-readable storage media also include data signals propagated in baseband or as part of a carrier wave, carrying readable program code. Such propagated data signals can take various forms, including but not limited to electromagnetic signals, optical signals, or any suitable combination thereof. A readable storage medium can also be any readable medium other than a readable storage medium that can send, propagate, or transmit a program for use by or in connection with an instruction execution system, apparatus, or device. The program code contained on the readable storage medium can be transmitted using any suitable medium, including but not limited to wireless, wired, optical fiber, radio frequency, etc., or any suitable combination thereof.

[0091] Program code for performing the operations of this invention can be written in any combination of one or more programming languages, including object-oriented programming languages ​​such as Java and C++, and conventional procedural programming languages ​​such as C or similar languages. The program code can execute entirely on the user's computing device, partially on the user's computing device, as a standalone software package, partially on the user's computing device and partially on a remote computing device, or entirely on a remote computing device or server. In cases involving remote computing devices, the remote computing device can be connected to the user's computing device via any type of network, including a local area network (LAN) or a wide area network (WAN), or it can be connected to an external computing device (e.g., via the Internet using an Internet service provider).

[0092] One or more instructions stored in a computer-readable storage medium can be loaded and executed by a processor to implement the corresponding steps of the gradient-based layer importance-aware KV Cache mixed-precision quantization method in the above embodiments; one or more instructions in the computer-readable storage medium are loaded and executed by the processor to perform the following steps: The algorithm calculates the L2 gradient norm of the projected weights of Key and Value in each layer relative to the model loss function to evaluate the differences in KV contributions between layers. Based on these differences, it allocates layer-specific bit widths to Key and Value in different layers. Based on the distribution characteristics of Key and Value, and according to the bit width allocation scheme, it performs per-channel quantization on Key and per-token quantization on Value, while compressing Key and Value into a low-bit space to obtain quantized KV cache data. Combining the inter-layer KV contribution differences, it adaptively retains the full-precision KV pairs of the most recent key tokens and performs low-bit compression on older KV pairs using the reuse quantization rules, resulting in a mixed-precision KV cache set. For the bit width allocation scheme, quantized data format, and mixed-precision KV cache set, it designs a CUDA kernel for low-bit quantization, dequantization, and matrix-vector multiplication. Operator fusion optimizes computational and memory access overhead to meet the inference computation requirements of the mixed-precision KV cache.

[0093] The databases involved in the embodiments provided in this application may include at least one type of relational database and non-relational database. Non-relational databases may include, but are not limited to, blockchain-based distributed databases. The processors involved in the embodiments provided in this application may be general-purpose processors, central processing units, graphics processing units, digital signal processors, programmable logic devices, quantum computing-based data processing logic devices, etc., and are not limited to these.

[0094] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of the present invention. The components of the embodiments of the present invention described and shown in the accompanying drawings can generally be arranged and designed in various different configurations. Therefore, the following detailed description of the embodiments of the present invention provided in the accompanying drawings is not intended to limit the scope of the claimed invention, but merely to illustrate selected embodiments of the invention. All other embodiments obtained by those skilled in the art based on the embodiments of the present invention without inventive effort are within the scope of protection of the present invention.

[0095] The proposed method was evaluated using Llama 2-7B-hf, Llama 3-8B-Instruct, Llama 3.1-8B, Mistral-7B-Instruct-v0.3, and Falcon-7B models. The datasets were selected based on three different evaluation schemes: ① Long Context Evaluation: The LongBench benchmark is used to evaluate performance on long context tasks. It covers several key long text application scenarios. Due to limited GPU memory, the maximum sequence length is set to 4096.

[0096] ② Language modeling: The perplexity on the Wikitext-2 dataset was measured to evaluate its language modeling capabilities.

[0097] ③ Mathematical Reasoning: The GSM8K dataset was used to evaluate the model's performance on mathematical reasoning tasks. An NVIDIA RTX 4090 GPU (24GB) was used to evaluate the model's inference efficiency and the compression ratio of the KV cache.

[0098] Layer importance analysis results: 3-bit and 2-bit mixed quantization was used for the Key, and 4-bit and 2-bit mixed quantization was used for the Value. When KV was quantized to 3-bit or 4-bit, the RPC ratio was set to 20%; for 2-bit quantization, the RPC ratio was set to 10%. When the RPC ratio exceeds 20%, its contribution to accuracy is negligible; therefore, 20% was chosen as the high-bit configuration for KVmix. The quantization group size was 32. Thirty prompts were selected from LongBench for KV importance analysis. The KV bit configuration and RPC ratio for each layer can be obtained using the KVmix profiler. When using the KVmix profiler, randomly selecting 20 to 30 prompts is sufficient to produce reliable importance analysis results; additional prompts do not significantly change the results. For the model and experimental environment used in this work, this process only takes 10 to 15 minutes, highlighting the efficiency of the KVmix profiler. Users can flexibly customize the quantization configuration by adjusting the ratio of different bit width layers in the KVmix profiler to meet different accuracy or memory requirements. Figure 7 This demonstrates the trends in model accuracy, key-value memory usage, and inference throughput as the quantization ratios to 3-bit and 4-bit models are varied. An optimal trade-off among these three metrics is achieved when the quantization ratio to 3-bit and 4-bit models is set to 20%. Under this configuration, the average quantization bit width for the key is 2.19 (exact value: 2.1875), and for the value it is 2.38 (exact value: 2.375). Detailed configuration obtained using KVmixprofiler is shown below. Figure 8 Unless otherwise specified, this configuration applies to k-2.19v2.38 quantization in subsequent experiments.

[0099] Long Context Performance Evaluation: Evaluate the performance of various quantization configurations on eight different datasets in the LongBench benchmark, using the FP16 model as the baseline.

[0100] Table 1. Inference accuracy of KVmix using different models and different parameter configurations based on LongBench in this invention.

[0101] Detailed experimental results are presented in Table 1, where KVmix-k2.19v2.38 was used. Figure 8 The configurations are as follows: KVmix-2bit uses asymmetric 2-bit quantization (key per channel and value per token) for all model layers. random-k2.19v2.38 randomly selects 20% of the model layers for mixed-precision quantization. KVmix-k2.19v2.38w / oRPC is the case where KVmix-k2.19v2.38 does not use RPC. Results show that KVmix-k2.19v2.38 has an average accuracy loss (average of four different models) of 1.67% relative to the FP16 baseline. In contrast, KVmix-2bit achieves an average accuracy loss of 4.53% relative to the baseline.

[0102] Furthermore, randomly selecting a high-bit quantization layer (random-k2.19v2.38) resulted in an average accuracy loss of 4.06%. This accuracy loss is significantly higher than that of KVmix-k2.19v2.38, demonstrating the advantages of KV importance-aware hybrid quantization. The average accuracy of KVmix-k2.19v2.38 with / o RPC decreased by 3.28% relative to KVmix-k2.19v2.38, proving the effectiveness of using RPC to improve model accuracy. KVmix was compared with previous state-of-the-art (SOTA) methods prior to KV caching, particularly the key-per-channel and value-per-token methods, namely KIVI and KVQuant, as they are known to minimize KV quantization errors.

[0103] Furthermore, KVmix was compared with the state-of-the-art (SOTA) method QJL. Table 2 shows the accuracy test results, where KIVI-2bit-r64 uses 2-bit quantization with a full-precision residual of length 64; KVQuant-3bit-1% uses 3-bit quantization with 1% outlier handling; QJL-3bit uses 3-bit quantization; and KVmix-k2.28v2.56 increases the proportion of high-bit quantization layers in KVmix-k2.19v2.38 to 30%.

[0104] Table 2 Performance comparison of KVmix with other methods in this invention

[0105] The results show that KVmix-k2.19v2.38 outperforms KIVI-2bit-r64 and QJL-3bit, reducing the average accuracy loss by 1.50% and 0.68%, respectively. While KVQuant-3bit-1% achieves comparable accuracy to KVmix-k2.19v2.38, its memory compression ratio and inference efficiency are inferior. Figure 9 , Figure 10 By increasing the quantization bit width, KVmix-k2.28v2.56 demonstrates a significant advantage in accuracy over KVQuant-3bit-1%, while maintaining comparable memory compression (4.8×) and better inference acceleration (5.23×). This flexibility in balancing accuracy and quantization bit width is a key advantage of KVmix.

[0106] Evaluation on GSM8K and Wikitext-2: Using the FP16 model as a baseline, the capabilities of the quantized model in language modeling and mathematical reasoning were evaluated. The evaluation was conducted using the lm_eval framework, where the quantized model replaced the Hugging Face model.

[0107] Specifically, the accuracy on the GSM8K dataset and the perplexity on the Wikitext-2 dataset were measured. The experimental results are shown in Table 3, where 2-bit (kT, vT) symmetric 2-bit (Key per token and Value per token) quantization was used for all model layers; 4-bit (kT, vT) symmetric 4-bit quantization was used; and their RPC ratios were all set to 0.

[0108] Table 3. Performance of KVmix on GSM8K and Wikitext-2 datasets in this invention.

[0109] The results show that the 2-bit (kT, vT) model suffered a catastrophic performance loss on GSM8K and Wikitext-2, almost losing its reasoning ability. For the 4-bit (kT, vT) model, the performance loss on GSM8K and Wikitext-2 also reached 9.17% and 5.28%, respectively. In contrast, on Wikitext-2, the perplexity score of KVmix-k2.19v2.38 was almost comparable to the baseline, while on the more challenging GSM8K mathematical reasoning task, KVmix-k2.19v2.38 achieved an accuracy loss of only 2.00%, significantly outperforming both 2-bit (kT, vT) and 4-bit (kT, vT). Furthermore, on GSM8K, KVmix-k2.19v2.38 demonstrated a significant accuracy improvement compared to KVmix-2bit and random-k2.19v2.38, which do not utilize KV importance analysis for more accurate quantization. Compared to Atom-4bit and other state-of-the-art methods, KVmix-k2.19v2.38 also exhibits a significant accuracy advantage. It's worth noting that Atom-4bit performs 4-bit quantization on both model weights and activations, leading to a greater loss of accuracy. These results demonstrate the superior performance of KVmix-k2.19v2.38 in complex task inference.

[0110] Inference Efficiency and Memory Usage Evaluation: KVmix's inference throughput and memory usage during inference were evaluated. To ensure fairness, the same input data was used across all evaluation methods. The number of input tokens was 688, the maximum number of new tokens was set to 1024, and the model used was Llama 2-7B-hf. KVmix was compared with KIVI-2bit-r64, KVQuant-3bit-1%, QJL-3bit, and Atom-4bit. Memory usage results are as follows: Figure 9 As shown, the batch size is fixed at 4.

[0111] The memory usage in the experimental results represents the peak memory usage during inference minus the memory occupied by the model before inference. To fully utilize GPU memory, the batch size was gradually increased to explore the maximum throughput of KVmix. The throughput results are as follows: Figure 10As shown, the baseline (FP16), Aotm-4bit, and KIVI-2bit-r64 reached memory overflow at batch sizes of 4, 18, and 28, respectively, while KVmix-k2.19v2.38 could achieve a maximum batch size of 30 and an inference throughput of 1032 tokens / sec. The results indicate that KVmix-k2.19v2.38 achieves a 4.9× reduction in memory usage and a 5.3× increase in throughput compared to the baseline. This effective memory compression stems from KVmix's extremely low bit quantization and dynamic RPC strategy, which gradually reduces the number of full-precision key-value pairs during inference. In contrast, KIVI employs a fixed full-precision residual strategy and cannot dynamically reduce the number of full-precision key-value pairs. Therefore, despite using Key-2.19 and Value-2.38 bit quantization, KVmix saves more memory than KIVI-2bit. Meanwhile, Atom quantizes model weights and activations and utilizes tensor cores for optimized kernel computation, achieving higher throughput at the same batch size, but resulting in a greater decrease in model accuracy (Table 3). While KVQuant achieves significant memory compression, its inference efficiency is hampered by substantial preprocessing requirements. QJL achieves “zero-overhead” quantization by eliminating the storage of extra constants such as zeros and scaling factors. This allows it to achieve slightly better memory compression than KVmix, but its inference efficiency and accuracy are lower than KVmix.

[0112] In summary, this invention presents a gradient-based layer importance-aware KV cache hybrid precision quantization method and system, solving the memory bottleneck problem caused by KV cache in LLM. This method creatively integrates layer importance analysis based on KV weight gradients into KV quantization and incorporates dynamic long context optimization, maintaining inference quality while reducing memory usage. It achieves significant memory compression and inference efficiency improvement with minimal accuracy loss and provides flexible selection of quantization strategies adaptable to various scenarios.

[0113] Those skilled in the art will clearly understand that, for the sake of convenience and brevity, the above-described division of functional units and modules is merely an example. In practical applications, the above functions can be assigned to different functional units and modules as needed, that is, the internal structure of the device can be divided into different functional units or modules to complete all or part of the functions described above. The functional units and modules in the embodiments can be integrated into one processing unit, or each unit can exist physically separately, or two or more units can be integrated into one unit. The integrated unit can be implemented in hardware or as a software functional unit. Furthermore, the specific names of the functional units and modules are only for easy differentiation and are not intended to limit the scope of protection of this application. The specific working process of the units and modules in the above system can be referred to the corresponding process in the foregoing method embodiments, and will not be repeated here.

[0114] In the above embodiments, the descriptions of each embodiment have different focuses. For parts that are not described in detail or recorded in a certain embodiment, please refer to the relevant descriptions of other embodiments.

[0115] Those skilled in the art will recognize that the units and algorithm steps of the various examples described in conjunction with the embodiments disclosed in this invention can be implemented in electronic hardware, or a combination of computer software and electronic hardware. Whether these functions are implemented in hardware or software depends on the specific application and design constraints of the technical solution. Those skilled in the art can use different methods to implement the described functions for each specific application, but such implementations should not be considered beyond the scope of this invention.

[0116] In the embodiments provided by this invention, it should be understood that the disclosed devices / terminals and methods can be implemented in other ways. For example, the device / terminal embodiments described above are merely illustrative. For instance, the division of modules or units is only a logical functional division, and in actual implementation, there may be other division methods. For example, multiple units or components may be combined or integrated into another system, or some features may be ignored or not executed. Furthermore, the coupling or direct coupling or communication connection shown or discussed may be through some interfaces; the indirect coupling or communication connection between devices or units may be electrical, mechanical, or other forms.

[0117] The units described as separate components may or may not be physically separate. The components shown as units may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the units can be selected to achieve the purpose of this embodiment according to actual needs.

[0118] Furthermore, the functional units in the various embodiments of the present invention can be integrated into one processing unit, or each unit can exist physically separately, or two or more units can be integrated into one unit. The integrated unit can be implemented in hardware or as a software functional unit.

[0119] If the integrated module / unit is implemented as a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, all or part of the processes in the methods of the above embodiments can also be implemented by a computer program instructing related hardware. The computer program can be stored in a computer-readable storage medium, and when executed by a processor, it can implement the steps of the various method embodiments described above. The computer program includes computer program code, which can be in the form of source code, object code, executable files, or certain intermediate forms. The computer-readable medium can include: any entity or device capable of carrying the computer program code, recording media, USB flash drives, portable hard drives, magnetic disks, optical disks, computer memory, read-only memory (ROM), random-access memory (RAM), electrical carrier signals, telecommunication signals, and software distribution media, etc. It should be noted that the content included in the computer-readable medium can be appropriately added or removed according to the requirements of legislation and patent practice in the jurisdiction. For example, in some jurisdictions, according to legislation and patent practice, computer-readable media do not include electrical carrier signals and telecommunication signals.

[0120] This application is described with reference to flowchart illustrations and / or block diagrams of methods, apparatus, and computer program products according to embodiments of this application. It will be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, special-purpose computer, embedded processor, or other programmable data processing apparatus to produce a machine, such that the instructions, which execute via the processor of the computer or other programmable data processing apparatus, generate instructions for implementing the flowchart... Figure 1 One or more processes and / or boxes Figure 1 A device that provides the functions specified in one or more boxes.

[0121] These computer program instructions may also be stored in a computer-readable storage medium that can direct a computer or other programmable data processing device to function in a particular manner, such that the instructions stored in the computer-readable storage medium produce an article of manufacture including instruction means, which are implemented in a process Figure 1One or more processes and / or boxes Figure 1 The function specified in one or more boxes.

[0122] These computer program instructions may also be loaded onto a computer or other programmable data processing equipment to cause a series of operational steps to be performed on the computer or other programmable equipment to produce a computer-implemented process, thereby providing instructions that execute on the computer or other programmable equipment for implementing the process. Figure 1 One or more processes and / or boxes Figure 1 The steps of the function specified in one or more boxes.

[0123] The above content is only for illustrating the technical concept of the present invention and should not be construed as limiting the scope of protection of the present invention. Any modifications made to the technical solution based on the technical concept proposed in this invention shall fall within the scope of protection of the claims of this invention.

Claims

1. A gradient-based method for layer importance-aware KV cache hybrid precision quantization, characterized in that, Includes the following steps: S1. Calculate the L2 gradient norm of the Key and Value projected weights of each layer relative to the model loss function, and evaluate the difference in KV contribution between layers; S2. Based on the inter-layer KV contribution difference results obtained in step S1, allocate layer-specific bit widths to the Key and Value of different layers; S3. Based on the distribution characteristics of Key and Value, and according to the bit width allocation scheme obtained in step S2, the Key is quantized per channel and the Value is quantized per token. At the same time, the Key and Value are compressed into a low bit space to obtain the quantized KV cache data. S4. Combining the inter-layer KV contribution difference results obtained in step S1, adaptively retain the full-precision KV pairs of the most recent key tokens, and perform low-bit compression on the old KV pairs using the quantization rules of step S3 to obtain a mixed-precision KV cache set. S5. For the bit width allocation scheme in step S2, the quantization data format in step S3, and the mixed precision KV cache set in step S4, design a CUDA kernel with low bit quantization, dequantization, and matrix-vector multiplication. Optimize computational and memory access overhead through operator fusion to adapt to the inference computation requirements of the mixed precision KV cache.

2. The gradient-based layer importance-aware KV cache mixed-precision quantization method according to claim 1, characterized in that, Step S1 is as follows: S101. Load the full-precision model, sample cues from the target dataset and convert them into input tensors, shift the input tensors one bit to the left to obtain the label tensors, and determine the loss value for each input through forward propagation; S102. For each layer of the model, independently calculate the gradient of the loss relative to the Key projection weights and the Value projection weights, calculate the L2 gradient norm of the gradients respectively, and obtain the importance scores of Key and Value for each layer. S103. Take the average gradient norm across multiple sampling prompts to generate the average importance score of Key and Value for each layer, as the result of the difference in KV contribution between layers.

3. The gradient-based layer importance-aware KV cache hybrid precision quantization method according to claim 2, characterized in that, In step S102, the L2 gradient norm of the i-th layer Key projection weights is expressed as: The L2 gradient norm of the projected weights of Value is expressed as: ,in, Let the Key projection weights be for the i-th layer. Let the Value be the projection weight of the i-th layer. The model loss function; Importance score of each Key level and the importance score of each layer of Value They are respectively: in, As a prompt word, The number of prompts, Let L2 be the gradient norm of the i-th layer Key projection weight corresponding to the p-th prompt word. Let L2 be the gradient norm of the i-th layer Value projection weight corresponding to the p-th prompt word.

4. The gradient-based layer importance-aware KV cache hybrid precision quantization method according to claim 1, characterized in that, In step S2, the specific method for allocating layer-specific bit widths is as follows: The layers with the highest importance scores (top 20%-30%) are designated as high-importance layers and assigned 3-bit or 4-bit quantization precision; the remaining layers are designated as low-importance layers and assigned 2-bit quantization precision. The proportion of high-importance layers can be dynamically adjusted.

5. The gradient-based layer importance-aware KV cache hybrid precision quantization method according to claim 1, characterized in that, Step S3 is as follows: S301. Group the Key by channel dimension and reshape the tensor as follows: The grouping format is as follows: each group contains all tokens for a single channel; the Value is grouped by the token dimension, with each group containing all channels for a single token, where B is the batch size. For the number of attention heads, The length of the token sequence. For head dimension; S302. For each group, calculate the scaling factor based on the minimum / maximum value of the elements within the group. ; S303, through Quantize the original elements and convert the quantized values. Cut in Within the range; S304. Store the quantized value in an int32 using bit operations. The dequantization process is performed via... accomplish, The original KV element values ​​within the group. For the original element value The quantized value obtained after quantization It is the minimum value among all the original KV element values ​​within the group.

6. The gradient-based layer importance-aware KV cache hybrid precision quantization method according to claim 5, characterized in that, When the bit width allocation scheme is 3 bits, a specific packing strategy is adopted: The quantized elements are organized into blocks of 11 elements, each stored in a 32-bit integer. The first 10 elements are quantized to 3 bits, and the 11th element is quantized to 2 bits. The clipping range is determined by the element index within the block. Adjusted to .

7. The gradient-based layer importance-aware KV cache hybrid precision quantization method according to claim 1, characterized in that, Step S4 is as follows: S401. Based on the importance scores of the Key and Value at each level, assign the most recent key context selection ratio to each level. ; S402, Through Calculate the number of RPCs, where, It is the sum of the number of new KV states and the number of historical RPCs at the current time step; S403. Retain the full-precision KV pairs corresponding to RPC, and perform low-bit compression on the old KV pairs of non-RPC using the quantization rules of step S3 to form a mixed-precision KV cache set.

8. The gradient-based layer importance-aware KV cache hybrid precision quantization method according to claim 1, characterized in that, In step S5, optimizing computational and memory access overhead through operator fusion specifically includes: merging the quantization operation of the KV state with the operation of connecting it to the historical KV cache into a CUDA kernel; and merging the dequantization operation of the KV cache with the matrix-vector multiplication operation in the attention computation into another CUDA kernel.

9. The gradient-based layer importance-aware KV cache hybrid precision quantization method according to claim 8, characterized in that, Also includes: For different quantization bit width configurations, corresponding CUDA kernels for quantization, dequantization, and matrix-vector multiplication are designed.

10. A gradient-based layer importance-aware KV cache hybrid precision quantization system, characterized in that, include: The analysis module is used to calculate the L2 gradient norm of the Key and Value projected weights of each layer relative to the model loss function, and to evaluate the differences in KV contributions between layers. A quantization module, connected to the analysis module, is used to allocate layer-specific bit widths to the Key and Value of different layers based on the inter-layer KV contribution difference results. The compression module, connected to the quantization module, is used to compress the KV values ​​to a low bit space by using per-channel quantization for the Key and per-token quantization for the Value based on the distribution characteristics of the Key and Value values ​​and the layer-specific bit width allocation scheme, thereby obtaining quantized KV cache data. The selection module is connected to the analysis module and the compression module respectively. It is used to combine the inter-layer KV contribution difference results, adaptively retain the full-precision KV pairs of the most recent key tokens, and perform low-bit compression on the old KV pairs using the reuse quantization rules to obtain a mixed-precision KV cache set. The optimization module, connected to the quantization module, compression module, and selection module respectively, is used to design a CUDA kernel for low-bit quantization, dequantization, and matrix-vector multiplication based on the bit width allocation scheme, quantization data format, and mixed-precision KV cache set. It optimizes computational and memory access overhead through operator fusion to adapt to inference computing requirements.

Images