HashEvict: A Pre-Attention KV Cache Eviction Strategy using Locality-Sensitive Hashing

📄 arXiv: 2412.16187v3 📥 PDF

作者: Minghui Liu, Tahseen Rabbani, Tony O'Halloran, Ananth Sankaralingam, Mary-Anne Hartley, Furong Huang, Cornelia Fermüller, Yiannis Aloimonos

分类: cs.LG, cs.AI, cs.CL, cs.DS, cs.PF

发布日期: 2024-12-13 (更新: 2025-06-04)

备注: 10 pages, 6 figures, 2 tables

💡 一句话要点

HashEvict：利用局部敏感哈希的预注意力KV缓存淘汰策略，降低LLM推理的GPU内存消耗。

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

关键词: 大型语言模型 KV缓存 局部敏感哈希 预注意力 内存压缩

📋 核心要点

  1. 大型语言模型推理时KV缓存占用大量GPU内存，成为性能瓶颈。
  2. HashEvict利用LSH快速找到与当前token不相似的缓存token，实现预注意力淘汰。
  3. 实验表明，HashEvict能在多种任务中压缩KV缓存30%-70%，同时保持高性能。

📝 摘要（中文）

本文提出了一种名为HashEvict的算法，该算法利用局部敏感哈希（LSH）来压缩Transformer大型语言模型（LLM）中的键值（KV）缓存，从而显著加速推理。HashEvict通过计算当前查询token和缓存token键的二值化高斯投影之间的汉明距离，快速定位缓存中与当前查询token余弦不相似的token。投影长度远小于嵌入维度，并在GPU内存中维护一个轻量级的二进制结构以加速计算。与现有通过计算注意力来决定token保留的压缩策略不同，HashEvict在注意力计算之前做出决策，从而降低了计算成本。此外，HashEvict是动态的，在每个解码步骤中，当前token的键和值会替换预计产生最低注意力分数的token嵌入。实验表明，HashEvict可以在推理、多项选择、长上下文检索和摘要任务中压缩KV缓存30%-70%，同时保持高性能。

🔬 方法详解

问题定义：大型语言模型（LLM）在推理过程中，需要维护一个键值（KV）缓存来存储过去token的键和值嵌入，以加速后续token的生成。然而，这个KV缓存会消耗大量的GPU内存，尤其是在处理长序列时，成为性能瓶颈。现有的压缩策略通常需要在注意力计算之后才能确定哪些token应该被保留，计算开销较大。

核心思路：HashEvict的核心思路是在注意力计算之前，利用局部敏感哈希（LSH）快速估计token之间的相似度，并淘汰那些与当前token不相似的token。通过预先淘汰不重要的token，可以有效压缩KV缓存，降低GPU内存占用，同时减少后续注意力计算的开销。

技术框架：HashEvict算法主要包含以下几个阶段：1) 高斯投影：将token的键嵌入进行高斯随机投影，降低维度。2) 二值化：将投影后的向量进行二值化，得到二进制表示。3) 汉明距离计算：计算当前查询token的二进制表示与缓存中token的二进制表示之间的汉明距离。4) 淘汰决策：根据汉明距离，淘汰与当前查询token最不相似的token。5) 动态更新：在每个解码步骤中，当前token的键和值替换被淘汰的token。

关键创新：HashEvict的关键创新在于其预注意力淘汰机制。与传统的后注意力淘汰方法相比，HashEvict能够在注意力计算之前做出淘汰决策，从而显著降低计算成本。此外，HashEvict使用LSH来快速估计token之间的相似度，避免了昂贵的余弦相似度计算。

关键设计：HashEvict的关键设计包括：1) 高斯投影长度：投影长度的选择需要在计算效率和相似度估计精度之间进行权衡。论文中可能给出了具体的选择方法或实验结果。2) 二值化阈值：二值化阈值的选择会影响哈希的碰撞概率。3) 淘汰比例：需要根据具体的任务和模型选择合适的淘汰比例，以在压缩率和性能之间取得平衡。4) 轻量级二进制结构：为了加速汉明距离计算，HashEvict在GPU内存中维护了一个轻量级的二进制结构。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，HashEvict可以在推理、多项选择、长上下文检索和摘要任务中压缩KV缓存30%-70%，同时保持高性能。这意味着在相同的GPU内存容量下，可以处理更长的上下文，或者运行更大的模型。具体的性能提升幅度取决于任务类型、模型大小和压缩比例等因素，论文中应该给出了详细的实验数据和对比结果。

🎯 应用场景

HashEvict可应用于各种需要处理长序列的大型语言模型推理场景，例如机器翻译、文本摘要、问答系统和代码生成等。通过降低GPU内存占用，HashEvict可以使更大的模型能够在资源受限的设备上运行，或者在相同的硬件资源下处理更长的序列，从而提高模型的性能和效率。该技术还有助于降低LLM部署和运行的成本，促进其更广泛的应用。

📄 摘要（原文）

Transformer-based large language models (LLMs) use the key-value (KV) cache to significantly accelerate inference by storing the key and value embeddings of past tokens. However, this cache consumes significant GPU memory. In this work, we introduce HashEvict, an algorithm that uses locality-sensitive hashing (LSH) to compress the KV cache. HashEvict quickly locates tokens in the cache that are cosine dissimilar to the current query token. This is achieved by computing the Hamming distance between binarized Gaussian projections of the current token query and cached token keys, with a projection length much smaller than the embedding dimension. We maintain a lightweight binary structure in GPU memory to facilitate these calculations. Unlike existing compression strategies that compute attention to determine token retention, HashEvict makes these decisions pre-attention, thereby reducing computational costs. Additionally, HashEvict is dynamic - at every decoding step, the key and value of the current token replace the embeddings of a token expected to produce the lowest attention score. We demonstrate that HashEvict can compress the KV cache by 30%-70% while maintaining high performance across reasoning, multiple-choice, long-context retrieval and summarization tasks.