大语言模型(LLM)近年来取得了令人瞩目的成就,但其部署成本高昂,特别是在对话系统和故事写作等长文本生成应用中。除模型参数外,大量瞬态状态信息(称为KV缓存)需要存储在GPU内存中,且随序列长度和批处理大小线性增长。为解决这一问题,研究人员提出了一种名为H2O(Heavy-Hitter Oracle)的新颖KV缓存实现方法,显著降低了内存占用。
H2O基于一个重要发现:在计算注意力分数时,只有一小部分词元贡献了大部分价值。这些词元被称为"重要访问"(Heavy Hitters, H2)。研究表明:
基于这些洞见,H2O提出了一种动态KV缓存淘汰策略,在保留最近词元和H2词元之间保持平衡。
研究团队将KV缓存淘汰问题建模为动态子模问题,并在温和假设下为其新颖的淘汰算法提供了理论保证,为未来的研究提供了指导。H2O的主要技术创新包括:
研究人员在OPT、LLaMA和GPT-NeoX等多个模型上对H2O进行了广泛的任务验证。实验结果表明:
这些性能提升使H2O成为提高大语言模型生成推理效率的有力工具。