Riepilogo Tecnico: RDKV - Allocazione dei Bit per Distorsione-Rate per l'Evizione e la Quantizzazione Congiunta della KV Cache

1. Enunciato del Problema

I Large Language Models (LLM) affrontano un collo di bottiglia critico durante l'inferenza con contesti lunghi: la larghezza di banda e la capacità della memoria. La dimensione della cache Key-Value (KV) cresce linearmente con la lunghezza della sequenza, richiedendo la rilettura di decine di gigabyte di dati dalla memoria off-chip ad alta larghezza di banda (HBM) alla memoria on-chip ad ogni passo di decodifica. Ciò rende l'inferenza limitata dalla memoria piuttosto che dal calcolo.

Le strategie di compressione esistenti rientrano generalmente in due categorie isolate:

1. Evizione: Valutazione dei token o dei canali e scarto completo di quelli con il punteggio più basso (mantenimento o eliminazione binario). Questo non lascia alcun livello di precisione per i token di importanza moderata.
2. Quantizzazione: Conservazione di tutti i token ma riduzione della loro larghezza in bit. Questo mantiene anche i token trascurabili, sprecando memoria.

Sebbene recenti ibridi tentino di combinare queste due strategie, spesso si basano su punteggi di instradamento euristici o rapporti pre-assegnati determinati tramite ricerca offline, fallendo nel trattare l'evizione e la quantizzazione come un problema di ottimizzazione unificato.

2. Metodologia: RDKV

Il documento propone RDKV (Compressione della KV cache basata su Rate-Distortion), un framework che formula la compressione della KV cache come un problema rate-distortion. In questa visione, l'evizione e la quantizzazione non sono azioni separate ma due estremi di una singola curva di allocazione dei bit.

2.1 Formulazione Teorica

RDKV deriva un peso di distorsione per ogni unità della cache (token nella cache V, canale nella cache K) basato sulla perturbazione indotta dalla sua rimozione o quantizzazione sul calcolo dell'attenzione.

• Cache V (Token): Il peso di distorsione $w_t$ per un token è la distanza di variazione totale della distribuzione di attenzione dopo l'evizione. È calcolato come la somma dei punteggi di attenzione ricevuti dal token attraverso tutte le posizioni di query: $w_t = \sum_{\tau} a_{\tau,t}$.
• Cache K (Canali): Il peso di distorsione $w_c$ per un canale è la norma spettrale della deviazione dei logit causata dallo zeroing di quel canale. È calcolato come il prodotto delle norme delle colonne corrispondenti di query e key: $w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$.

Questi pesi fungono da coefficienti di costo in una funzione obiettivo rate-distortion:
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{s.t.} \quad \sum_{u} b_u \leq B$$
dove $b_u$ è la larghezza in bit assegnata all'unità $u$, e $B$ è il budget totale di bit.

2.2 Pipeline di Allocazione dei Bit

L'allocazione è risolta tramite reverse water-filling (rilassamento lagrangiano) adattato per larghezze in bit discrete supportate dall'hardware $\mathcal{B} = \{0, 2, 4, 8, 16\}$.

1. Fase 1 (Ponderazione): Calcolo dei pesi dei token ($w_t$) e dei pesi dei canali ($w_c$) una sola volta dopo la fase di prefill utilizzando la cache completa non compressa.
2. Fase 2 (Allocazione dei Token): Risoluzione di un Problema dello Zaino a Scelta Multipla (MCKP) per la cache V. Ai token vengono assegnate larghezze in bit da $\{0, 2, 4, 8, 16\}$. Una larghezza in bit di 0 corrisponde all'evizione.
3. Fase 3 (Allocazione dei Canali): Risoluzione dell'MCKP per la cache K, ma solo per i token mantenuti nella Fase 2. Ai canali vengono assegnate larghezze in bit indipendentemente dalla larghezza in bit V del token, consentendo a un singolo token di avere componenti K e V a precisione mista.

2.2 Decodifica Imballata Efficiente (TriZone)

Per realizzare effettivi risparmi di memoria e latenza, RDKV introduce un layout della cache TriZone che fonde la dequantizzazione nel kernel di attenzione, evitando la necessità di decomprimere i dati in FP16 prima dell'elaborazione.

