FastKV: KV-cachecompressie voor snelle verwerking van lange context met token-selectieve propagatie

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

Samenvatting

Hoewel grote taalmodellen (LLM's) uitblinken in het verwerken van lange-contextreeksen, hebben ze aanzienlijke key-value (KV) caches nodig om contextuele informatie op te slaan, wat de computationele efficiëntie en geheugenverbruik zwaar kan belasten. Eerdere inspanningen om deze KV-caches te comprimeren richtten zich voornamelijk op het verminderen van geheugenvereisten, maar waren beperkt in het verbeteren van latentie. Om dit probleem aan te pakken, introduceren we FastKV, een KV-cachecompressiemethode die is ontworpen om de latentie voor lange-contextreeksen te verbeteren. Om de verwerkingssnelheden te verbeteren met behoud van nauwkeurigheid, maakt FastKV gebruik van een nieuw Token-Selective Propagation (TSP) benadering die de volledige contextinformatie behoudt in de initiële lagen van LLM's en selectief slechts een deel van deze informatie doorgeeft in diepere lagen, zelfs in de prefill-fase. Daarnaast maakt FastKV gebruik van een KV-cachecompressie die zich bewust is van gegroepeerde-query-aandacht (GQA) om de voordelen van GQA te benutten op zowel geheugen- als computationeel gebied. Onze experimentele resultaten tonen aan dat FastKV verbeteringen van respectievelijk 2,00 keer en 1,40 keer behaalt in time-to-first-token (TTFT) en throughput in vergelijking met HeadKV, de state-of-the-art KV-cachecompressiemethode. Bovendien behoudt FastKV met succes de nauwkeurigheid op lange-context benchmarks op niveaus die vergelijkbaar zijn met de baselines. Onze code is beschikbaar op 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.