Lossless KV Cache Compression to 2%

Recently, transformer-based models have garnered significant attention due to their powerful performance in various fields, including NLP and multimodal tasks. Large language models such as GPT-4 OpenAI et al. (2023) and LLaMA3 AI@Meta (2024) have achieved remarkable results across a range of natural language processing tasks.

However, the traditional decoder-only architecture incurs substantial training and online inference costs. To address this issue, Palm Chowdhery et al. (2023) and StarCode Li et al. (2023b) utilize MQA Shazeer (2019), where all attention heads share a group of KV. This approach aims to reduce the computational overhead associated with the attention mechanism. LLaMA2 Touvron et al. (2023) employs a different strategy within its attention blocks by grouping keys and values (GQA) Ainslie et al. (2023) to minimize the KV cache during inference. Concurrently, it increases the feed-forward network hidden size to maintain model performance. FlashAttention Dao (2023) optimizes training and inference efficiency from the perspective of GPU I/O operations. This method specifically targets the bottlenecks in data movement and computation within the GPU architecture. Mixture-of-Experts models (Jiang et al., 2024; Muennighoff et al., 2024; Wang et al., 2024a) utilize sparse architectures to achieve improved performance.

Previous research has utilized techniques such as reducing the dimensions of attention heads and employing low-rank matrix approximations to enhance model efficiency. Some recent work reduces the KV cache through the layer dimension(Wu and Tu, 2024; Zuhri et al., 2024; He and Wu, 2024; Liu et al., ). Notably, Yu effectively compressed models by minimizing errors in MHA-to-GQA transitions Yu et al. (2024), while maintaining compatibility with RoPE Su et al. (2023). Some research efforts concentrate on reducing KV cache requirements by decreasing the sequence length, using methods such as token dropping Zhang et al. (2024b); Li et al. (2024); Xiao et al. (2023); Shi et al. (2024) and prompt compression Jiang et al. (2023); Chuang et al. (2024). However, these methods often lead to a notable decline in performance.

Recently, many research efforts have employed quantization compression to reduce KV cache, thereby improving inference efficiency (Banner et al., 2018; Wu et al., 2020; Yang et al., 2024; He et al., 2024; Liu et al., 2024a). Gear Kang et al. (2024) utilizes low-rank matrices and low-precision quantization to achieve near-lossless 4-bit KV cache compression. Coupled Quantization Zhang et al. (2024a) encodes multiple KV channels into a low-bit code. KIVI Liu et al. (2024b) performs quantization on keys and values along the channel and token dimensions, respectively. These works have primarily focused on inference. In contrast, we propose CLLA-quant method to quantize the KV cache in 4-bit integers during training and deployment.

Currently, transformer model is the most prevalent LLM architecture. It is composed of many blocks stacked together. Each block includes an attention layer and a MLP layer. Multi-head Attention (MHA) Vaswani et al. (2017) is a typical design for the attention layer. Given the sequence hidden states $\mathbf{H}=[h_{1},h_{2},...,h_{t}]$, the computations of MHA for each head are performed as follows:

where $\mathbf{W}_{i}^{Q}$, $\mathbf{W}_{i}^{K}$, and $\mathbf{W}_{i}^{V}$ are weight matrices in $\mathbb{R}^{D_{h}\times D_{h}}$ for $\text{head}_{i}$. $N$ denotes the number of heads. $\text{Att}(\cdot)$ represents the calculation of scaled dot-product attention Vaswani et al. (2017).

To increase the speed of autoregressive decoding, it is necessary to store the key and value vectors ($K$ and $V$) of previous input hidden states. In Transformer, the memory requirement for KV cache is influenced by multiple parameters: batch size ($B$), number of layers ($L$), sequence length ($S$), number of attention heads ($N$), and the dimension of each head ($D_{h}$). This dependency can be formulated as:

Classical methods such as GQA Ainslie et al. (2023), MLA DeepSeek-AI (2024), and CLA Brandon et al. (2024) usually compress KV cache from $N$, $D_{h}$, and $L$ aspects as shown in Table 1. In this work, we strive to jointly consider all these aspects for performance-lossless KV cache compression.

CLLA involves compressing KV cache into low-rank latent states and subsequently allowing multiple layers to share these states. The key research point is how to effectively share compressed latent states across layers.

To further reduce memory overhead associated with KV cache, we implement low-bit quantization on latent states by converting them from 16-bit floating-point values into lower-bit integers using symmetric quantization Nagel et al. (2021). The quantization function can be expressed as follows:

