LeanK: Learnable K Cache Channel Pruning for Efficient Decoding
2508.02215v1
cs.LG, cs.AI, cs.CL
2025-08-09
Авторы:
Yike Zhang, Zhiyuan He, Huiqiang Jiang, Chengruidong Zhang, Yuqing Yang, Jianyong Wang, Lili Qiu
Резюме на русском
Large language models (LLMs) обеспечивают высокую эффективность для задач с длинным контекстом, однако сталкиваются с проблемами эффективности в связи с быстрым увеличением размера ключ-значение (KV) кэша. Предлагаемое решение — LeanK: Learnable K Cache Channel Pruning — — метод, основанный на технике канальной сжатия KV-кэша, использующий новшество двухэтапной обучаемой сжимаемости. LeanK удаляет незначительные канали K-кэша с помощью чистого обучения и удовлетворяет требованиям к статической сжимаемости и аппаратной алгоритмичности. Для повышения производительности включена пользовательская шина аттенции, которая обеспечивает до 70% снижения размера K-кэша и ускорение обработки в 1,3 раза. На основе экспериментов показано, что LeanK эффективно улучшает время выполнения LLM в задачах с длинным контекстом, не приводя к потере точности. Анализ также дает понимание влияния подхода на модельные каналы и заголовки аттенции во время работы с длинным контекстом. Детали реализации доступны по адресу https://aka.ms/LeanK.
Abstract
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.
Ссылки и действия
Дополнительные ресурсы: