AlignedKV: Reducing Memory Access of KV-Cache with Precision-Aligned Quantization

Model quantization is a technique commonly used to reduce memory usage and accelerate inference speed when deploying large models. During quantization, outliers have a greater impact on the result and require higher precision than normal parameters (Dettmers et al. 2022). To deal with these outliers, some approaches (Li et al. 2023; Lee et al. 2024; Dettmers et al. 2023; Kim et al. 2023) choose to treat outlier and normal weights in a different manner. They decompose the weight matrix into a dense quantified matrix to store the regular parameters and a sparse high-precision matrix to store the outliers, then calculate with them separately. Besides, some methods blend outliers into normal parameters. Rotation (Lin et al. 2024a) treats matrices as a combination of vectors and randomly rotates them to hide outliers behind normal values. SmoothQuant (Xiao et al. 2023a) and AWQ (Lin et al. 2024b) scale outliers to fit other elements and multiply them by coefficients after dequantization.

However, the methods above can only identify important parameters through qualitative analysis and manual experiments. Only the largest elements are selected as outliers to retain full precision, lacking quantitative support. On the other hand, methods that analyze the importance of parameters quantitatively (Park et al. 2024; Kloberdanz and Le 2023) only focus on the importance of a matrix or a layer from a macroscopic perspective. To the best of the authors’ knowledge, this is the first work to systematically analyzes the importance of every parameter in a quantitative manner.

KV-Cache is an essential mechanism to save computations, but at the expense of huge storage consumption and memory access latency. It’s necessary to reduce the expense. Early works (Liu et al. 2024b; Ge et al. 2023; Zhang et al. 2024b) evites useless tokens to save memory. Some methods (Zhang et al. 2024a; Yang et al. 2024a) based on this way find that the required number of tokens decreases as the layer becomes deeper. Some quantization methods like KIVI (Liu et al. 2024c) also take different measures to quantize KV-Cache to four or even two bits. KIVI suggests that the key cache should be quantified per channel, and the value cache should be quantified per token.

As the self-attention mechanism naturally assigns importance to tokens through the attention score, mixed-precision quantization is introduced for the quantization of KV-Cache. There are some methods (Yang et al. 2024b; He et al. 2024; Yue et al. 2024; Liu et al. 2024a) quantize tokens to different precision based on their scores, while other methods (Kang et al. 2024; Dong et al. 2024) pick outliers out of matrices and save them with high precision.

Currently, most methods aim to save memory occupied by KV-Cache. The latency associated with quantization and dequantization, nevertheless, is little attended. This paper is the first work targeting to reduce the memory access latency of KV-Cache. As the memory latency is higher than the computational latency by order of magnitude (Megha et al. 2023), our work proves that KV-Cache latency is worth significant optimization effort.

In physics, the measurements of physical quantities generally cannot be exact. Therefore, in numerical representation, they are usually expressed in the form of $number\pm uncertainty$, indicating that the result falls within this range. Physics has also established a set of rules for the calculation of such numerical values, which is called uncertainty calculation (Johnson 2020).

Uncertainty calculation is closely related to LLM quantization. From such a perspective, we can represent the pre-quantified matrix W as the sum of a quantified matrix $W_{quant}$ and an uncertainty $\Delta W$. Therefore, we can express the matrix multiplication of the activation values A and the quantified weights W in the following form of uncertainty computation ($QK^{T}$ and $SV$ follow the same pattern):

We can consider the matrix multiplication above as a combination of two fundamental operations:

• •

  Multiplication of a number from A and a number from W, which corresponds to scalar multiplication in uncertainty calculation:

  $$s\times(x\pm\mathrm{\Delta}x)=sx\pm s\mathrm{\Delta}x$$

• •

  Addition of the multiplication results, which corresponds to addition in uncertainty calculation:

  $$(x\pm\mathrm{\Delta}x)+(y\pm\mathrm{\Delta}y)=(x+y)\pm(\mathrm{\Delta}x+\mathrm{\Delta}y)$$

In uncertainty addition, we can find that the uncertainty of the result equals $\Delta x+\Delta y$. On the other hand, the cost of representing $x$ is proportional to $-\log\Delta x+constant$, because we can reduce the uncertainty to $0.5\Delta x$ when we use one more bit on $x$’s representation. So we can keep the uncertainty of the result in $\Delta x+\Delta y$ at the cost of $\alpha(-\log\Delta x-\log\Delta y+constant)$, which equals $\alpha(-\log(\Delta x\Delta y)+constant)$.

We have the mean value inequality $\Delta x+\Delta y\geq 2\sqrt{\Delta x\Delta y}$, where the conditions for equality is $\Delta x=\Delta y$. So we conclude that the addition is most efficient when $\Delta x=\Delta y$. We call this the “precision alignment criterion”.

Floating-point numbers are widely used in computers. Particularly, in deep learning, the float16 (also known as half-precision floating point, FP16) is the most frequently used format. The float16 format consists of 1 sign bit, 5 exponent bits, and 10 mantissa bits, as shown in Figure 1.