Here $\mathbf{C}$ represents the latent states, and $Q\left(\mathbf{C}\right)$ represents their quantized version stored as $b$-bit integers. $\lfloor\cdot\rceil$ denotes the rounding operator, while $S$ serves as the scaling factor. The latent states are scaled to the restricted range $\left[I_{min},I_{max}\right]$, which is set to $\left[-\left(2^{b-1}-1\right),2^{b}-1\right]$. The quantized latent states are stored in cache and dequantized only when needed during computations. The dequantization function is defined as follows,

The recovered latent states $\mathbf{\hat{C}}$ are then utilized in computations. We adopt 4-bit for CLLA-quant.

The training dataset was collected in-house and consists of a cleaned combination of Chinese and English datasets. We trained all models on 300 billion tokens. For evaluation, we aim to achieve robust conclusions across diverse domains in both English and Chinese. We conducted comprehensive evaluations involving 7 English tasks, including MMLU, Hellaswag, ARC-C, Winogrande, BoolQ, PIQA, and SIQA. For Chinese tasks, we evaluated models on 8 tasks: CMMLU, CEval, CLUEWSC2020, Thuc-news, OCNLI, C3, FewCLUE-CHID, and CMNLI. More details are noted in Appendix B.

In order to ensure equitable comparison, the intermediate dimensionality of the feed-forward neural network (FFN) was adjusted to maintain comparable activation parameters across all models. The Appendix elaborates on the implementation specifics of each model. The training process was executed utilizing 64 NVIDIA H800 GPUs, with the same hardware resources employed for the inference evaluations. Complete details of the training configurations are provided in the Appendix.

For every competitor, we utilize the same 1.44 billion activation parameters with a Mixture of Experts (MoE) architecture Fedus et al. (2022); Lepikhin et al. (2020). Specifically, we employ a model with a 1Bx16 MoE configuration, comprising nearly 11 billion parameters in total, to ensure fair comparisons.The following methods were evaluated:

MHA: The foundational transformer architecture Vaswani et al. (2017); Radford et al. (2019) resembles that of the Llama2 model Touvron et al. (2023). It features SwiGLU Shazeer (2020) as the activation function for the feed-forward layer and utilizes RoPE Su et al. (2023) for the embedding of relative positions. All subsequent models are variations of this foundational model.

GQA: GQA Ainslie et al. (2023) groups query heads together, allowing each group to share keys and values. We set the number of KV heads to 8 out of a total of 16 query heads to reduce KV cache by half.

MLA: In MLA DeepSeek-AI (2024), the KV cache is mapped into a low-rank space. The latent dimension for this low-rank space is set at 512.

CLLA: Our method combines the settings from MLA and CLA Brandon et al. (2024) while adhering to the aforementioned sharing configurations.

CLLA-quant: We applied int4 scalar quantization to compress latent states with a quantization group size of 32 on the original CLLA.

In each model, the hidden size is configured to 1536, with 32 layers and 16 attention heads. We incorporated 17 experts in the MoE layer, consisting of one shared expert alongside 16 experts, from which the top-1 is selected for routing. To ensure the alignment of activation parameters and total parameters across various methods, we modified the expansion ratio of the feed-forward network (FFN) as needed for fair comparisons. Further details regarding training and model hyperparameters can be found in Appendix A.

After pre-training on 300 billion tokens, we evaluated all models across 15 benchmarks that encompass a wide range of tasks. All models are compared in the same tokenizer and data. On English benchmarks shown in Table 2, MLA are comparable with GQA baseline. The proposed models CLLA and CLLA-quant achieve better results with higher compress ratio. We can achieve similar conclusion on Chinese benchmarks presented in Table 3. Additionally, theoretical KV cache memory bytes usages for each model are shown in Table 4. From these tables, several interesting observations emerge regarding the effectiveness of our methods in achieving high KV cache compression ratios without performance degradation:

First, CLLA notably surpassed MHA in performance while also achieving a significant reduction in memory usage of nearly 95%. For English tasks, CLLA’s performance was about 1.24 point higher than that of MHA, and it showed similar enhancements for Chinese tasks. This improvement is possibly attributed to our model design. Under the same amount of activated parameters, we were able to increase the FFN hidden size by utilizing parameters saved from the CLLA approach, which likely contributed to the improved outcomes.

Second, CLLA-quant exhibited results comparable to CLLA while further achieving a fourfold reduction in memory usage. These results verify the impressive conclusion that our CLLA-quant method compresses KV cache memory to 2% without performance loss generally on 15 benchmarks compared to MHA. We hypothesize that the minimal impact of quantization on performance is due to its application during pretraining, allowing the entire model to adapt to the quantized KV cache. The quantization likely enforces a form of robustness in the representations stored in the KV cache, ensuring that minor precision losses do not adversely affect the model’s ability.

