LeanK: Learnable K Cache Channel Pruning for Efficient Decoding
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866912518385958912 |
|---|---|
| author | Zhang, Yike He, Zhiyuan Jiang, Huiqiang Zhang, Chengruidong Yang, Yuqing Wang, Jianyong Qiu, Lili |
| author_facet | Zhang, Yike He, Zhiyuan Jiang, Huiqiang Zhang, Chengruidong Yang, Yuqing Wang, Jianyong Qiu, Lili |
| contents | Large language models (LLMs) enable long-context tasks but face efficiency challenges due to the growing key-value (KV) cache. We propose LeanK, a learning-based method that prunes unimportant key (K) cache channels by leveraging static channel sparsity. With a novel two-stage training process, LeanK learns channel-wise static mask that could satisfy specific sparsity ratio and hardware alignment requirement. LeanK reduces GPU memory and accelerates decoding without sacrificing accuracy. Experiments demonstrate up to 70% K cache and 16%-18% V cache memory reduction. Custom decoding kernel enables 1.3x speedup for attention computation. We also provide insights into model channels and attention heads during long-context inference by analyzing the learned importance distribution. Our code is available at https://aka.ms/LeanK. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2508_02215 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | LeanK: Learnable K Cache Channel Pruning for Efficient Decoding Zhang, Yike He, Zhiyuan Jiang, Huiqiang Zhang, Chengruidong Yang, Yuqing Wang, Jianyong Qiu, Lili Machine Learning Artificial Intelligence Computation and Language Large language models (LLMs) enable long-context tasks but face efficiency challenges due to the growing key-value (KV) cache. We propose LeanK, a learning-based method that prunes unimportant key (K) cache channels by leveraging static channel sparsity. With a novel two-stage training process, LeanK learns channel-wise static mask that could satisfy specific sparsity ratio and hardware alignment requirement. LeanK reduces GPU memory and accelerates decoding without sacrificing accuracy. Experiments demonstrate up to 70% K cache and 16%-18% V cache memory reduction. Custom decoding kernel enables 1.3x speedup for attention computation. We also provide insights into model channels and attention heads during long-context inference by analyzing the learned importance distribution. Our code is available at https://aka.ms/LeanK. |
| title | LeanK: Learnable K Cache Channel Pruning for Efficient Decoding |
| topic | Machine Learning Artificial Intelligence Computation and Language |
| url | https://arxiv.org/abs/2508.02215 |