The principle outlined in the previous subsection applies to floating point addition if we treat uncertainty as the precision of floating point. This is more comprehensible in the context of vertical addition, as shown in Figure 2. The figure demonstrates a case in which the precisions of two addends are different. We can find that the last few digits of the second addend do not have corresponding positions for the first addend, nor do they have corresponding positions for the results. Thus, these digits are wasted for their absence in result computation.

To avoid loading such wasted bits, we want the following two conditions to be held:

• •

  Precision Alignment on Attends: We expect the addends to have the same number of decimal places. Otherwise, some digits have no digits to add.

• •

  Precision Alignment on Results: We expect each addend to have identical numbers of decimal places with the final result. Otherwise, some digits cannot be stored in the result.

These two expectations are equivalent in the vast majority of cases. They simply come from different perspectives.

As mentioned earlier, matrix multiplication consists of two fundamental operators—scalar multiplication and addition. We can understand matrix multiplication from a 3D perspective, as Basil Hosmer described (Hosmer 2023). Matrix multiplication corresponds to a large cube, with the two matrices involved in the multiplication and the result corresponding to the three faces of the cube, as Figure 3 shows. Based on this correspondence, the scalar multiplication operation represents the smaller cubes, while summing up the results of the scalar multiplications corresponds to the addition between the smaller cubes in the horizontal dimension.

The previous subsection suggests that addition with consistent precision among the addends is the most efficient. This principle is also suitable in additions inside General Matrix Multiply (GEMM). We can treat the addition in GEMM from two perspectives:

• •

  Precision Alignment on Attends: The precision of the intermediate results in GEMM should be consistent.

• •

  Precision Alignment on Results: The precision of each intermediate result in GEMM should be the same as that of the final result.

Currently, the vast majority of quantization methods are static, with only a few works employing dynamic quantization based on data (Wang, Zhang, and Han 2021) as most quantization techniques aim to save storage space, and hence a static quantization suffices.

However, in the context of KV-cache, dynamic quantization offers compelling advantages. Firstly, dynamic quantization can reduce memory access latency without adding additional delays. It’s important because memory access latency is the bottleneck for the inference speed in this stage. Secondly, static quantization introduces a significant problem for the usage of mixed-precision quantization in KV-cache, as it requires determining the importance of elements at the time of quantization. Although there are numerous methods available to predict the importance of elements to address this issue (i.e., the persistence of importance hypothesis)(Ge et al. 2023; Liu et al. 2024b), these predictions are always subject to counterexamples (Dong et al. 2024).

We propose a method to dynamically quantize KV-cache, which quantizes data when it is used. When storing the KV-cache, we maintain additional data structures to store the approximate magnitudes of the KV-cache elements, which is one of the preconditions for precision alignment. Afterwards, during the computation of $QK^{T}$ and $SV$, we use the precision alignment criterion to determine the required precision for each element in K and V.

We then retrieve only the necessary number of bits from memory on an as-needed basis, as illustrated in Figure 6. According to the results of the precision alignment principle, we read all 16 bits for important parameters, the first 12 bits for less important parameters, and the first 8 bits for particularly unimportant parameters. The remaining bits are filled with 100… to minimize the precision loss to the greatest possible extent.

At the same time, as the GPU hardware can only read memory at a minimum unit of several bytes, we rearrange the storage structure to ensure continuity for hardware access, as shown in Figure 8.

AlignedKV follows the process outlined below, applicable to both K-Cache and V-Cache:

1. 1.

   Obtain the approximate magnitude information of the values in KV-Cache by an additional data structure, preparing the data for precision alignment. We only need the order of magnitude information.

2. 2.

   Apply our framework to determine the required precision for loading each data element.

3. 3.

   Perform truncation on the data elements using dynamic quantization methods while loading data from GPU memory to L1 cache.

Among these, the latter two have already been addressed in the previous sections, leaving only the first step, which requires specific strategies based on the characteristics of the data.

Based on observations of the K-Cache, we found that the data distribution exhibits a columnar characteristic, where the data within a single column is relatively similar and the data between different columns varies significantly, as shown in Figure 7. In other words, we can use a single value from a column (specifically, the maximum value) to represent the order of magnitude for that column.

Thus, we maintain the information of the maximum value for each column of the K-Cache, which is recorded as $K_{colmax}$. When computing $QK^{T}$, we first use the value of $Q$ along with the values of $K_{colmax}$ to estimate the order of magnitude of the intermediate results. Then, we can determine the aligned precision of the intermediate results by Rule 1 (Addend Precision Alignment), which allows us to infer the required bit width for each element in the K-Cache. All of this is derived quantitatively, rather than merely through qualitative analysis.

Then, we load the required bits of each parameter from memory using the proposed dynamic quantization methods. As a result, we can avoid loading a large number of redundant bits from GPU memory, which significantly reduces the memory access latency for this part.

