In-context KV-Cache Eviction for LLMs via Attention-Gate

TL;DR

This paper introduces Attention-Gate, a lightweight module for dynamic KV-Cache eviction in large language models, reducing memory bottlenecks and improving inference efficiency without significant overhead.

Contribution

It proposes a novel Attention-Gate mechanism that enables dynamic, tunable eviction of tokens in KV-Cache during LLM inference, adaptable to pre-trained models.

Findings

Reduces memory usage during inference.

Improves model efficiency and performance.

Minimal additional computational overhead.

Abstract

The KV-Cache technique has become the standard for the inference of large language models (LLMs). Yet, it is widely criticized that KV-Cache can become a bottleneck of the LLM inference system. This paper enables a novel dynamic KV-Cache eviction policy by injecting a lightweight module called Attention-Gate to the model. It accepts the global context as input and yields eviction flags for each token. The self-attention modules in the model proceed according to the flags and cache only a subset of the KV states for next token prediction. The Attention-Gates can yield various flags for different heads and layers and be easily tuned on top of a pre-trained LLM via continual pre-training or supervised fine-tuning. The computational and memory overhead introduced by Attention-Gates can be minimal. We empirically evaluate the proposed approach across multiple scenarios, showing that effective eviction of redundant tokens can not only improve efficiency but also enhance performance.