FastKV : Compression de cache KV pour un traitement rapide de long contexte avec propagation sélective de jetons

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

February 3, 2025

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

cs.AI

Résumé

Alors que les grands modèles de langage (LLM) excellent dans la gestion de séquences à long contexte, ils nécessitent des caches substantiels clé-valeur (KV) pour stocker des informations contextuelles, ce qui peut lourdement impacter l'efficacité computationnelle et l'utilisation de la mémoire. Les efforts précédents pour compresser ces caches KV se sont principalement concentrés sur la réduction des besoins en mémoire mais étaient limités dans l'amélioration de la latence. Pour résoudre ce problème, nous introduisons FastKV, une méthode de compression de cache KV conçue pour améliorer la latence des séquences à long contexte. Pour améliorer les vitesses de traitement tout en maintenant la précision, FastKV adopte une approche novatrice de Propagation Sélective de Token (TSP) qui conserve l'ensemble des informations contextuelles dans les couches initiales des LLM et propage sélectivement seulement une partie de ces informations dans les couches plus profondes même à l'étape de préremplissage. De plus, FastKV intègre une compression de cache KV consciente de l'attention de requête groupée (GQA) pour exploiter les avantages de GQA en termes d'efficacité mémoire et computationnelle. Nos résultats expérimentaux montrent que FastKV obtient des améliorations de 2,00 fois et 1,40 fois respectivement en temps jusqu'au premier token (TTFT) et en débit par rapport à HeadKV, la méthode de compression de cache KV de pointe. De plus, FastKV maintient avec succès la précision sur des référentiels à long contexte à des niveaux comparables aux baselines. Notre code est disponible sur 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.