FastKV: Compresión de Caché KV para un Procesamiento Rápido de Contextos Largos con Propagación Selectiva 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

Resumen

Si bien los modelos de lenguaje grandes (LLMs) sobresalen en el manejo de secuencias de largo contexto, requieren un considerable almacenamiento en cachés clave-valor (KV) para guardar información contextual, lo cual puede sobrecargar la eficiencia computacional y el uso de memoria. Los esfuerzos previos para comprimir estas cachés KV se centraron principalmente en reducir las demandas de memoria, pero estaban limitados en mejorar la latencia. Para abordar este problema, presentamos FastKV, un método de compresión de caché KV diseñado para mejorar la latencia en secuencias de largo contexto. Para mejorar las velocidades de procesamiento manteniendo la precisión, FastKV adopta un enfoque novedoso de Propagación Selectiva de Tokens (TSP) que conserva la información de contexto completa en las capas iniciales de LLMs y propaga selectivamente solo una parte de esta información en capas más profundas incluso en la etapa de precarga. Además, FastKV incorpora compresión de caché KV consciente de la atención de consultas agrupadas (GQA) para explotar las ventajas de GQA en eficiencia tanto de memoria como computacional. Nuestros resultados experimentales muestran que FastKV logra mejoras de 2.00 veces y 1.40 veces en tiempo hasta el primer token (TTFT) y rendimiento, respectivamente, en comparación con HeadKV, el método de compresión de caché KV de vanguardia. Además, FastKV mantiene con éxito la precisión en benchmarks de largo contexto a niveles comparables con los baselines. Nuestro código está disponible en 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.