P yramid I nfer: Pyramid KV Cache Compression for High-throughput LLM Inference (2024)

Abstract: Large Language Models (LLMs) have shown remarkable comprehension abilities but face challenges in GPU memory usage during inference, hindering their scalability for real-time applications like chatbots. To accelerate inference, we store computed keys and values (KV cache) in the GPU memory. Existing methods study the KV cache compression to reduce memory by pruning the pre-computed KV cache. However, they neglect the inter-layer dependency between layers and huge memory consumption in pre-computation. To explore these deficiencies, we find that the number of crucial keys and values that influence future generations decreases layer by layer and we can extract them by the consistency in attention weights. Based on the findings, we propose PyramidInfer, a method that compresses the KV cache by layer-wise retaining crucial context. PyramidInfer saves significant memory by computing fewer keys and values without sacrificing performance. Experimental results show PyramidInfer improves 2.2x throughput compared to Accelerate with over 54% GPU memory reduction in KV cache.

## introduction large language models (llms) @xcite @xcite like gpt4 have demonstrated the unprecedented ability of remarkable comprehension in human languages. however, these large models meet up with a substantial challenge of immense gpu memory usage in the inference, due to the model and computational complexity. this hinders deploying llms at scale to meet the thousands of demands for chatting with chatbots. different from training, models in the inference do not need to record the optimizer states, activations, or gradients. as llms are mostly transformer-based auto-regressive models, the gpu memory usage mainly consists of two parts: model parameters and kv cache. kv cache presents the keys and values previously computed in the attention. we store the kv cache in the gpu memory and reuse it in future generations to avoid re-computation. the kv cache mechanism has been widely used to improve the inference speed @xcite . however, the kv cache consumes huge gpu memory, especially for llms. for example, in figure 1 , for a model with 7 billion parameters, the parameters only consume 14 gb of memory but the kv cache requires around 72 gb. the kv cache has the potential to consume memory several times the size of the model. it demonstrates a great challenge that the throughput of llm inference is constrained by how much data (kv cache) we can put in the gpu besides the model. we break down llm inference into two phases: prefill phase and generation phase @xcite . in the prefill phase, the prompt is computed in parallel to generate the first token, and the initial kv cache is pre-filled. in the generation phase, the model decodes the next token one by one and appends the keys and values of the newly decoded token to the old kv cache. recent studies @xcite @xcite compress the kv cache to reduce gpu memory usage. however, as shown in figure 2 , they all only reduce the kv cache that has been already computed rather than reducing the kv cache to be computed. they have to prefill the initial kv cache before they can start to compress, which neglects the great gpu memory consumption of computing the initial kv cache, especially for longer prompts and larger models. if the model can not process the prompt in the prefill phase, these methods are no longer applicable as their compression starts in the generation phase. in this paper, we focus on how to further compress the kv cache in the prefill phase besides the generation phase. we give out our findings and then propose our method pyramidinfer inspired by these findings. during the training, all input tokens predict the tokens next to themselves in an one-to-one teacher-forcing way @xcite . during the inference, the tokens except for the last token no longer need to predict the next tokens but they still record this redundant information in keys and values. we call this inference context redundancy (icr) hypothesis. it inspires us to compress the kv cache by only computing the keys and values that record the context information. another challenge arises as the initial kv cache is reused multiple times for generating future tokens, necessitating careful retention of context information during compression. inspired by the work @xcite , we further explore what parts of the kv cache are always crucial for future generations. we observe that queries of recent tokens closer to the last token are more consistent in attending to the same context keys and values, denoted as the pivotal context (pvc). we call this phenomenon as recent attention consistency (rac). the consistency of attention weights in recent tokens indicates that we can leverage it as the oracle to select the crucial kv cache for future generations in advance. based on our observations, we propose the pyramidinfer, an effective method of reducing the kv cache both in the prefill and generation phase by layer-wise selecting the pvcs. in pyramidinfer, the pvcs are gradually reduced as the layers get deeper where the kv cache is like a pyramid. we showcase the capability of pyramidinfer on a wide range of tasks using opencompass @xcite on models of different types and sizes. the results show that pyramidinfer has higher throughput than the full cache method accelerate and deepspeed by 2.2x and 1.4x, kv cache compression method h 2 o by 2.4x with over 54% less gpu memory in kv cache. ## related work due to the increasing demands for chatting with chatbots, efficient strategies are required to process thousands of queries to maximize the throughput. the fundamental way to improve the throughput is to put more data (larger batch) into the gpu memory to utilize the gpu parallelism better. inference parallelism one way is to enlarge the gpu memory. we can borrow the techniques used in training to accelerate the inference, e.g., pipeline parallelism @xcite , kv cache offload @xcite , etc. these methods leverage multiple gpus or even ram to make up bigger space for input data. ## observation and insight we verify the hypotheses of inference context redundancy and recent attention consistency, which inspire us to design the method pyramidinfer. ## layer-wise pvc selection based on the observations, we design the pyramid-infer, a method to highly increase the inference throughput by layer-wise selecting the pvcs to compress the kv cache for each layer. ## conclusion we alleviate the difficulty of deploying llms at scale by introducing pyramidinfer, a novel method that efficiently compresses the kv cache during both prefill and generation phases. inspired by icr and rac, pyramidinfer significantly reduces gpu memory usage without compromising model performance. experimental results present pyramid-infer is a promising solution for optimizing llm deployment in resource-constrained environments.