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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Liu, Minghui, Rabbani, Tahseen, O'Halloran, Tony, Sankaralingam, Ananth, Hartley, Mary-Anne, Huang, Furong, Fermüller, Cornelia, Aloimonos, Yiannis
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908393908731904
author Liu, Minghui
Rabbani, Tahseen
O'Halloran, Tony
Sankaralingam, Ananth
Hartley, Mary-Anne
Huang, Furong
Fermüller, Cornelia
Aloimonos, Yiannis
author_facet Liu, Minghui
Rabbani, Tahseen
O'Halloran, Tony
Sankaralingam, Ananth
Hartley, Mary-Anne
Huang, Furong
Fermüller, Cornelia
Aloimonos, Yiannis
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2412_16187
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle HashEvict: A Pre-Attention KV Cache Eviction Strategy using Locality-Sensitive Hashing
Liu, Minghui
Rabbani, Tahseen
O'Halloran, Tony
Sankaralingam, Ananth
Hartley, Mary-Anne
Huang, Furong
Fermüller, Cornelia
Aloimonos, Yiannis
Machine Learning
Artificial Intelligence
Computation and Language
Data Structures and Algorithms
Performance
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.
title HashEvict: A Pre-Attention KV Cache Eviction Strategy using Locality-Sensitive Hashing
topic Machine Learning
Artificial Intelligence
Computation and Language
Data Structures and Algorithms
Performance
url https://arxiv.org/abs/2412.16187