Large Language Models (LLMs), despite their recent impressive accomplishments, are notably cost-prohibitive to deploy, particularly for applications involving long-content generation, such as dialogue systems and story writing. Often, a large amount of transient state information, referred to as the KV cache, is stored in GPU memory in addition to model parameters, scaling linearly with the sequence length and batch size. In this paper, we introduce a novel approach for implementing the KV cache which significantly reduces its memory footprint. Our approach is based on the noteworthy observation that a small portion of tokens contributes most of the value when computing attention scores. We call these tokens Heavy Hitters (H$_2$). Through a comprehensive investigation, we find that (i) the emergence of H$_2$ is natural and strongly correlates with the frequent co-occurrence of tokens in the text, and (ii) removing them results in significant performance degradation. Based on these insights, we propose Heavy Hitter Oracle (H$_2$O), a KV cache eviction policy that dynamically retains a balance of recent and H$_2$ tokens. We formulate the KV cache eviction as a dynamic submodular problem and prove (under mild assumptions) a theoretical guarantee for our novel eviction algorithm which could help guide future work. We validate the accuracy of our algorithm with OPT, LLaMA, and GPT-NeoX across a wide range of tasks. Our implementation of H$_2$O with 20% heavy hitters improves the throughput over three leading inference systems DeepSpeed Zero-Inference, Hugging Face Accelerate, and FlexGen by up to 29$\times$, 29$\times$, and 3$\times$ on OPT-6.7B and OPT-30B. With the same batch size, H2O can reduce the latency by up to 1.9$\times$. The code is available at https://github.com/FMInference/H2O.
翻译:大语言模型(LLMs)尽管近期取得了令人瞩目的成就,但其部署成本高昂,尤其适用于涉及长内容生成的应用场景,如对话系统和故事写作。通常,除模型参数外,大量瞬态状态信息(称为KV缓存)会存储在GPU内存中,其大小随序列长度和批处理规模线性增长。本文提出一种实现KV缓存的新方法,显著降低了其内存占用。该方法基于一个值得注意的观察:在计算注意力分数时,一小部分token贡献了绝大部分价值。我们将这些token称为重击者(H$_2$)。通过全面研究,我们发现:(i)H$_2$的出现是自然现象,且与文本中token的频繁共现强相关;(ii)移除这些token会导致性能显著下降。基于这些见解,我们提出重击者预言器(H$_2$O),一种动态保留近期token与H$_2$ token平衡的KV缓存驱逐策略。我们将KV缓存驱逐形式化为动态子模问题,并在温和假设下证明了我们新颖驱逐算法的理论保证,这有助于指导未来工作。我们通过OPT、LLaMA和GPT-NeoX在广泛任务中验证了算法的准确性。在OPT-6.7B和OPT-30B上,使用20%重击者的H$_2$O实现相较于DeepSpeed Zero-Inference、Hugging Face Accelerate和FlexGen三大主流推理系统,吞吐量提升分别高达29倍、29倍和3倍。在相同批处理规模下,H2O可将延迟降低高达1.9倍。代码开源地址:https://github.com/FMInference/H2O。