Riepilogo Tecnico: Tamarin (HSQ) – Compressione del KV-Cache Gerarchica Reversibile

1. Il Problema

Nell'inferenza di modelli linguistici di grandi dimensioni (LLM) autoregressivi, la cache Key-Value (KV) è il principale consumatore di memoria GPU, particolarmente negli scenari a contesto lungo. A differenza dei pesi del modello, che possono essere quantizzati una sola volta, la cache KV cresce linearmente con la lunghezza del contesto. Ad esempio, un modello Qwen3-8B in fp16 richiede circa 4,5 GiB di VRAM per un contesto di 32K token, superando spesso la capacità delle GPU standard.

Le soluzioni esistenti affrontano una dicotomia con un punto cieco condiviso:

• Quantizzazione KV: Comprime tutte le voci a bassa precisione (2–4 bit). Sebbene sia sicura, il rapporto di compressione è limitato (tipicamente ~4× rispetto a fp16) e può degradare la qualità a larghezze di bit inferiori.
• Evizione/Fusione dei Token: Metodi come StreamingLLM, H2O e SnapKV scartano o fondono permanentemente i token ritenuti non importanti durante la fase di prefill. Questi metodi ottengono elevati rapporti di compressione ma sono irreversibili. Se un token evitto precocemente diventa critico per una query generata in un passaggio di decoding successivo, l'informazione viene persa.

Il limite fondamentale dei lavori precedenti è che l'importanza del token è una funzione della query futura, che i metodi basati sull'evizione devono prevedere al momento della compressione senza il beneficio di quel contesto futuro.

2. Metodologia: Tamarin (HSQ)

Tamarin (Hierarchical Semantic Quantization, o HSQ) inquadra la compressione della cache KV come un problema di indicizzazione piuttosto che un problema di eliminazione. Mantiene tre rappresentazioni sincronizzate del contesto, garantendo che ogni token compresso rimanga indirizzabile e recuperabile.

L'Architettura a Tre Livelli

1. L1 (Livello Focale, GPU): Un piccolo insieme di token non riassunti memorizzati come vettori quantizzati a 4 bit. Questo include i sink di attenzione, una finestra recente opzionale e i token esplicitamente contrassegnati come importanti da un lettore appreso.
2. L2 (Livello di Riassunto, GPU): Gli altri token sono partizionati in gruppi (dimensione $N \in \{8, 16\}$). Ogni gruppo è compresso in due vettori rappresentativi appresi ($\hat{k}, \hat{v}$) memorizzati a precisione a 3 bit. Fondamentalmente, questi vettori agiscono come chiavi di routing, non come sostituti del contenuto.
3. L3 (Livello di Archivio, CPU RAM): Le chiavi e i valori originali (quantizzati) per ogni token compresso sono memorizzati nella RAM della CPU a precisione a 4 bit. Questo funge da fonte di verità per il contenuto.

Il Processo di Decoding

Durante il decoding, il processo opera come un lookup di indice reversibile:

1. Routing: Un router per head valuta le chiavi di riassunto L2 rispetto alla query corrente.
2. Fetching: In base a un budget di routing ($\beta$), il router seleziona i gruppi con il punteggio più alto. Il sistema recupera i token originali corrispondenti dall'archivio L3 della CPU, li dequantizza sulla GPU e li ruota nelle loro posizioni semantiche.
3. Attenzione: Il meccanismo di attenzione opera sull'unione dei token L1, dei riassunti L2 e degli originali L3 ripristinati. I riassunti L2 sono mascherati per i gruppi in cui sono stati recuperati gli originali per evitare il doppio conteggio.

Componenti Tecnici Chiave

• Reader: Un piccolo MLP che valuta l'importanza del token basandosi sugli stati nascosti del decoder precoce. Dirige i token verso L1, L2 (gruppi medi) o L2 (gruppi bassi).
• Compressor: Un pooler cross-attention in stile Perceiver che genera le chiavi di routing a 3 bit. È addestrato per classificare efficacemente i gruppi sotto le proiezioni della query del modello, non per ricostruire il contenuto.
• Quantizzazione (TurboQuant): Utilizza una quantizzazione scalare randomizzata con una rotazione ortogonale casuale fissa per decorrelare le coordinate. L1 e L3 utilizzano una precisione a 4 bit; L2 utilizza una precisione a 3 bit.
• Implementazione: Implementato come una libreria Python di circa 9K righe sopra HuggingFace Transformers, non richiede kernel CUDA personalizzati o modifiche ai pesi del modello.

3. Contributi Chiave

