ParisKV: Fast and Drift-Robust KV-Cache Retrieval for Long-Context LLMs

Riepilogo Tecnico: ParisKV

Problematica

I Large Language Models (LLM) stanno rapidamente estendendo le loro finestre di contesto a milioni di token, creando una pressione senza precedenti sull'efficienza dell'inferenza. La KV-cache, che memorizza chiavi e valori per la decodifica autoregressiva, domina rapidamente l'impronta di memoria e la latenza man mano che il contesto cresce. La decodifica a lungo contesto è intrinsecamente limitata dalla memoria (memory-bound), poiché ogni passaggio richiede la lettura di un grande volume di vettori KV, con costi di banda che scalano linearmente con la lunghezza del contesto.

I metodi di attenzione sparsa/selettiva tentano di mitigare questo problema prestando attenzione solo a un sottoinsieme Top-k di token passati rilevanti, ma gli approcci esistenti al recupero della KV-cache affrontano tre sfide critiche negli scenari di generazione lunga e sensibili alla latenza:

1. Compromesso Velocità–Qualità: Il recupero approssimativo leggero (ad esempio, clustering grossolano, quantizzazione a basso bit) spesso sacrifica il richiamo (recall) per la velocità. Recuperare l'accuratezza richiede tipicamente di aumentare il budget di recupero, il che erode i benefici della sparsità.
2. Drift della Decodifica (Decoding Drift): I centroidi appresi dal clustering delle chiavi di prefill diventano obsoleti durante la generazione. Man mano che le chiavi di decodifica si accumulano, la distribuzione cambia, causando un disallineamento tra i centroidi di prefill e la vera distribuzione delle chiavi. Ciò porta a un grave collasso del richiamo durante i lunghi passi di decodifica nei metodi precedenti.
3. Recupero Vincolato dalla CPU (CPU-Bound Retrieval): Quando l'intera KV-cache viene scaricata (offloaded) nella memoria della CPU per risparmiare spazio sulla GPU, il recupero spesso comporta la ricerca lato CPU e il trasferimento dati da CPU a GPU. Ciò introduce errori di approssimazione e rende l'orchestrazione della CPU il fattore dominante nella latenza end-to-end.

Metodologia: Il Framework ParisKV

ParisKV è un algoritmo–sistema co-design consapevole dell'hardware per il recupero della KV-cache che raggiunge una decodifica veloce e un richiamo stabile sotto input lunghi e/o generazione lunga (fino alla scala di un milione di token), anche in presenza di drift della distribuzione.

Innovazioni Algoritmiche Core

1. Centroidi Analitici Robusti al Drift:
A differenza dei metodi precedenti che apprendono i centroidi dalle chiavi di prefill, ParisKV trasforma query e chiavi in uno spazio stabile e indipendente dai dati.

• Normalizzazione e Rotazione: Le chiavi e le query sono normalizzate $\ell_2$ e sottoposte a una rotazione ortogonale casuale condivisa (implementata tramite SRHT). Questo le mappa su un ipersfera unitaria.
• Centroidi Analitici: Invece di apprendere i centroidi, ParisKV definisce un insieme di centroidi distribuiti uniformemente sull'ipersfera unitaria basandosi su pattern di segno (ad esempio, $\pm 1/\sqrt{m}$). Poiché questi centroidi sono indipendenti dai dati e distribuiti uniformemente, qualsiasi nuova chiave generata durante la decodifica sarà sempre vicina ad almeno un centroide, prevenendo il problema del "centroide obsoleto".

2. Pipeline di Recupero Nativa su GPU in Due Stadi:
ParisKV implementa una pipeline di recupero da grossolano a raffinato interamente sulla GPU per evitare i colli di bottiglia della CPU.

• Generazione di Candidati Grossolana (Basata su Collisioni): Le chiavi ruotate vengono suddivise in $B$ sottospazi. Ogni direzione di sottospazio è assegnata al centroide analitico più vicino. Per una query, il sistema seleziona i top-$\rho$ centroidi più simili in ogni sottospazio. Le chiavi in questi bucket selezionati ricevono voti (con pesi a livelli). Il sistema somma i voti attraverso i sottospazi per generare un punteggio intero grossolano, potando lo spazio di ricerca a un piccolo set di candidati (ad esempio, top 5–10%).
• Reranking dei Candidati (RSQ-IP): I candidati vengono riordinati utilizzando un estimatore del prodotto interno calibrato senza accedere alle chiavi a piena precisione.
  • Quantizzazione: Ogni direzione di sottospazio è codificata usando un codice a 4 bit (1 bit di segno + 3 bit di magnitudo) derivato dai priori Beta indotti dalla rotazione.
  • Calibrazione: Un fattore di scala leggero ($w_{i,b}$) è precomputato durante la fase di prefill per correggere gli errori di allineamento della quantizzazione.
  • Stima: Il sistema stima il punteggio di attenzione grezzo $\langle k, q \rangle$ usando la formula: $\|q\|_2 \sum_{b=1}^B w_{i,b} \langle v_{i,b}, \tilde{q}_b \rangle$, dove $v_{i,b}$ è la direzione dequantizzata e $\tilde{q}_b$ è la query ruotata. Ciò consente un reranking ad alta fedeltà utilizzando solo metadati compatti.

