Technische Samenvatting: RDKV - Rate-Distortion Bit-toewijzing voor Gecombineerde Evictie en Quantisatie van de KV-cache

1. Probleemstelling

Grote Taalmodellen (LLM's) ondervinden tijdens inferentie met lange contexten een kritieke knelpunt: geheugenbandbreedte en -capaciteit. De grootte van de Key-Value (KV)-cache groeit lineair met de sequentielengte, waardoor bij elke decodestap tientallen gigabytes aan data opnieuw moeten worden gelezen van off-chip High-Bandwidth Memory (HBM) naar on-chip geheugen. Dit maakt inferentie geheugen-gebonden in plaats van rekenkracht-gebonden.

Bestaande compressiestrategieën vallen over het algemeen in twee geïsoleerde categorieën:

1. Evictie: Tokens of kanalen scoren en de laagst gerangschikte volledig verwijderen (binair behouden-of-verwijderen). Dit laat geen precisieniveau over voor tokens van gemiddeld belang.
2. Quantisatie: Alle tokens behouden maar hun bit-breedte verkleinen. Hierdoor blijven zelfs verwaarloosbare tokens bewaard, wat geheugen verspillen.

Hoewel recente hybride methoden proberen deze te combineren, vertrouwen ze vaak op heuristische routeer-scores of vooraf toegewezen verhoudingen bepaald via offline zoekopdrachten, en behandelen ze evictie en quantisatie niet als een geünificeerd optimalisatieprobleem.

2. Methodologie: RDKV

Het artikel stelt RDKV (Rate-Distortion KV-cache compressie) voor, een kader dat KV-cache compressie formuleert als een rate-distortion-probleem. In dit perspectief zijn evictie en quantisatie geen aparte acties, maar twee eindpunten van één enkele bit-toewijzingscurve.

2.1 Theoretische Formulering

RDKV leidt een vervormingsgewicht af voor elke cache-eenheid (token in de V-cache, kanaal in de K-cache) op basis van de verstoring die door het verwijderen of quantiseren ervan wordt veroorzaakt op de attentieberekening.

• V-cache (Tokens): Het vervormingsgewicht $w_t$ voor een token is de totale variatieafstand van de attentieverdeling na evictie. Het wordt berekend als de som van de attentiescores die het token ontvangt over alle query-posities: $w_t = \sum_{\tau} a_{\tau,t}$.
• K-cache (Kanalen): Het vervormingsgewicht $w_c$ voor een kanaal is de spectrale norm van de afwijking in logit veroorzaakt door het op nul zetten van dat kanaal. Het wordt berekend als het product van de normen van de corresponderende query- en key-kolommen: $w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$.

Deze gewichten dienen als kostencoëfficiënten in een rate-distortion-doelfunctie:
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{onder voorwaarde} \quad \sum_{u} b_u \leq B$$
waarbij $b_u$ de toegewezen bit-breedte is voor eenheid $u$, en $B$ het totale bit-budget.

2.2 Bit-toewijzingspijplijn

De toewijzing wordt opgelost via reverse water-filling (Lagrangiaanse relaxatie), aangepast voor discrete, door hardware ondersteunde bit-breedtes $\mathcal{B} = \{0, 2, 4, 8, 16\}$.

1. Fase 1 (Gewichten): Bereken token-gewichten ($w_t$) en kanaal-gewichten ($w_c$) één keer na de prefill-fase met behulp van de volledige, niet-gecomprimeerde cache.
2. Fase 2 (Token-toewijzing): Los een Multiple-Choice Knapsack-probleem (MCKP) op voor de V-cache. Tokens krijgen bit-breedtes toegewezen uit $\{0, 2, 4, 8, 16\}$. Een bit-breedte van 0 komt overeen met evictie.
3. Fase 3 (Kanaaltoewijzing): Los het MCKP op voor de K-cache, maar alleen voor de tokens die in Fase 2 behouden zijn. Kanalen krijgen onafhankelijk van de V-bit-breedte van het token bit-breedtes toegewezen, waardoor een enkel token gemengde precisie-K- en V-componenten kan hebben.

2.3 Efficiënte Gepakte Decode (TriZone)

Om daadwerkelijke geheugen- en latentiewinst te realiseren, introduceert RDKV een TriZone-cache-indeling die dequantisatie fuseert in de attentie-kernel, waardoor de noodzaak om data te unpacken naar FP16 voordat verwerking plaatsvindt, wordt vermeden.

• Zone A: Een gepakte oude-cache-stroom met behouden K-rijen en V-rijen met bit-breedtes $\{2, 4, 8\}$. Deze zijn gegroepeerd op bit-breedte om uniforme dequantisatiepaden mogelijk te maken.
• Zone B: FP16 V-rijen voor tokens die 16-bits precisie toegewezen hebben gekregen.
• Zone C: FP16 K- en V-rijen voor nieuw gegenereerde tokens.

De attentie-kernel berekent logits en outputs door deze drie zones te combineren, waarbij dequantisatie direct wordt gefuseerd in de matrixvermenigvuldiging, wat extra overhead voor geheugenbandbreedte elimineert.

3. Belangrijkste Bijdragen

1. Rate-Distortion Formulering: Het artikel presenteert KV-cache compressie als een geünificeerd rate-distortion-probleem, en bewijst dat de optimale oplossing inherent een mix is van zero-rate (evictie) en finite-rate (quantisatie) eenheden. Dit elimineert de behoefte aan heuristische verhoudingen of offline zoekopdrachten.
2. Discrete Toewijzing via Reverse Water-Filling: De auteurs concretiseren de toewijzing als een discreet knapsack-probleem dat oplosbaar is via Lagrangiaanse relaxatie, en wijzen optimale bit-breedtes toe uit $\{0, 2, 4, 8, 16\}$ op basis van vervormingsgewichten die zijn afgeleid uit attentiedynamica.
3. TriZone Gepakte Indeling: Een hardware-bewuste implementatie die dequantisatie fuseert in de attentie-kernel, waardoor de theoretische bit-breedtebesparingen worden omgezet in daadwerkelijke reducties van piekgeheugen en decode-latentie.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op LLaMA-3.1-8B-Instruct, Mistral-7B, Qwen3-4B, LLaMA-2-13B en Qwen2.5-72B over LongBench, RULER en InfiniteBench.

• Nauwkeurigheid: RDKV presteert beter dan de best beoordeelde baseline (Ada-SnapKV) met gemiddeld 9,1%. Op LongBench herstelt het 97,81% van de nauwkeurigheid van de volledige cache met slechts 2,48% cachebehoud.
• Prestaties bij Lange Context:
  • RULER: Bij een context van 128K bereikt RDKV 66,95% nauwkeurigheid vergeleken met 58,43% voor de sterkste baseline (AdaKV), een verschil van 8,5 punten.
  • InfiniteBench: RDKV staat eerste onder compressiemethoden met een gemiddelde score van 39,46, en overtreft AdaKV met 1,40 punten.
  • Needle-in-a-Haystack: RDKV behoudt een bijna uniforme ophaalnauwkeurigheid (gemiddeld 0,99) over dieptes, terwijl baselines die alleen evictie toepassen significant falen bij intermediaire dieptes.
• Efficiëntie:
  • Latentie: Met een context van 128K bereikt RDKV een 4,5× versnelling in decode vergeleken met decode van volledige cache met FlashAttention-2.
  • Geheugen: Het bereikt een 1,9× reductie in piekgeheugen bij een context van 128K en kan succesvol een context van 256K draaien op een enkele A100 64GB GPU, waar FullKV uit het geheugen loopt (OOM).

5. Betekenis en Claims

Het artikel stelt dat het gezamenlijk behandelen van evictie en quantisatie binnen één optimalisatiekader superieur is aan gefaseerde of op verhoudingen gebaseerde benaderingen. Door toewijzingsscores direct af te leiden uit de vervorming die wordt veroorzaakt op de attentieberekening, vermijdt RDKV de "binaire valkuil" van methoden die alleen evictie toepassen (die tokens van gemiddeld belang verwijderen) en de "verspilling" van methoden die alleen quantisatie toepassen (die verwaarloosbare tokens behouden).

De auteurs benadrukken dat hun methode:

• Geen modelhertraining vereist.
• Architectuur-agnostisch is en consistent presteert over verschillende modelgroottes en families.
• Een Pareto-optimale afweging biedt tussen latentie en nauwkeurigheid, en baselines strikt overtreft op beide metrieken bij lange contextlengtes.
• Aantoont dat het "continue optimum" van rate-distortion-theorie, wanneer gediscretiseerd voor hardware, van nature een strategie met gemengde bit-breedtes oplevert die huidige state-of-the-art hybriden overtreft.

Het artikel concludeert dat deze gezamenlijke aanpak de nauwkeurigheid verbetert zonder efficiëntie op te offeren, en inferentie met lange contexten mogelijk maakt op hardware met beperkte middelen.