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.
6894106d741a16f544fbcea6
Yike Zhang
claimed_verified
2025-08-07T10:39:04.154Z
6541159afcbd1aa00682ec66
/avatars/94fd9ce0ea534e5d84736c1a1eb949e7.svg
zhangyik21
user
zhangyik21
6894106d741a16f544fbcea7
Zhiyuan He
claimed_verified
2025-08-07T10:38:57.434Z
6455f5cababbbbd3486d6ee3
/avatars/b6c8f65fd2bef8a00aa3269856ea238e.svg
Zhiyuan He
user
hzy46
6894106d741a16f544fbcea8
Huiqiang Jiang
6894106d741a16f544fbcea9
Chengruidong Zhang
6894106d741a16f544fbceaa
Yuqing Yang
6894106d741a16f544fbceab
Jianyong Wang
6894106d741a16f544fbceac
Lili Qiu
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.
LeanK, a learning-based method, prunes unimportant key cache channels in large language models to reduce memory usage and accelerate decoding without sacrificing accuracy.
large language models key-value cache channel sparsity two-stage training channel-wise static mask GPU memory decoding speedup attention computation long-context inference importance distribution