LeanK: Learnable K Cache Channel Pruning for Efficient Decoding
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.
Code (0)
등록된 구현이 없습니다.
Similar Papers 제목 키워드 기반
Unifying KV Cache Compression for Large Language Models with LeanKV
Large language models (LLMs) exhibit exceptional performance but incur significant serving costs due to their substantial memory requirements, with the key-value (KV) cache being a primary bottleneck. Existing KV cache c…
GPUQuantizationRotation-Aligned Key Channel Pruning for Efficient Vision-Language Model Inference
Vision-Language Models suffer severe KV cache pressure at inference, as a single image often encodes into thousands of tokens. Most existing methods exploit token sparsity through token pruning, but permanently discardin…
Residual vector quantization for KV cache compression in large language model
KV cache compression methods have mainly relied on scalar quantization techniques to reduce the memory requirements during decoding. In this work, we apply residual vector quantization, which has been widely used for hig…
Audio CompressionLanguage ModelingLanguage ModellingLarge Language Model+1KVzap: Fast, Adaptive, and Faithful KV Cache Pruning
Growing context lengths in transformer-based language models have made the key-value (KV) cache a critical inference bottleneck. While many KV cache pruning methods have been proposed, they have not yet been adopted in m…
PermLLM: Learnable Channel Permutation for N:M Sparse Large Language Models
Channel permutation is a powerful technique for enhancing the accuracy of N:M sparse models by reordering the channels of weight matrices to prioritize the retention of important weights. However, traditional channel per…