LeanK: Poda de Canales de Caché K Aprendible para una Decodificación Eficiente

LeanK: Learnable K Cache Channel Pruning for Efficient Decoding

August 4, 2025

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

cs.AI

Resumen

Los modelos de lenguaje de gran escala (LLMs) permiten tareas de contexto extenso, pero enfrentan desafíos de eficiencia debido al crecimiento de la caché clave-valor (KV). Proponemos LeanK, un método basado en aprendizaje que poda los canales de la caché clave (K) no importantes aprovechando la esparsidad estática de canales. Con un novedoso proceso de entrenamiento en dos etapas, LeanK aprende una máscara estática a nivel de canal que puede satisfacer un ratio de esparsidad específico y los requisitos de alineación de hardware. LeanK reduce la memoria de GPU y acelera la decodificación sin sacrificar precisión. Los experimentos demuestran una reducción de hasta el 70% en la caché K y del 16%-18% en la memoria de la caché V. Un núcleo de decodificación personalizado permite una aceleración de 1.3x en el cálculo de atención. También proporcionamos insights sobre los canales del modelo y las cabezas de atención durante la inferencia de contexto extenso mediante el análisis de la distribución de importancia aprendida. Nuestro código está disponible en 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.