SnapKV: L'LLM Sa Cosa Stai Cercando Prima della Generazione

Abstract

I modelli linguistici di grandi dimensioni (LLMs) hanno compiuto progressi significativi nell'elaborazione di contesti estesi, con la cache Key-Value (KV) che svolge un ruolo cruciale nel migliorare le loro prestazioni. Tuttavia, la crescita della cache KV in risposta all'aumento della lunghezza dell'input pone sfide all'efficienza della memoria e del tempo. Per affrontare questo problema, questo articolo introduce SnapKV, un approccio innovativo e privo di fine-tuning che riduce efficacemente le dimensioni della cache KV mantenendo prestazioni comparabili nelle applicazioni reali. Abbiamo scoperto che ogni testa di attenzione nel modello si concentra costantemente su specifiche caratteristiche di attenzione del prompt durante la generazione. Nel frattempo, questo schema robusto può essere ottenuto da una finestra di "osservazione" situata alla fine dei prompt. Basandoci su questa intuizione, SnapKV comprime automaticamente le cache KV selezionando posizioni KV importanti e raggruppate per ogni testa di attenzione. Il nostro approccio riduce significativamente il sovraccarico computazionale e l'impronta di memoria in crescita durante l'elaborazione di sequenze di input lunghe. Nello specifico, SnapKV raggiunge una velocità di decodifica costante con un aumento di 3,6x nella velocità di generazione e un miglioramento di 8,2x nell'efficienza della memoria rispetto alla baseline quando elabora input di 16K token. Allo stesso tempo, mantiene prestazioni comparabili ai modelli baseline su 16 dataset di sequenze lunghe. Inoltre, SnapKV può elaborare fino a 380K token di contesto su una singola GPU A100-80GB utilizzando l'implementazione di HuggingFace con modifiche minime, mostrando solo un calo di precisione trascurabile nel test Needle-in-a-Haystack. Ulteriori studi approfonditi suggeriscono il potenziale di SnapKV per applicazioni pratiche.

English

Large Language Models (LLMs) have made remarkable progress in processing extensive contexts, with the Key-Value (KV) cache playing a vital role in enhancing their performance. However, the growth of the KV cache in response to increasing input length poses challenges to memory and time efficiency. To address this problem, this paper introduces SnapKV, an innovative and fine-tuning-free approach that efficiently minimizes KV cache size while still delivering comparable performance in real-world applications. We discover that each attention head in the model consistently focuses on specific prompt attention features during generation. Meanwhile, this robust pattern can be obtained from an `observation' window located at the end of the prompts. Drawing on this insight, SnapKV automatically compresses KV caches by selecting clustered important KV positions for each attention head. Our approach significantly reduces the growing computational overhead and memory footprint when processing long input sequences. Specifically, SnapKV achieves a consistent decoding speed with a 3.6x increase in generation speed and an 8.2x enhancement in memory efficiency compared to baseline when processing inputs of 16K tokens. At the same time, it maintains comparable performance to baseline models across 16 long sequence datasets. Moreover, SnapKV can process up to 380K context tokens on a single A100-80GB GPU using HuggingFace implementation with minor changes, exhibiting only a negligible accuracy drop in the Needle-in-a-Haystack test. Further comprehensive studies suggest SnapKV's potential for practical applications.

SnapKV: L'LLM Sa Cosa Stai Cercando Prima della Generazione

SnapKV: LLM Knows What You are Looking for Before Generation

Abstract

Support