CoKV: Optimizing KV Cache Allocation via Cooperative Game
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2025
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866910843121172480 |
|---|---|
| author | Sun, Qiheng Zhang, Hongwei Xia, Haocheng Zhang, Jiayao Liu, Jinfei Ren, Kui |
| author_facet | Sun, Qiheng Zhang, Hongwei Xia, Haocheng Zhang, Jiayao Liu, Jinfei Ren, Kui |
| contents | Large language models (LLMs) have achieved remarkable success on various aspects of human life. However, one of the major challenges in deploying these models is the substantial memory consumption required to store key-value pairs (KV), which imposes significant resource demands. Recent research has focused on KV cache budget allocation, with several approaches proposing head-level budget distribution by evaluating the importance of individual attention heads. These methods, however, assess the importance of heads independently, overlooking their cooperative contributions within the model, which may result in a deviation from their true impact on model performance. In light of this limitation, we propose CoKV, a novel method that models the cooperation between heads in model inference as a cooperative game. By evaluating the contribution of each head within the cooperative game, CoKV can allocate the cache budget more effectively. Extensive experiments show that CoKV achieves state-of-the-art performance on the LongBench benchmark using LLama-3-8B-Instruct and Mistral-7B models. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_17501 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | CoKV: Optimizing KV Cache Allocation via Cooperative Game Sun, Qiheng Zhang, Hongwei Xia, Haocheng Zhang, Jiayao Liu, Jinfei Ren, Kui Machine Learning Artificial Intelligence Large language models (LLMs) have achieved remarkable success on various aspects of human life. However, one of the major challenges in deploying these models is the substantial memory consumption required to store key-value pairs (KV), which imposes significant resource demands. Recent research has focused on KV cache budget allocation, with several approaches proposing head-level budget distribution by evaluating the importance of individual attention heads. These methods, however, assess the importance of heads independently, overlooking their cooperative contributions within the model, which may result in a deviation from their true impact on model performance. In light of this limitation, we propose CoKV, a novel method that models the cooperation between heads in model inference as a cooperative game. By evaluating the contribution of each head within the cooperative game, CoKV can allocate the cache budget more effectively. Extensive experiments show that CoKV achieves state-of-the-art performance on the LongBench benchmark using LLama-3-8B-Instruct and Mistral-7B models. |
| title | CoKV: Optimizing KV Cache Allocation via Cooperative Game |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2502.17501 |