To gain a deeper understanding of the strategy for combining CLA and MLA, we conduct a thorough analysis of both intra-layer and inter-layer sharing methods. The superiority between inter-layer sharing methods, such as CLA, and intra-layer sharing strategies, such as MQA and GQA, has been a subject of contradictory conclusions in current studies. In Brandon et al. (2024), the authors claim that with proper design, the combination of intra-layer and inter-layer strategies can make inter-layer sharing better than intra-layer sharing in some situations. However, a later study Zuhri et al. (2024) found that intra-layer strategies are generally better than inter-layer strategies. To explore these design choices, a series of experiments were conducted. The experimental models are in the same setting with our main experiment, except we only train 50B tokens for a fast comparison on validation perplexity. The results are presented in Table 5:

(1) Intra-layer strategies consistently achieve lower validation perplexity compared to inter-layer strategies when using the same memory budget. Specifically, comparing GQA-Group8 with CLA-Share2, GQA-Group4 with CLA-Share4, and GQA-Group2 with CLA-Share8, intra-layer strategies consistently outperform inter-layer strategies in terms of validation perplexity. Furthermore, even with equivalent memory reductions, the performance loss associated with the CLA strategy is significantly greater than that of the GQA strategy. This conclusion is supported by comparisons between CLA configurations with varying share factors and GQA configurations with different group sizes.

(2) Neither the combination of GQA and CLA strategies nor adjustments to the head dimensions improved the performance. Specifically, when we combined both strategies, we observed that GQA-Group2+CLA-Share2 performed worse than MQA, while GQA-Group4+CLA-Share2 was inferior to GQA-Group2. Additionally, maintaining the total hidden dimension while adjusting head dimensions failed to surpass simple intra-layer sharing strategies. This is evident from comparisons between MQA-Dim192-CLA-Share2 and MQA as well as MQA-Dim384-CLA-Share2 and GQA-Group2.

Based on these observations, we can outline a design strategy for utilizing intra-layer and inter-layer sharing to conserve KV-cache memory. First, prioritize intra-layer strategies if the KV-cache memory budget can be met solely through intra-layer approaches. Second, resort to inter-layer strategies only after significantly compressing intra-layer memory strategies, such as MQA and MLA, when further reduction is necessary.

In addition to the aforementioned ablation studies, inspired by MLA, we explored the possibility of sharing input hidden states for attention blocks instead of KV heads. This approach necessitates caching input hidden states for inference while computing KV heads in real-time. We can view the MLA strategy as a LoRA variant of this concept. When applying CLA within this framework, we must determine whether to share KV projection weights as in Figure 3. Our findings indicate that allowing each layer its own KV projection weights yields slightly better performance than sharing KV projections, with validation perplexities of 8.97 versus 8.99. Although this decision is critical for CLLA, the difference is not substantial in this context.

Upon decomposing the components eligible for cross-layer sharing within the MLA strategy, we identified three elements: KV latent vectors, K rope vectors, and KV projection matrices. These three configurations are illustrated in Figure 2. Our final design decision favors the first configuration due to its superior performance. The second configuration offers additional KV-cache savings but only marginally so because of the relatively small size of K rope vectors compared to KV latent vector sizes. The third configuration may reduce computation and model parameters.

We evaluated these three models based on downstream tasks under identical experimental settings as those used in the main experiments. The results are presented in Table 6, which displays average scores for both English and Chinese tasks. Detailed results for all English and Chinese task can be found in the Appendix C.

Our analysis reveals that CLLA with exclusive sharing of KV latent vectors achieved the best performance. The strategy of sharing K rope vectors significantly detracted from performance with minimal benefits. While the KV projection sharing strategy yielded better results for English tasks, it was lower overall. Nevertheless, this approach can provide moderate improvements during training and inference by eliminating half of the KV projection calculations and matrices required. Given our primary focus on downstream performance, we opted for the first configuration as our CLLA strategy. Another reason for selecting this configuration is its favorable interaction with quantization techniques.

Based on our previous findings, we selected CLLA and CLLA+share-kvproj for latent quantization training and compared their downstream performance as shown in Table 7. All experimental settings were consistent with those used in the main experiments. Our results indicate that the strategy without sharing KV projection matrices achieved significantly higher performance when combined with quantization techniques. The observed differences may be attributed to the fact that having multiple KV projections can better compensate for quantization loss compared to relying on a shared KV projection.