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.

翻译：巨大的内存消耗一直是阻碍高吞吐量大型语言模型在实际应用中部署的主要瓶颈。除了庞大的参数数量外，Transformer架构中注意力机制的关键值（KV）缓存也消耗了大量内存，尤其是在深度语言模型中层数较多的情况下。本文提出了一种新颖方法，仅计算并缓存少数层的KVs，从而显著节省内存消耗并提升推理吞吐量。我们在大型语言模型上的实验表明，该方法相比标准Transformer实现了高达26倍的吞吐量提升，并在语言建模及下游任务中保持了具有竞争力的性能。此外，本方法与现有Transformer内存节省技术正交，因此可轻松将其集成到我们的模型中，进一步提升推理效率。我们的代码开源在https://github.com/whyNLP/LCKV。