FastKV: Compressão de Cache KV para Processamento Rápido de Longo Contexto com Propagação Seletiva de Tokens

FastKV: KV Cache Compression for Fast Long-Context Processing with Token-Selective Propagation

February 3, 2025

Autores: Dongwon Jo, Jiwon Song, Yulhwa Kim, Jae-Joon Kim

cs.AI

Resumo

Embora os grandes modelos de linguagem (LLMs) se destaquem no manuseio de sequências de longo contexto, eles requerem caches substanciais de chave-valor (KV) para armazenar informações contextuais, o que pode sobrecarregar significativamente a eficiência computacional e o uso de memória. Esforços anteriores para comprimir esses caches KV focaram principalmente em reduzir as demandas de memória, mas eram limitados em melhorar a latência. Para abordar essa questão, apresentamos o FastKV, um método de compressão de cache KV projetado para melhorar a latência de sequências de longo contexto. Para aprimorar a velocidade de processamento mantendo a precisão, o FastKV adota uma abordagem inovadora de Propagação Seletiva de Tokens (TSP) que mantém todas as informações de contexto nas camadas iniciais dos LLMs e propaga seletivamente apenas uma parte dessas informações em camadas mais profundas, mesmo na fase de pré-preenchimento. Além disso, o FastKV incorpora compressão de cache KV consciente de agrupamento de consultas (GQA) para explorar as vantagens do GQA em eficiência de memória e computacional. Nossos resultados experimentais mostram que o FastKV alcança melhorias de 2,00 vezes e 1,40 vezes no tempo até o primeiro token (TTFT) e na taxa de transferência, respectivamente, em comparação com o HeadKV, o método de compressão de cache KV de última geração. Além disso, o FastKV mantém com sucesso a precisão em benchmarks de longo contexto em níveis comparáveis aos baselines. Nosso código está disponível em https://github.com/dongwonjo/FastKV.

English

While large language models (LLMs) excel at handling long-context sequences, they require substantial key-value (KV) caches to store contextual information, which can heavily burden computational efficiency and memory usage. Previous efforts to compress these KV caches primarily focused on reducing memory demands but were limited in enhancing latency. To address this issue, we introduce FastKV, a KV cache compression method designed to enhance latency for long-context sequences. To enhance processing speeds while maintaining accuracy, FastKV adopts a novel Token-Selective Propagation (TSP) approach that retains the full context information in the initial layers of LLMs and selectively propagates only a portion of this information in deeper layers even in the prefill stage. Additionally, FastKV incorporates grouped-query attention (GQA)-aware KV cache compression to exploit the advantages of GQA in both memory and computational efficiency. Our experimental results show that FastKV achieves 2.00times and 1.40times improvements in time-to-first-token (TTFT) and throughput, respectively, compared to HeadKV, the state-of-the-art KV cache compression method. Moreover, FastKV successfully maintains accuracy on long-context benchmarks at levels comparable to the baselines. Our code is available at https://github.com/dongwonjo/FastKV.