LeanK: 효율적인 디코딩을 위한 학습 가능한 K 캐시 채널 프루닝

초록

대규모 언어 모델(LLMs)은 장문맥 작업을 가능하게 하지만, 키-값(KV) 캐시의 증가로 인해 효율성 문제에 직면하고 있습니다. 우리는 LeanK를 제안합니다. 이는 정적 채널 희소성을 활용하여 중요하지 않은 키(K) 캐시 채널을 제거하는 학습 기반 방법입니다. LeanK는 새로운 두 단계 학습 프로세스를 통해 특정 희소성 비율과 하드웨어 정렬 요구 사항을 충족할 수 있는 채널별 정적 마스크를 학습합니다. LeanK는 정확도를 희생하지 않으면서 GPU 메모리를 절약하고 디코딩 속도를 가속화합니다. 실험 결과, 최대 70%의 K 캐시와 16%-18%의 V 캐시 메모리 감소를 보여줍니다. 맞춤형 디코딩 커널은 어텐션 계산에서 1.3배의 속도 향상을 가능하게 합니다. 또한 학습된 중요도 분포를 분석함으로써 장문맥 추론 과정에서의 모델 채널과 어텐션 헤드에 대한 통찰을 제공합니다. 우리의 코드는 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.

LeanK: 효율적인 디코딩을 위한 학습 가능한 K 캐시 채널 프루닝

LeanK: Learnable K Cache Channel Pruning for Efficient Decoding

초록

Support