In-context KV-Cache Eviction for LLMs via Attention-Gate
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866912332011012096 |
|---|---|
| author | Zeng, Zihao Lin, Bokai Hou, Tianqi Zhang, Hao Deng, Zhijie |
| author_facet | Zeng, Zihao Lin, Bokai Hou, Tianqi Zhang, Hao Deng, Zhijie |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2410_12876 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | In-context KV-Cache Eviction for LLMs via Attention-Gate Zeng, Zihao Lin, Bokai Hou, Tianqi Zhang, Hao Deng, Zhijie Computation and Language Machine Learning 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. |
| title | In-context KV-Cache Eviction for LLMs via Attention-Gate |
| topic | Computation and Language Machine Learning |
| url | https://arxiv.org/abs/2410.12876 |