Innovazioni a Livello di Sistema

1. Offloading tramite Indirizzamento Virtuale Unificato (UVA):
ParisKV scarica l'intera KV-cache a piena precisione nella memoria della CPU, ma sfrutta l'UVA per consentire ai kernel della GPU di recuperare direttamente solo le coppie KV Top-k selezionate su richiesta. Questo evita operazioni esplicite di memcpy e l'overhead di scheduling della CPU, permettendo alla GPU di accedere alla memoria della CPU come se fosse locale.

2. Meccanismo di Aggiornamento a Finestra Scorrevole (Sliding-Window):
Per supportare una generazione illimitata, la KV-cache è gestita in quattro regioni: Sink (token iniziali su GPU), Local (token recenti su GPU), Retrieval (token storici scaricati) e Update Buffer. Quando il buffer si riempie, i token locali più vecchi vengono espulsi nella zona di Retrieval (scaricati sulla CPU), e i token del buffer vengono promossi a Local. Ciò assicura che i metadati di recupero rimangano freschi preservando al contempo l'accesso rapido sia al contesto recente che alla storia a lungo raggio.

Contributi Chiave

1. Centroidi Analitici Robusti al Drift: Una trasformazione indipendente dai dati (normalizzazione + rotazione) che mantiene un richiamo di recupero stabile anche sotto significativi shift di distribuzione alla scala di un milione di token, eliminando la necessità di centroidi appresi che degradano nel tempo.
2. Pipeline Coarse-to-Fine Nativa su GPU: Un sistema di recupero che combina il pruning basato su collisioni e il reranking calibrato (RSQ-IP) che ottiene una selezione Top-k ad alta fedeltà direttamente su codici compressi senza ricerca lato CPU.
3. Offloading Scalabile Basato su UVA: Un'implementazione che consente il fetching on-demand della KV-cache con un intervento minimo della CPU, permettendo al sistema di scalare a contesti di un milione di token dove l'attenzione completa esaurisce la memoria.
4. Valutazione Estesa: Prestazioni dimostrate su benchmark di ragionamento a lunga generazione (GPQA-Diamond, MATH500, AIME25) e comprensione a lungo contesto (LongBench-V2, RULER).

Risultati Sperimentali

Gli autori hanno valutato ParisKV su benchmark di ragionamento a lunga generazione (GPQA-Diamond, MATH500, AIME25) e comprensione a lungo contesto (LongBench-V2, RULER) su molteplici famiglie di modelli (Qwen-3-8B, DeepSeek-R1-Llama-8B, ecc.).

Accuratezza:

• ParisKV eguaglia o supera la qualità della full attention su benchmark di lungo input e lunga generazione.
• Su AIME25, migliora l'accuratezza pass@8 di +40 a +77 punti rispetto ai baseline allo stato dell'arte MagicPIG e PQCache.
• In 7 casi su 9, ParisKV eguaglia o supera l'accuratezza della full attention, superando significativamente i precedenti baseline di recupero.

Efficienza:

• Throughput: ParisKV raggiunge fino a 2.8× un throughput maggiore rispetto alla full attention nel suo range operativo.
• Latenza alla Scala: Alla scala di un milione di token (1M token), ParisKV riduce la latenza di decodifica di 17× rispetto a MagicPiga e di 44× rispetto a PQCache.
• Scalabilità: Mentre la full attention va in Out-of-Memory (OOM) con batch size $\ge 4$ per contesti da 128K e $\ge 2$ per 256K, ParisKV scala rispettivamente a batch size di 8 e 5. Rimane eseguibile a 384K e 1024K di contesto dove la full attention fallisce completamente.

Significato e Rivendicazioni

Il paper sostiene che ParisKV affronti le sfide fondamentali del drift della distribuzione e della latenza di recupero nell'inferenza LLM a lungo contesto. Trasformando le chiavi su un ipersfera unitaria e impiegando centroidi analitici indipendenti dai dati, il sistema elimina l'obsolescenza dei centroidi che affligge i metodi esistenti.

Gli autori affermano che un recupero della KV-cache attentamente progettato può migliorare simultaneamente sia l'efficienza che l'accuratezza per l'inferenza a lungo contesto, anche in presenza di drift. Ciò apre nuove possibilità per il deployment di LLM in applicazioni di inferenza a latenza sensibile (come i compiti di ragionamento complesso) dove i metodi precedenti o non riuscivano a scalare o subivano un grave degrado dell'accuratezza. Il lavoro dimostra che è possibile ottenere una qualità quasi lossless mantenendo i benefici di velocità dell'attenzione sparsa, anche quando la KV-cache è scaricata nella memoria della CPU.