LeanK : Élagage des canaux du cache K apprenable pour un décodage efficace

LeanK: Learnable K Cache Channel Pruning for Efficient Decoding

August 4, 2025

papers.authors: Yike Zhang, Zhiyuan He, Huiqiang Jiang, Chengruidong Zhang, Yuqing Yang, Jianyong Wang, Lili Qiu

cs.AI

papers.abstract

Les grands modèles de langage (LLMs) permettent de réaliser des tâches à contexte long, mais rencontrent des défis d'efficacité en raison de la croissance du cache clé-valeur (KV). Nous proposons LeanK, une méthode basée sur l'apprentissage qui élimine les canaux de cache clé (K) non importants en exploitant la sparsité statique des canaux. Grâce à un processus d'entraînement en deux étapes innovant, LeanK apprend un masque statique par canal qui peut satisfaire un ratio de sparsité spécifique et les exigences d'alignement matériel. LeanK réduit la mémoire GPU et accélère le décodage sans sacrifier la précision. Les expériences démontrent une réduction allant jusqu'à 70 % du cache K et de 16 % à 18 % du cache V. Un noyau de décodage personnalisé permet une accélération de 1,3x pour le calcul de l'attention. Nous fournissons également des insights sur les canaux du modèle et les têtes d'attention lors de l'inférence à contexte long en analysant la distribution d'importance apprise. Notre code est disponible à l'adresse https://aka.ms/LeanK.

English

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.