LeanK: Poda de Canais de Cache K Aprendível para Decodificação 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

Resumo

Modelos de linguagem de grande escala (LLMs) permitem tarefas de contexto longo, mas enfrentam desafios de eficiência devido ao crescimento do cache de chave-valor (KV). Propomos o LeanK, um método baseado em aprendizado que poda canais de cache de chave (K) irrelevantes, aproveitando a esparsidade estática de canais. Com um processo de treinamento em duas etapas inovador, o LeanK aprende uma máscara estática por canal que pode satisfazer uma taxa de esparsidade específica e requisitos de alinhamento de hardware. O LeanK reduz a memória da GPU e acelera a decodificação sem sacrificar a precisão. Experimentos demonstram uma redução de até 70% no cache K e de 16%-18% no cache V. Um kernel de decodificação personalizado permite uma aceleração de 1,3x no cálculo de atenção. Também fornecemos insights sobre os canais do modelo e cabeças de atenção durante a inferência de contexto longo, analisando a distribuição de importância aprendida. Nosso código está disponível em 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.