A Systematic Study of Cross-Layer KV Sharing for Efficient LLM Inference

📄 arXiv: 2410.14442v2 📥 PDF

作者: You Wu, Haoyi Wu, Kewei Tu

分类: cs.CL

发布日期: 2024-10-18 (更新: 2025-02-05)

备注: Accepted to NAACL2025 main conference

💡 一句话要点

系统性研究跨层KV共享，提升LLM推理效率

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 LLM推理 KV缓存 跨层共享 模型优化

📋 核心要点

  1. 大型语言模型推理效率受限于KV缓存，跨层共享KV是有效方法，但缺乏系统性研究。
  2. 论文提出统一框架，涵盖多种跨层KV共享策略及其变体，便于系统性分析。
  3. 实验表明，适度减少KV缓存可提升吞吐量，过度减少需权衡训练成本与性能。

📝 摘要（中文）

本文系统性地研究了跨层键值(KV)共享技术，该技术旨在提高大型语言模型(LLM)的推理效率。我们提出了一个统一的框架，涵盖了多种现有的跨层KV共享方法及其变体。通过对框架中所有配置进行全面的实验，评估了它们的生成吞吐量以及在语言建模和下游任务中的性能。实验结果表明，在KV缓存大小减少2倍时，大多数配置都能在保持竞争力的性能的同时，实现比标准Transformer更高的吞吐量。然而，当进一步减小KV缓存大小时，将所有层的查询与上层的KV配对表现更好，但会增加额外的训练成本和预填充延迟。我们希望这项工作能够帮助用户更明智地选择跨层KV共享方法，并促进未来对高效LLM推理的研究。

🔬 方法详解

问题定义：大型语言模型（LLM）的推理效率受到KV缓存大小的限制。传统的Transformer架构中，每一层都需要维护自己的KV缓存，这导致了巨大的内存占用和计算开销，尤其是在长序列生成时。现有的跨层KV共享方法旨在通过在不同层之间共享KV缓存来减少内存占用，但缺乏系统性的研究和比较，难以指导实际应用。

核心思路：本文的核心思路是构建一个统一的框架，将多种现有的跨层KV共享方法及其变体纳入其中，从而能够系统地研究不同策略的优缺点。通过对框架中所有配置进行全面的实验，分析它们在生成吞吐量、语言建模性能和下游任务性能方面的表现，从而为用户提供更明智的选择。

技术框架：该研究提出了一个统一的跨层KV共享框架，该框架允许配置不同的KV共享策略。具体来说，该框架允许配置哪些层的查询与哪些层的KV进行配对。例如，可以将所有层的查询与上层的KV进行配对，也可以只将相邻层的查询与KV进行配对。此外，该框架还允许配置KV缓存的大小。

关键创新：本文的关键创新在于提出了一个统一的框架，能够系统地研究不同的跨层KV共享策略。通过对框架中所有配置进行全面的实验，揭示了不同策略的优缺点，并为用户提供了更明智的选择。此外，该研究还发现，在KV缓存大小减少2倍时，大多数配置都能在保持竞争力的性能的同时，实现比标准Transformer更高的吞吐量。

关键设计：该研究的关键设计包括：1) 统一的跨层KV共享框架，允许配置不同的KV共享策略；2) 全面的实验评估，涵盖了生成吞吐量、语言建模性能和下游任务性能；3) 对不同策略的优缺点进行深入分析，为用户提供更明智的选择。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在KV缓存大小减少2倍时，大多数跨层KV共享配置都能实现比标准Transformer更高的吞吐量，同时保持具有竞争力的性能。当进一步减少KV缓存大小时，将所有层的查询与上层的KV配对表现更好，但会增加额外的训练成本和预填充延迟。这些发现为实际应用中选择合适的跨层KV共享策略提供了重要的指导。

🎯 应用场景

该研究成果可应用于各种需要高效LLM推理的场景，例如：在线对话系统、文本生成、机器翻译等。通过选择合适的跨层KV共享策略，可以在保证性能的同时，显著降低内存占用和计算开销，从而提高LLM的推理效率，降低部署成本，并支持更大规模的LLM应用。

📄 摘要（原文）

Recently, sharing key-value (KV) cache across layers has been found effective in efficient inference of large language models (LLMs). To systematically investigate different techniques of cross-layer KV sharing, we propose a unified framework that covers several recent methods and their novel variants. We conduct comprehensive experiments on all the configurations of the framework, evaluating their generation throughput and performance in language modeling and downstream tasks. We find that when reducing the size of the KV cache by 2$\times$, most configurations can achieve higher throughput than standard transformers while maintaining competitive performance. When further reducing the size of the KV cache, however, pairing queries of all layers with KVs of upper layers performs better, at the expense of additional training cost and prefilling latency. We hope that this work will help users make more informed choices of cross-layer KV sharing approaches and facilitate future research on efficient LLM inference.