Technische Samenvatting: ParisKV

Probleemstelling

Large Language Models (LLM's) breiden hun contextvensters snel uit naar miljoenen tokens, wat ongekende druk legt op de efficiëntie van inferentie. De KV-cache, die keys en values opslaat voor autoregressieve decodering, domineert het geheugenvoetafdruk en de latentie snel naarmate de context groeit. Decodering met lange context is inherent geheugenbeperkt (memory-bound), aangezien elke stap een grote hoeveelheid KV-vectoren moet lezen, waarbij de bandbreedtekosten lineair schalen met de contextlengte.

Methoden voor ijle (sparse) of selectieve aandacht proberen dit te verzachten door alleen aandacht te besteden aan een Top-k subset van relevante eerdere tokens, maar bestaande benaderingen voor KV-cache retrieval kampen met drie kritieke uitdagingen in latency-gevoelige scenario's met lange generaties:

1. Snelheid–Kwaliteit Tradeoff: Lichtgewicht benaderende retrieval (bijv. grove clustering, lage-bit kwantisatie) offert vaak recall op voor snelheid. Het herstellen van nauwkeurigheid vereist doorgaans een groter retrieval-budget, wat de voordelen van ijle aandacht (sparsity) tenietdoet.
2. Decoding Drift: Centroids die geleerd zijn van clustering tijdens de prefill-fase worden verouderd (stale) naarmate de generatie voortduurt. Naarmate decodings-keys accumuleren, verschuift de distributie, wat zorgt voor een mismatch tussen de prefill-only centroids en de werkelijke key-distributie. Dit leidt tot een ernstige instorting van de recall tijdens lange decodingsstappen bij eerdere methoden.
3. CPU-gebonden Retrieval: Wanneer volledige KV-caches naar het CPU-geheugen worden verplaatst om GPU-ruimte te besparen, omvat retrieval vaak CPU-zijde zoekopdrachten en CPU-naar-GPU datatransfers. Dit introduceert benaderingsfouten en maakt CPU-orkestratie de dominante factor in de end-to-end latentie.

Methodologie: Het ParisKV Framework

ParisKV is een hardware-bewuste algoritme–systeem co-design voor KV-cache retrieval die snelle decodering en stabiele recall bereikt onder lange input en/of lange generatie (tot op de schaal van een miljoen tokens), zelfs in de aanwezigheid van distributie-drift.

Kerninnovaties in het Algoritme

1. Drift-Robuuste Analytische Centroids:
In tegenstelling tot eerdere methoden die centroids leren van prefill-keys, transformeert ParisKV queries en keys naar een stabiele, data-onafhankelijke ruimte.

• Normalisatie en Rotatie: Keys en queries worden $\ell_2$-genormaliseerd en onderworpen aan een gedeelde willekeurige orthogonale rotatie (geïmplementeerd via SRHT). Dit brengt ze in kaart op een eenheids-hypersfeer.
• Analytische Centroids: In plaats van centroids te leren, definieert ParisKV een set uniform verdeelde centroids op de eenheids-hypersfeer gebaseerd op tekenpatronen (bijv. $\pm 1/\sqrt{m}$). Omdat deze centroids data-onafhankelijk en uniform verdeeld zijn, zal elke nieuw gegenereerde key tijdens decodering altijd dicht bij ten minste één centroid liggen, wat het probleem van "verouderde centroids" voorkomt.

2. Tweestaps GPU-Native Retrieval Pipeline:
ParisKV implementeert een coarse-to-fine retrieval pipeline volledig op de GPU om CPU-bottlenecks te vermijden.

• Grove Candidatengeneratie (Collision-Based): De geroteerde keys worden gesplitst in $B$ subspaces. Elke subspace-richting wordt toegewezen aan de dichtstbijzijnde analytische centroid. Voor een query selecteert het systeem de top-$\rho$ meest gelijkaande centroids in elke subspace. Keys in deze geselecteerde buckets ontvangen stemmen (met getrapte gewichten). Het systeem telt de stemmen over de subspaces op om een grove integer score te genereren, waardoor de zoekruimte wordt teruggebracht tot een kleine kandidaatverzameling (bijv. top 5–10%).
• Kandidaat Reranking (RSQ-IP): De kandidaten worden geherrankt met behulp van een gekalibreerde inwendig product-schatter zonder toegang te hebben tot full-precision keys.
  • Kwantisatie: Elke subspace-richting wordt gecodeerd met een 4-bit code (1-bit teken + 3-bit magnitude) afgeleid van de door rotatie geïnduceerde Beta-priors.
  • Kalibratie: Een lichtgewicht schalingsfactor ($w_{i,b}$) wordt tijdens de prefill-fase vooraf berekend om kwantisatie-uitlijnfouten te corrigeren.
  • Estimatie: Het systeem schat de ruwe aandacht-score $\langle k, q \rangle$ met de formule: $\|q\|_2 \sum_{b=1}^B w_{i,b} \langle v_{i,b}, \tilde{q}_b \rangle$, waarbij $v_{i,b}$ de gedequantiseerde richting is en $\tilde{q}_b$ de geroteerde query. Dit maakt high-fidelity reranking mogelijk met enkel compacte metadata.

Systeem-niveau Innovaties

1. Unified Virtual Addressing (UVA) Offloading:
ParisKV verplaatst de full-precision KV-cache naar het CPU-geheugen, maar maakt gebruik van UVA om GPU-kernels direct de geselecteerde Top-k KV-paren op aanvraag te laten ophalen. Dit omzeilt expliciete memcpy-operaties en CPU-zijde scheduling overhead, waardoor de GPU toegang heeft tot het CPU-geheugen alsof het lokaal is.

2. Sliding-Window Update Mechanisme:
Om onbegrensde generatie te ondersteunen, wordt de KV-cache beheerd in vier regio's: Sink (vroege tokens op GPU), Local (recente tokens op GPU), Retrieval (offloaded historische tokens) en een Update Buffer. Wanneer de buffer vol raakt, worden de oudste lokale tokens verwijderd naar de Retrieval-zone (offloaded naar CPU), en de buffer-tokens worden gepromoveerd naar Local. Dit zorgt ervoor dat de retrieval-metadata vers blijft terwijl snelle toegang tot zowel recente context als lange-afstandsgeschiedenis behouden blijft.

Belangrijkste Bijdragen

1. Drift-Robuuste Analytische Centroids: Een data-onafhankelijke transformatie (normalisatie + rotatie) die een stabiele retrieval recall behoudt, zelfs onder significante distributieverschuivingen op de schaal van een miljoen tokens, waardoor de noodzaak voor geleerde centroids die over tijd degraderen, wordt geëlimineerd.
2. GPU-Native Coarse-to-Fine Pipeline: Een retrieval-systeem dat collision-based pruning combineert met gekalibreerde reranking (RSQ-IP), wat direct op gecomprimeerde codes een high-fidelity Top-k selectie bereikt zonder CPU-zijde zoekopdrachten.
3. Schaalbare UVA-gebaseerde Offloading: Een implementatie die on-demand KV-fetching mogelijk maakt met minimale CPU-interventie, waardoor het systeem kan schalen naar contexten van een miljoen tokens waar full attention tekortschiet qua geheugen.
4. Uitgebreide Evaluatie: Gedemonstreerde prestaties over long-generation reasoning benchmarks (GPQA-Diamond, MATH500, AIME25) en long-context understanding (LongBench-V2, RULER).

Experimentele Resultaten

De auteurs evalueerden ParisKV op long-generation reasoning benchmarks (GPQA-Diamond, MATH500, AIME25) en long-context understanding (LongBench-V2, RULER) over meerdere modelfamilies (Qwen-3-8B, DeepSeek-R1-Llama-8B, etc.).

Nauwkeurigheid:

• ParisKV evenaart of overtreft de kwaliteit van full attention op long-input en long-generation benchmarks.
• Op AIME25 verbetert het de pass@8 nauwkeurigheid met +40 tot +77 punten ten opzichte van de state-of-the-art baselines MagicPIG en PQCache.
• In 7 van de 9 settings evenaart of overtreft ParisKV de nauwkeurigheid van full attention, wat aanzienlijk beter is dan eerdere retrieval baselines.

Efficiëntie:

• Throughput: ParisKV bereikt tot 2,8× hogere throughput dan full attention binnen zijn bruikbare bereik.
• Latency op Schaal: Op de schaal van een miljoen tokens (1M tokens) vermindert ParisKV de decode-latentie met 17× vergeleken met MagicPIG en met 44× vergeleken met PQCache.
• Schaalbaarheid: Terwijl full attention out-of-memory (OOM) gaat bij batch sizes $\ge 4$ voor 128K contexten en $\ge 2$ voor 256K contexten, schaalt ParisKV naar batch sizes van respectievelijk 8 en 5. Het blijft bruikbaar bij 384K en 1024K contexten waar full attention volledig faalt.

Betekenis en Claims

Het artikel claimt dat ParisKV de fundamentele uitdagingen van distributie-drift en retrieval-latentie in long-context LLM-inferentie aanpakt. Door de keys op een eenheids-hypersfeer te transformeren en analytische, data-onafhankelijke centroids te gebruiken, elimineert het systeem de centroid-veroudering die bestaande methoden teistert.

De auteurs stellen dat zorgvuldig ontworpen KV-cache retrieval simultaan zowel de efficiëntie als de nauwkeurigheid kan verbeteren voor long-context inferentie, zelfs in de aanwezigheid van drift. Dit opent nieuwe mogelijkheden voor het inzetten van LLM's in latency-gevoelige, lang-genererende toepassingen (zoals complexe redeneertaken) waar eerdere methoden ofwel niet konden schalen of ernstige nauwkeurigheidsdegradatie leden. Het werk demonstreert dat het mogelijk is om bijna verliesloze kwaliteit te bereiken terwijl de snelheidvoordelen van sparse attention behouden blijven, zelfs wanneer de KV-cache naar de CPU is verplaatst.