FastKV: Decoupling of Context Reduction and KV Cache Compression for Prefill-Decoding Acceleration

Probleemstelling

Grote Taalmodellen (LLMs) worden steeds bekwaamer in het verwerken van lange contexten (tot 128K tokens of meer), maar dit brengt aanzienlijke rekenkundige en geheugenkosten met zich mee in twee fasen:

1. Prefill-fase: De berekening van de aandacht (attention) schaalt kwadratisch met de invoerlengte, wat leidt tot hoge latentie bij het verwerken van lange prompts.
2. Decoding-fase: De Key-Value (KV) cache groeit lineair met de contextlengte. Bij het genereren van nieuwe tokens moet deze grote cache herhaaldelijk worden opgehaald, wat de doorvoer beperkt en het GPU-geheugen verzadigt.

Bestaande oplossingen kampen met een fundamenteel compromis:

• Decoding-only methoden (zoals StreamingLLM, SnapKV) comprimeren de KV-cache na het genereren, maar versnellen de prefille-fase niet.
• Prefill-aware methoden (zoals GemFilter, PyramidInfer) proberen de context tijdens de prefille te reduceren, maar koppelen dit te strak aan de KV-cache grootte. Dit betekent dat agressieve reductie tijdens de prefill (voor snelheid) direct leidt tot een kleiner KV-budget voor decoding, wat vaak resulteert in een significante daling van de nauwkeurigheid. Ze negeren vaak dat de "kritieke context" (de tokens waar het model op moet focussen) per laag verschilt.

Methodologie: FastKV

FastKV introduceert een raamwerk dat de contextreductie en de KV-cache compressie ontkoppelt. Het maakt gebruik van twee kernobservaties:

1. Laagafhankelijke dynamiek: In de vroege lagen van het model is de aandacht voor tokens onstabiel; kritieke tokens verschuiven sterk. In latere lagen stabiliseert deze selectie.
2. Sparce gebruik: Tijdens de decoding-fase gebruikt elk model slechts een klein subset van de tokens (ongeveer 0,38% voor een 128K context) voor de meeste aandachtswaarden.

Op basis hiervan introduceert FastKV twee technieken:

1. Token-Selective Propagation (TSP)

In plaats van direct alle tokens te filteren, voert FastKV een twee-staps prefill-strategie uit:

• Fase 1 (Vroege lagen): Het model verwerkt de volledige context tot een specifieke TSP-laag. Dit zorgt ervoor dat vroege lagen volledige contextinformatie kunnen opnemen voordat tokens worden verwijderd.
• Fase 2 (Latere lagen): Bij de TSP-laag wordt een subset van de meest informatieve tokens geselecteerd op basis van hun "saliency score" (gemiddelde aandachtswaarde over de attention heads). Alleen deze geselecteerde tokens (plus een vast venster van recente tokens) worden doorgegeven aan de daaropvolgende lagen.
• Resultaat: Dit vermindert de rekenlast in de latere lagen van de prefill-fase zonder de vroege lagen te beperken.

2. Ontkoppeling van KV-cache Retentie

FastKV breekt de koppeling tussen hoeveel context er tijdens de prefill wordt verwerkt en hoeveel KV-cache er voor de decoding wordt bewaard.

• Er worden twee onafhankelijke hyperparameters gebruikt: de TSP-snelheid (hoeveel tokens worden doorgegeven na de TSP-laag) en de KV-retentie-snelheid (hoeveel KV-entries per laag worden bewaard voor decoding).
• Elke laag comprimeert zijn eigen KV-cache onafhankelijk op basis van belang, ongeacht of de context al is gereduceerd door TSP. Dit stelt onderzoekers in staat om de prefill-snelheid en de decoding-geheugenefficiëntie apart te optimaliseren.

Belangrijkste Bijdragen

• Decoupling: Het eerste framework dat de reductie van de prefill-berekening en de grootte van de KV-cache voor decoding onafhankelijk van elkaar regelt, waardoor het compromis tussen snelheid en nauwkeurigheid wordt doorbroken.
• Laag-specifieke strategie: Het inzicht dat vroege lagen volledige context nodig hebben, terwijl latere lagen kunnen werken met een gereduceerde set, wat leidt tot de TSP-architectuur.
• Flexibiliteit: Het toestaan van een conservatieve TSP-snelheid (voor nauwkeurigheid) gecombineerd met een agressieve KV-retentie (voor geheugenbesparing), of vice versa.

Resultaten

De prestaties zijn getest op modellen zoals LLaMA-3.1-8B-Instruct en Ministral-8B-Instruct met contexten tot 128K tokens, vergeleken met baselines zoals StreamingLLM, SnapKV, GemFilter en PyramidInfer.

• Snelheid:
  • Prefill: Tot 1,82x versnelling ten opzichte van de volledige context-baseline.
  • Decoding: Tot 2,87x versnelling.
• Nauwkeurigheid:
  • FastKV behoudt de nauwkeurigheid van de volledige context-baseline (minder dan 1% daling op de LongBench-benchmark).
  • Het presteert aanzienlijk beter dan GemFilter en PyramidInfer, die bij vergelijkbare snelheidstijdingen grote nauwkeurigheidsdalingen vertonen (tot wel 11% daling bij GemFilter).
  • Op de "Needle-in-a-Haystack" test behaalt FastKV een score van 99,9% (tegenover 99,0% voor de volledige context), wat suggereert dat het filteren zelfs helpt bij het focussen op de meest relevante informatie.
• Efficiëntie: FastKV bereikt een betere afweging tussen snelheid en nauwkeurigheid dan alle bestaande methoden, waarbij het zowel de prefill- als decoding-fase versnelt zonder de GPU-geheugenvraag te laten exploderen.

Betekenis en Impact

FastKV biedt een schaalbare oplossing voor de groeiende contextvensters van moderne LLM's. Door de dynamiek van kritieke context over de lagen van het model te respecteren en de twee fasen van inferentie (prefill en decoding) onafhankelijk te optimaliseren, maakt het lang-context inferentie praktischer en kosteneffectiever. Dit is cruciaal voor toepassingen zoals documentanalyse, code-generatie en multi-document redenering, waar de huidige rekenkosten vaak een bottleneck vormen. De methode is compatibel met bestaande frameworks (zoals FlashAttention-2) en kan worden gecombineerd met andere optimalisatietechnieken zoals kwantisatie.