Similar to the quantization of K-Cache, the quantization of V-Cache also requires obtaining approximate order of magnitude information for the addends or results in order to apply precision alignment criterion. However, unlike K-Cache, V-Cache does not exhibit a columnar distribution characteristic, which prevents us from applying the quantization method used for K-Cache. Fortunately, we observed that the magnitude differences among the elements in Score matrix (recorded as S) are quite significant, as shown in Figure 7, which means that we only need to compute the product of a subset of the larger elements in S with the corresponding V to obtain an approximate estimation of final results.

We first select the top k largest values from the score matrix S using an approximate top-k algorithm (we don’t use the exact top-k algorithm for speed) and multiply them by their corresponding V values to obtain an estimate of the SV result, referred to as $O_{est}$. Subsequently, we can determine the precision alignment of the intermediate results by Precision Alignment on Result using $O_{est}$. Similar to the quantization of K-Cache, the alignment of the intermediate results allows us to infer the required bit width for each element in V-Cache. We can then reduce memory access latency by reading only the necessary bits for each element.

We performed comprehensive experiments to validate the effectiveness of the proposed method. We conducted our experiments on the Llama-2-7b model with our implementation of KV-cache coded in CUDA. The experiments were conducted on a Nvidia V100 GPU. It must be noted that the proposed method is friendly to efficient hardware implementation and the dedicated hardware can be significantly more efficient.

In this experiment, we apply AlignedKV on the actual KV-Cache data generated by Llama-2-7B. We statistically analyze the alignment of actual accuracy and the number of bits required for each element in the KV-Cache. The results are illustrated in Figure 9. They clearly demonstrate that after using AlignedKV, the average bit width for each element decreases from 16 to approximately 12, resulting in a significant reduction of the memory access latency during the KV-Cache access. It drops with the increase of context length, indicating that our method has significant potential for longer context.

In quantization, it is essential to consider not only the compression rate but also the accuracy. In the previous experiment, we also analyze the impact of applying AlignedKV on the accuracy of the results. We record the changes in $QK^{T}$ and $SV$ GEMM after using AlignedKV and represent these changes through the distribution of relative errors.

We use the following formula to calculate the relative error for each value in the results:

We calculate the relative error for each value in the result matrix and observe their distribution, as shown in Table 1. Notably, although AlignedKV reduces the bit widths of KV-Cache by 4 bits, most of the results remain unchanged. In contrast, uniformly removing 3 bits from each element in KV-Cache yields significantly poorer results, despite this approach achieving a lower compression rate than AlignedKV. This indicates that by quantitatively analyzing the bit-width requirements for each element and letting them achieve a state of “alignment”, we can take both a high compression rate and high accuracy.

We also evaluated the actual runtime. We run the inference of a Llama-2-7B model with 128-token input and up to 8192-token completion (disregarding the model’s maximum context length of 4096 because we only care about the time cost). To highlight the effectiveness of AlignedKV, we only focus on the computations involved in the KV-Cache component. The computations are as follows:

When the context length reaches hundreds of thousands or even millions, the memory access latency of the KV-Cache will become the bottleneck for the entire model’s running speed. At this point, the benefits of AlignedKV will be substantial.

We compare AlignedKV with the native implementation in Torch and two other quantization methods—KIVI (Liu et al. 2024c) and GEAR (Kang et al. 2024) (these quantization methods are set to 10 bits for fairness). We also set a control group using AlignedKV’s code to evaluate the memory access reduction AlignedKV brings, ignoring the additional overhead associated with CUDA program execution. The only difference between the control group and AlignedKV is that the latter always loads elements as 16-bit values.

The runtime results are displayed in Figure 10. With the increasing position of tokens, the length of the KV-Cache also grows, at which point the advantages of AlignedKV in reducing memory access latency become apparent. When the generating length reaches around 7000 to 8000, AlignedKV demonstrated the shortest execution time among all experimental groups, surpassing the performance of the native implementation in Torch. The CUDA implementation of Torch is highly optimized. We expect a higher level of performance improvement of our implementation after optimization.

We are willing to present the end-to-end accuracy performance of AlignedKV, as AlignedKV doesn’t cause any loss of precision. We benchmark it on CoQA(Reddy, Chen, and Manning 2019), TruthfulQA(Lin, Hilton, and Evans 2021), and GSM8K(Cobbe et al. 2021) tasks using default settings, as KIVI does. We utilize an open-source evaluation repository (Gao et al. 2024) to conduct this experiment. The results are shown in Table 2.

Through experiments, we find the results of AlignedKV are very close to the origin model, which means AlignedKV has no impact on the result of LLMs. The results are as we expected since the errors in GEMM results will be diluted on a higher level, and the results of GEMM by AlignedKV are already accurate.

To further demonstrate that our method incurs exactly no loss in end-to-end accuracy, we used the model to generate sentences by greedy search. We found that the results generated by AlignedKV are completely consistent with those produced by the original model. In contrast, although the sentences generated by KIVI are coherent, the initial words differ from those generated by the original model. A brief example is shown in Table 3.