arXiv:2412.16187cs.LGcs.AI2024-12被引 9

用哈希技术压缩大模型推理缓存,省内存还快

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

  • 用局部敏感哈希快速找不相关的缓存项
  • 压缩率30%-70%仍保持高精度
  • 适合需要长上下文的部署场景

基于Transformer的大语言模型通过键值(KV)缓存加速推理,但占用大量GPU内存。本文提出HashEvict算法,利用局部敏感哈希(LSH)压缩KV缓存。该方法通过计算当前查询词与缓存键的二值化高斯投影之间的汉明距离,快速定位与当前查询余弦不相似的缓存项,投影长度远小于嵌入维度。在GPU内存中维护轻量级二进制结构以支持快速计算。与现有依赖注意力计算决定保留的压缩策略不同,HashEvict在注意力计算前做出淘汰决策,降低计算开销。此外,HashEvict为动态机制:每解码一步,当前词的键值替换预计产生最低注意力得分的缓存项。实验表明,HashEvict可在推理、多选、长文本检索和摘要任务中实现30%-70%的缓存压缩率,同时保持高性能。

原文摘要 · Abstract (English)

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.

缓存优化哈希技术大模型推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。