FastKV: Кэш KV сжатия для быстрой обработки длинного контекста с выборочным распространением токенов.

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

February 3, 2025

Авторы: Dongwon Jo, Jiwon Song, Yulhwa Kim, Jae-Joon Kim

cs.AI

Аннотация

Хотя большие языковые модели (LLM) отлично справляются с обработкой последовательностей длинного контекста, им требуются значительные кэши ключ-значение (KV) для хранения контекстуальной информации, что может серьезно обременять вычислительную эффективность и использование памяти. Предыдущие усилия по сжатию этих кэшей KV в основном сосредотачивались на снижении требований к памяти, но оказались ограничены в улучшении задержки. Для решения этой проблемы мы представляем FastKV, метод сжатия кэша KV, разработанный для улучшения задержки для последовательностей длинного контекста. Для повышения скорости обработки при сохранении точности FastKV принимает новый подход Token-Selective Propagation (TSP), который сохраняет всю контекстную информацию в начальных слоях LLM и выборочно передает только часть этой информации в более глубокие слои даже на этапе предварительного заполнения. Кроме того, FastKV включает в себя сжатие кэша KV, осведомленного о групповых запросах (GQA), для использования преимуществ GQA как в памяти, так и в вычислительной эффективности. Наши экспериментальные результаты показывают, что FastKV достигает улучшений во времени до первого токена (TTFT) и пропускной способности в 2,00 раза и 1,40 раза соответственно по сравнению с HeadKV, методом сжатия кэша KV последнего поколения. Более того, FastKV успешно сохраняет точность на длинных контекстных бенчмарках на уровне, сравнимом с базовыми показателями. Наш код доступен по адресу 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.