LeanK: Lernbare K-Cache-Kanalreduzierung für effiziente Dekodierung

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

Große Sprachmodelle (LLMs) ermöglichen Aufgaben mit langem Kontext, stehen jedoch aufgrund des wachsenden Key-Value (KV)-Caches vor Effizienzproblemen. Wir schlagen LeanK vor, eine lernbasierte Methode, die unwichtige Key (K)-Cache-Kanäle durch die Nutzung statischer Kanalsparsamkeit beschneidet. Mit einem neuartigen zweistufigen Trainingsprozess lernt LeanK kanalweise statische Masken, die ein bestimmtes Sparsamkeitsverhältnis und Hardware-Ausrichtungsanforderungen erfüllen können. LeanK reduziert den GPU-Speicher und beschleunigt die Decodierung, ohne die Genauigkeit zu beeinträchtigen. Experimente zeigen eine Reduktion des K-Caches um bis zu 70 % und des V-Cache-Speichers um 16 %–18 %. Ein benutzerdefinierter Decodierungskernel ermöglicht eine 1,3-fache Beschleunigung der Aufmerksamkeitsberechnung. Wir bieten auch Einblicke in Modellkanäle und Aufmerksamkeitsköpfe während der Inferenz mit langem Kontext durch die Analyse der gelernten Wichtigkeitsverteilung. Unser Code ist verfügbar unter 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.