InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management

📄 arXiv: 2406.19707v1 📥 PDF

作者: Wonbeom Lee, Jungi Lee, Junghwan Seo, Jaewoong Sim

分类: cs.LG, cs.DC

发布日期: 2024-06-28

备注: OSDI 2024

💡 一句话要点

InfiniGen：通过动态KV缓存管理实现大语言模型的高效生成式推理

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

关键词: 大语言模型 KV缓存管理 生成式推理 长文本生成 卸载系统

📋 核心要点

  1. 现有LLM推理在生成长文本时，KV缓存占用大量内存，导致效率瓶颈，尤其是在卸载系统中。
  2. InfiniGen通过预测并仅预取对后续注意力层至关重要的少量token的KV缓存，减少了内存访问开销。
  3. 实验表明，InfiniGen在卸载系统中性能提升高达3.00倍，并显著提高了模型准确性。

📝 摘要（中文）

基于Transformer的大语言模型(LLM)在各种自然语言处理任务中表现出令人印象深刻的性能。然而，为生成长内容提供LLM推理服务面临挑战，因为瞬态状态（即键值(KV)缓存）的巨大内存占用随序列长度和批大小而扩展。本文提出了InfiniGen，一种为长文本生成量身定制的新型KV缓存管理框架，它与现代基于卸载的推理系统协同工作。InfiniGen利用关键洞察力，即通过对当前层的输入以及后续层的部分查询权重和键缓存执行最小排练，可以推测出对于计算Transformer中后续注意力层至关重要的少量重要token。这允许我们仅预取必要的KV缓存条目（而无需全部获取），从而减轻了基于卸载的LLM服务系统中从主机内存获取的开销。在我们对几个代表性LLM的评估中表明，与先前的KV缓存管理方法相比，InfiniGen将现代基于卸载的系统的整体性能提高了高达3.00倍，同时提供了明显更好的模型准确性。

🔬 方法详解

问题定义：现有的大语言模型在生成长文本时，KV缓存会占用大量的内存空间，并且随着序列长度和batch size的增加而线性增长。这导致了推理效率的瓶颈，尤其是在基于卸载的系统中，频繁地从主机内存中读取KV缓存会带来巨大的开销。因此，如何有效地管理KV缓存，减少内存访问，提高长文本生成的推理效率是一个关键问题。

核心思路：InfiniGen的核心思路是，并非所有的KV缓存条目都是同等重要的。通过预测哪些token对于计算后续注意力层是至关重要的，并仅预取这些token的KV缓存，可以显著减少内存访问量。这种选择性预取策略避免了不必要的内存读取，从而提高了推理效率。

技术框架：InfiniGen与现有的基于卸载的推理系统协同工作。整体流程如下：1. 当前层接收输入。2. 使用当前层的输入和后续层的部分查询权重和键缓存，进行最小排练，预测重要的token。3. 仅预取预测到的重要token的KV缓存。4. 使用预取的KV缓存计算注意力层。5. 将结果传递给下一层。

关键创新：InfiniGen的关键创新在于提出了一种基于最小排练的KV缓存预测方法。与现有方法需要预取所有KV缓存条目相比，InfiniGen只预取必要的条目，从而显著减少了内存访问开销。这种方法的核心在于利用了Transformer的结构特性，即后续层的计算依赖于前几层的关键token。

关键设计：InfiniGen的关键设计包括：1. 最小排练：使用当前层的输入和后续层的部分查询权重和键缓存，进行少量计算，以预测重要的token。2. 预取策略：根据预测结果，仅预取必要的KV缓存条目。3. 与卸载系统的集成：InfiniGen可以无缝集成到现有的基于卸载的推理系统中，无需修改底层架构。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

InfiniGen在多个代表性LLM上的评估结果表明，与先前的KV缓存管理方法相比，InfiniGen将现代基于卸载的系统的整体性能提高了高达3.00倍，同时提供了明显更好的模型准确性。这表明InfiniGen能够有效地减少内存访问开销，提高推理效率，并且不会牺牲模型性能。

🎯 应用场景

InfiniGen适用于需要生成长文本的各种应用场景，如长篇小说生成、代码生成、对话系统、机器翻译等。通过提高LLM的推理效率，InfiniGen可以降低部署成本，并支持更大规模的应用。未来，该技术可以进一步扩展到其他类型的Transformer模型和任务中，具有广阔的应用前景。

📄 摘要（原文）

Transformer-based large language models (LLMs) demonstrate impressive performance across various natural language processing tasks. Serving LLM inference for generating long contents, however, poses a challenge due to the enormous memory footprint of the transient state, known as the key-value (KV) cache, which scales with the sequence length and batch size. In this paper, we present InfiniGen, a novel KV cache management framework tailored for long-text generation, which synergistically works with modern offloading-based inference systems. InfiniGen leverages the key insight that a few important tokens that are essential for computing the subsequent attention layer in the Transformer can be speculated by performing a minimal rehearsal with the inputs of the current layer and part of the query weight and key cache of the subsequent layer. This allows us to prefetch only the essential KV cache entries (without fetching them all), thereby mitigating the fetch overhead from the host memory in offloading-based LLM serving systems. Our evaluation on several representative LLMs shows that InfiniGen improves the overall performance of a modern offloading-based system by up to 3.00x compared to prior KV cache management methods while offering substantially better model accuracy.