• Architettura: Una gerarchia reversibile a tre livelli che separa il routing (L2) dal contenuto (L3), permettendo elevati rapporti di compressione senza perdita permanente di informazioni.
• Metodologia di Validazione: Un rigoroso framework statistico utilizzando prove accoppiate ($n=500$ per cella), test esatti di McNemar, TOST (Two One-Sided Tests) per l'equivalenza nel recupero binario, e margini di non-inferiorità a lato singolo per la perplexity.
• Risultati sulla Scalabilità: Il metodo dimostra che i costi di qualità diminuiscono all'aumentare delle dimensioni del modello, e si compone liberamente con la quantizzazione dei pesi a 4 bit per modelli $\ge$ 4B di parametri.
• Open Source: Rilascio del pacchetto beta HSQ sotto licenza Apache 2.0, inclusi adapter del modello e harness di valutazione.

4. Risultati

Valutato sulla famiglia Qwen3 (0.6B–14B) attraverso pesi fp16, int8 e int4, e contesti fino a 32K:

• Rapporto di Compressione: Tamarin riduce la cache KV residente in GPU del 12–28× (ad esempio, Qwen3-8B a 32K scende da 4,5 GiB a 172 MiB su GPU, con 2,29 GiB su CPU).
• Prestazioni di Recupero:
  • Sui modelli da 4B e 8B, il recupero è statisticamente equivalente alla baseline non compressa (equivalenza TOST) in tutti i profondità.
  • Al contesto di 32K, Qwen3-8B-int4 raggiunge la parità di recupero della "needle" (ago) di 500/500 con zero coppie discordanti.
• Overhead di Perplexity (PPL):
  • Per i modelli da 4B–14B, l'overhead di PPL è minimo (0,1–0,4%).
  • Sui modelli da 4B, la perplexity multilingue è stata misurata al di sotto della baseline (rapporto 0,9972), suggerendo che il costo della compressione è al di sotto del pavimento di misurazione.
  • I modelli piccoli (0.6B) incorrono in costi più elevati (4–6% in fp16, 6–10% in int4).
• Accordo dei Token: L'accordo greedy a livello di token è del 96–97% con la baseline per i modelli da 4B–14B.
• Velocità:
  • L'overhead di prefill si ammortizza a contesti più lunghi (1,2–1,3× a 32K).
  • Il throughput di decode è attualmente un collo di bottiglia, operando al 40–73% della baseline per contesti fino a 8K, e scendendo al 17–18% a 3K a causa del costo di fetching e dequantizzazione dalla CPU.

5. Claim e Significato

Il paper sostiene che Tamarin riesce a spostare il paradigma della gestione della cache KV dalla eliminazione all'indicizzazione. Ritardando la decisione di quali token siano "importanti" fino al tempo della query (tramite il router), il sistema evita l'irreversibilità dei metodi di evizione.

Claim Specifici:

• Capacità rispetto alla Latenza: Tamarin è presentato come una soluzione per far entrare contesti più grandi e batch size maggiori in una memoria GPU limitata, non per ridurre la latenza.
• Reversibilità: Ogni token compresso è indirizzabile; il sistema non scarta permanentemente le informazioni.
• Scaling del Modello: Il metodo diventa più efficace all'aumentare della capacità del modello, con i modelli da 4B e 8B che mostrano un'equivalenza quasi perfetta rispetto alle baseline.
• Composizione: Funziona perfettamente con la quantizzazione dei pesi a 4 bit per modelli di almeno 4B di parametri senza costi aggiuntivi misurabili.

Limitazioni Esplicite e Confini:

• Modelli Piccoli: I modelli sotto i 4B di parametri (specificamente lo 0.6B) soffrono di un compounding del rumore, portando a un maggiore overhead di PPL.
• Checkpoint Fragili: Il checkpoint Qwen3-14B-int4 ha mostrato "vuoti" stretti nel recupero di singola needle a profondità specifiche, attribuiti alla fragilità della quantizzazione piuttosto che al metodo stesso.
• Famiglie di Modelli: Il metodo è validato statisticamente su Qwen3. Un test iniziale su Mistral-7B ha mostrato un "effetto famiglia" con prestazioni degradate, ipotizzato come causato dall'assenza di normalizzazione delle chiavi per head.
• Throughput: L'implementazione attuale scambia la velocità di decode per la capacità di memoria, con penalità significative di throughput a contesti molto lunghi (32K).

Il paper conclude che Tamarin offre una via pratica e reversibile per l'inferenza a lungo contesto su GPU consumer e datacenter, a patto che l'utente accetti l'attuale compromesso di una riduzione del throughput di decode per massicci risparmi di memoria.