SnapKV: O LLM Sabe o Que Você Está Procurando Antes da Geração

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) têm feito progressos notáveis no processamento de contextos extensos, com o cache Chave-Valor (KV) desempenhando um papel vital na melhoria de seu desempenho. No entanto, o crescimento do cache KV em resposta ao aumento do comprimento da entrada apresenta desafios para a eficiência de memória e tempo. Para resolver esse problema, este artigo introduz o SnapKV, uma abordagem inovadora e livre de ajuste fino que minimiza eficientemente o tamanho do cache KV enquanto ainda oferece desempenho comparável em aplicações do mundo real. Descobrimos que cada cabeça de atenção no modelo consistentemente se concentra em características específicas de atenção do prompt durante a geração. Enquanto isso, esse padrão robusto pode ser obtido a partir de uma janela de `observação' localizada no final dos prompts. Com base nessa percepção, o SnapKV comprime automaticamente os caches KV selecionando posições importantes agrupadas de KV para cada cabeça de atenção. Nossa abordagem reduz significativamente a sobrecarga computacional crescente e a pegada de memória ao processar sequências de entrada longas. Especificamente, o SnapKV alcança uma velocidade de decodificação consistente com um aumento de 3,6x na velocidade de geração e uma melhoria de 8,2x na eficiência de memória em comparação com a linha de base ao processar entradas de 16K tokens. Ao mesmo tempo, mantém um desempenho comparável aos modelos de linha de base em 16 conjuntos de dados de sequências longas. Além disso, o SnapKV pode processar até 380K tokens de contexto em uma única GPU A100-80GB usando a implementação do HuggingFace com pequenas alterações, exibindo apenas uma queda de precisão insignificante no teste Needle-in-a-Haystack. Estudos abrangentes adicionais sugerem o potencial do SnapKV para aplicações práticas.

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: O LLM Sabe o Que Você Está Procurando Antes da Geração

SnapKV: LLM Knows What You are Looking for Before Generation

Resumo

Support