LeanK: Potatura Apprendibile dei Canali della Cache K per una Decodifica Efficiente

LeanK: Learnable K Cache Channel Pruning for Efficient Decoding

August 4, 2025

Autori: Yike Zhang, Zhiyuan He, Huiqiang Jiang, Chengruidong Zhang, Yuqing Yang, Jianyong Wang, Lili Qiu

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) abilitano compiti a contesto lungo ma affrontano sfide di efficienza a causa della crescente cache chiave-valore (KV). Proponiamo LeanK, un metodo basato sull'apprendimento che pota i canali della cache chiave (K) non importanti sfruttando la sparsità statica dei canali. Con un innovativo processo di addestramento in due fasi, LeanK apprende una maschera statica a livello di canale che può soddisfare specifici requisiti di rapporto di sparsità e allineamento hardware. LeanK riduce la memoria GPU e accelera la decodifica senza sacrificare l'accuratezza. Gli esperimenti dimostrano una riduzione fino al 70% della cache K e del 16%-18% della memoria della cache V. Un kernel di decodifica personalizzato consente un'accelerazione di 1,3x nel calcolo dell'attenzione. Forniamo inoltre approfondimenti sui canali del modello e sulle teste di attenzione durante l'inferenza a contesto lungo analizzando la distribuzione appresa dell'importanza. Il nostro codice è disponibile all'indirizzo 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.