LeanK: Leerbare K Cache Kanaalpruning voor Efficiënt Decoderen

Samenvatting

Grote taalmodellen (LLMs) maken lang-context taken mogelijk, maar kampen met efficiëntieproblemen door de groeiende key-value (KV) cache. Wij stellen LeanK voor, een op leren gebaseerde methode die onbelangrijke key (K) cache-kanalen snoeit door gebruik te maken van statische kanaalsparsheid. Met een innovatief tweestaps trainingsproces leert LeanK een kanaalgewijs statisch masker dat aan specifieke sparsityratio's en hardware-uitlijningsvereisten voldoet. LeanK vermindert het GPU-geheugen en versnelt het decoderen zonder de nauwkeurigheid op te offeren. Experimenten tonen een reductie van tot 70% in K-cache en 16%-18% in V-cache geheugen aan. Een aangepaste decoderingkernel zorgt voor een 1,3x versnelling van de aandachtberekening. We bieden ook inzichten in modelkanalen en aandachtkoppen tijdens lang-context inferentie door de geleerde belangrijkheidsverdeling te analyseren. Onze code is beschikbaar op 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: Leerbare K Cache Kanaalpruning voor Efficiënt Decoderen

LeanK: Learnable K Cache Channel Pruning for Efficient Decoding

Samenvatting

Support