• Zona A: Flusso di vecchia cache imballato contenente le righe K e le righe V mantenute con larghezze in bit $\{2, 4, 8\}$. Queste sono raggruppate per larghezza in bit per consentire percorsi di dequantizzazione uniformi.
• Zona B: Righe V in FP16 per i token assegnati a precisione a 16 bit.
• Zona C: Righe K e V in FP16 per i token appena generati.

Il kernel di attenzione calcola i logit e gli output combinando queste tre zone, con la dequantizzazione fusa direttamente nella moltiplicazione di matrici, eliminando l'overhead aggiuntivo di larghezza di banda della memoria.

3. Contributi Chiave

1. Formulazione Rate-Distortion: Il documento presenta la compressione della KV cache come un problema rate-distortion unificato, dimostrando che la soluzione ottimale mescola intrinsecamente unità a rate zero (evizione) e unità a rate finito (quantizzazione). Questo elimina la necessità di rapporti euristici o ricerche offline.
2. Allocazione Discreta tramite Reverse Water-Filling: Gli autori istanziano l'allocazione come un problema dello zaino discreto risolvibile tramite rilassamento lagrangiano, assegnando le larghezze in bit ottimali da $\{0, 2, 4, 8, 16\}$ in base ai pesi di distorsione derivati dalle dinamiche di attenzione.
3. Layout Imballato TriZone: Un'implementazione consapevole dell'hardware che fonde la dequantizzazione nel kernel di attenzione, convertendo i risparmi teorici di larghezza in bit in riduzioni effettive della memoria di picco e della latenza di decodifica.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su LLaMA-3.1-8B-Instruct, Mistral-7B, Qwen3-4B, LLaMA-2-13B e Qwen2.5-72B attraverso LongBench, RULER e InfiniteBench.

• Accuratezza: RDKV supera la migliore baseline valutata (Ada-SnapKV) del 9,1% in media. Su LongBench, recupera il 97,81% dell'accuratezza della cache completa con solo il 2,48% di ritenzione della cache.
• Prestazioni con Contesti Lunghi:
  • RULER: A un contesto di 128K, RDKV raggiunge un'accuratezza del 66,95% rispetto al 58,43% della baseline più forte (AdaKV), con un divario di 8,5 punti.
  • InfiniteBench: RDKV si classifica primo tra i metodi di compressione con un punteggio medio di 39,46, superando AdaKV di 1,40 punti.
  • Needle-in-a-Haystack: RDKV mantiene un'accuratezza di recupero quasi uniforme (media 0,99) attraverso le profondità, mentre le baselines basate solo sull'evizione falliscono significativamente a profondità intermedie.
• Efficienza:
  • Latenza: Con un contesto di 128K, RDKV raggiunge un speedup di decodifica di 4,5× rispetto alla decodifica FlashAttention-2 con cache completa.
  • Memoria: Raggiunge una riduzione della memoria di picco di 1,9× a un contesto di 128K ed esegue con successo un contesto di 256K su una singola GPU A100 da 64GB, dove FullKV va incontro a esaurimento della memoria (OOM).

5. Significato e Affermazioni

Il documento afferma che trattare l'evizione e la quantizzazione congiuntamente all'interno di un singolo framework di ottimizzazione è superiore agli approcci a stadi o basati su rapporti. Derivando i punteggi di allocazione direttamente dalla distorsione indotta sul calcolo dell'attenzione, RDKV evita la "trappola binaria" dei metodi basati solo sull'evizione (che scartano i token moderati) e lo "spreco" dei metodi basati solo sulla quantizzazione (che mantengono i token trascurabili).

Gli autori sottolineano che il loro metodo:

• Non richiede nessun riaddestramento del modello.
• È agnostico rispetto all'architettura, performando in modo coerente attraverso diverse dimensioni e famiglie di modelli.
• Fornisce un compromesso Pareto-ottimale tra latenza e accuratezza, superando rigorosamente le baselines in entrambe le metriche a lunghezze di contesto elevate.
• Dimostra che l'"ottimo continuo" della teoria rate-distortion, quando discretizzato per l'hardware, produce naturalmente una strategia a bit misti che supera gli ibridi attuali allo stato dell'arte.

Il documento conclude che questo approccio congiunto migliora l'accuratezza senza sacrificare l'efficienza, abilitando l'inferenza con contesti lunghi su hardware con risorse limitate.