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

基于Transformer的大型语言模型在各种自然语言处理任务中表现出卓越的性能。然而，为生成长内容提供LLM推理面临挑战，因为瞬态状态（称为键值（KV）缓存）的内存占用量随着序列长度和批处理大小而增加。本文提出了InfiniGen，一种专为长文本生成量身定制的KV缓存管理框架，与现代基于卸载的推理系统协同工作。InfiniGen利用的关键洞察是，在Transformer中，计算后续注意力层所需的少量重要标记可以通过使用当前层的输入和部分查询权重和后续层的键缓存来进行最小的排练来推断。这使我们只能预取必要的KV缓存条目（而不是全部获取），从而减轻了基于卸载的LLM服务系统中来自主机内存的获取开销。我们对几个典型的LLM进行评估，结果显示，与先前的KV缓存管理方法相比，InfiniGen将现代基于卸载的系统的整体性能提高了高达3.00倍，同时提供了更好的模型准确性。

本文旨在解决使用Transformer进行长文本生成时，由于key-value（KV）缓存的巨大内存占用量而导致的推理速度下降的问题。

本文提出了一种名为InfiniGen的KV缓存管理框架，通过在当前层的输入和下一层的查询权重和键缓存的一部分上执行最小排练，来预测计算后续注意层所需的重要令牌。这样可以仅预取必要的KV缓存条目（而不是全部获取），从而减少基于卸载的LLM服务系统中来自主机内存的获取开销。

实验结果表明，与之前的KV缓存管理方法相比，InfiniGen可以将现代基于卸载的系统的整体性能提高高达3.00倍，同时提供更好的模型准确性。本文使用了几个代表性的LLM进行评估，并且开源了代码。

与此相关的最新研究包括《Large Memory Layers with Product Keys》和《DeepSpeed IO: Efficient IO Training for Large Neural Networks》等。