FastKV: Compressione della cache KV per un rapido processamento a lungo contesto con propagazione selettiva dei token.

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

February 3, 2025

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

cs.AI

Abstract

Mentre i grandi modelli di linguaggio (LLM) eccellono nel gestire sequenze di contesto lunghe, richiedono cache sostanziali di chiave-valore (KV) per memorizzare informazioni contestuali, il che può gravare pesantemente sull'efficienza computazionale e sull'utilizzo della memoria. Gli sforzi precedenti per comprimere queste cache KV si sono concentrati principalmente sulla riduzione delle richieste di memoria ma erano limitati nell'aumentare la latenza. Per affrontare questo problema, presentiamo FastKV, un metodo di compressione della cache KV progettato per migliorare la latenza per sequenze di contesto lunghe. Per migliorare le velocità di elaborazione pur mantenendo l'accuratezza, FastKV adotta un approccio innovativo di Propagazione Selettiva dei Token (TSP) che conserva le informazioni di contesto completo nei livelli iniziali dei LLM e propaga selettivamente solo una parte di queste informazioni nei livelli più profondi anche nella fase di pre-riempimento. Inoltre, FastKV incorpora la compressione della cache KV consapevole dell'attenzione delle query raggruppate (GQA) per sfruttare i vantaggi di GQA sia in termini di memoria che di efficienza computazionale. I nostri risultati sperimentali mostrano che FastKV ottiene miglioramenti di 2,00 volte e 1,40 volte nel tempo al primo token (TTFT) e nella capacità di elaborazione, rispettivamente, rispetto a HeadKV, il metodo di compressione della cache KV all'avanguardia. Inoltre, FastKV mantiene con successo l'accuratezza su benchmark di contesto lungo a livelli paragonabili alle baselines. Il nostro codice è disponibile su 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.