Layer-Condensed KV Cache for Efficient Inference of Large Language Models (2024)

Abstract: Huge memory consumption has been a major bottleneck for deploying high-throughput large language models in real-world applications. In addition to the large number of parameters, the key-value (KV) cache for the attention mechanism in the transformer architecture consumes a significant amount of memory, especially when the number of layers is large for deep language models. In this paper, we propose a novel method that only computes and caches the KVs of a small number of layers, thus significantly saving memory consumption and improving inference throughput. Our experiments on large language models show that our method achieves up to 26 \times higher throughput than standard transformers and competitive performance in language modeling and downstream tasks. In addition, our method is orthogonal to existing transformer memory-saving techniques, so it is straightforward to integrate them with our model, achieving further improvement in inference efficiency. Our code is available at https://github.com/whyNLP/LCKV.

## introduction high throughput and low latency are essential for deploying large language models (llms) in realworld applications @xcite . however, the huge memory consumption of llms has been a major bottleneck, preventing a large batch size and high throughput generation. among the memory-consuming components, the key-value (kv) cache is one of the most significant parts @xcite that takes over 30% of the gpu memory during deployment @xcite . the kv cache is used to store the keys and values in each transformer layer during generation to avoid re-computation. its memory consumption is proportional to both the sequence length and the number of layers. there have been substantial works on reducing the memory consumption of the kv cache in llms. most of them focus on compressing the kv cache by reducing the length of the cached kv sequence. for example, @xcite ; @xcite ; @xcite compress the prompts to save the memory consumption. @xcite incrementally compress a specified span of tokens into compact ones to reduce the kv cache length. @xcite ; @xcite propose to store only the kvs of initial and recent tokens in the kv cache. @xcite propose a dynamic kv cache eviction policy to only keep a small portion of the kv cache in memory. in this paper, we propose to reduce the memory consumption of the kv cache from a novel perspective that is orthogonal to previous efforts: reducing the number of layers. specifically, we present a new variant of transformer decoders in which queries of all layers are paired with keys and values of just the top layer, as illustrated in figure 1 . in this way, we only need to cache the keys and values of one layer (vs. tens of layers in a typical llm), significantly saving memory consumption while introducing no computation overheads dur-ing inference. in fact, since we only need the keys and values of the top layer, we can omit the keyvalue computation and discard the key-value parameters for all the other layers, further improving the throughput and reducing the memory consumption as well as the mode size. we draw our inspiration from the interpretation of the stacking layer structure of a transformer as an iterative process of improving token representation @xcite . in this interpretation, the representation at the top layer is the most informative, so it makes sense to attend only to the top layer. we also note the similarity of our idea to the cross-attention mechanism in a standard transformer encoder-decoder, in which all the decoder layers attend to the top encoder layer. however, applying this idea to a decoder has never been attempted before as far as we know. although our model presented above achieves very efficient inference, its performance in language modeling and downstream tasks degrades in comparison with standard transformers. therefore, we further propose to retain standard attention for a small number of layers in our model, which slightly diminishes our saving of the kv cache memory consumption but leads to almost no performance degradation. another challenge faced by our model is training. when training a standard transformer decoder, the computation of all the tokens can be fully parallelized. in our model, however, the computation at each token depends on the top layer of the previous token, creating sequential dependencies that spoil parallel training. we address the challenge by deriving an approximate training method that supports parallel training. our experiments on llama @xcite ) show that our model achieves up to 32× larger batch sizes and up to 26× higher throughput than standard transformers for llms of 1b-30b parameters; at the same time, our model has competitive performance to standard transformers in language modeling and downstream tasks. we further empirically demonstrate that it is straightforward to integrate our model with other memory-saving techniques like streamingllm @xcite , achieving further improvements in inference efficiency. we summarize our contributions as follows: 1) we propose a new variant of transformer decoders that reduces the memory consumption of the kv cache by dramatically reducing the number of cached layers; 2) we make parallel training of our model feasible by designing a novel approximate training method; 3) we conduct extensive experiments to verify and analyze the effectiveness and efficiency of our method. ## experiments we empirically verify the effectiveness of our method on the llama model @xcite . we show that our method achieves significant memory reduction and throughput improvement as well as competitive performance in language modeling and downstream tasks compared with standard transformers. we also show that our method could effectively integrate with other memory-saving techniques. ## analyses in this section, we empirically analyze design choices in our method. for experiment details, please refer to appendix b. ## related work extensive research has been done on reducing kv cache for efficient inference of llms. with the exception of vllm @xcite , which proposes paged attention to reduce memory fragmentation of the kv cache from a system perspective, most recent works focus on compressing the kv cache by reducing the length of the cached kv sequence. @xcite accelerate model 0 5 10 15 20 25 30 # of iterations 0.0 0.5 1.0 1.5 2.0 2.5 mse tinyllama (1.1b) random trained figure 10: mse of the kv before and after the ith iteration. the models are tested with 2048 tokens. inference by compressing document-level prompts into short prompts. li et al. (2023) remove the redundancy in the input context. mu et al. (2023) train gist tokens to replace the reusable system prompts. ren et al. (2023) incrementally compress a specified span of tokens into compact ones to reduce the kv cache length. liu et al. (2023) find that only pivotal tokens have a significant influence at a future step, so pruning unimportant tokens does not affect the performance. ge et al. (2023) argue that the attention heads could be classified into different types and propose to apply different kv pruning strategies to different types of attention heads. xiao et al. (2024); han et al. (2023) find that only initial tokens and recent tokens are crucial and propose to store only the kvs of these tokens to enable infinite-length context for llms. zhang et al. (2023) propose a kv cache eviction policy based on the summation of attention scores to only keep a small portion of the kv cache in memory. unlike these previous methods, our method reduces the memory consumption of the kv cache by reducing the number of layers, which is orthogonal to these methods and can potentially be combined with these methods to further reduce the kv cache and improve inference efficiency. feedback transformers @xcite aggregate hidden representations from all layers and use them as token memory. this is followed by kv projections to obtain the key-value pairs for all layers. their experimental results show improved performance with this strategy, even when using only the hidden representation from the top layer. however, their sequential training process is time-costly and not practical for large models. our method supports parallel training, which is more efficient and scalable. ## conclusion in this paper, we propose a novel method to reduce the memory consumption and improve the throughput of llms by reducing the number of layers whose keys and values need to be computed and cached. we empirically show that our method achieves significant memory reduction and throughput improvement with negligible performance degradation. we also show that our method could effectively integrate with other memory-saving techniques like streamingllm. future work includes designing more efficient training approaches, developing large-batch-friendly kernels, and verifying our method on larger and more complex llms. we hope that our work could provide a new perspective for improving inference efficiency of llms and inspire more research in this direction.