PolyKV: A Shared Asymmetrically-Compressed KV Cache Pool for Multi-Agent LLM Inference

1. Enunciato del Problema

L'inferenza dei Modelli Linguistici di Grande Dimensione (LLM) è dominata da colli di bottiglia nella memoria causati dalla cache Chiave-Valore (KV), che memorizza gli stati di attenzione per ogni token nel contesto. Questo problema è esacerbato nei sistemi di inferenza multi-agente, dove $N$ agenti che elaborano lo stesso contesto documentale condiviso richiedono tradizionalmente $N$ cache KV indipendenti e in precisione completa. Ciò comporta una scalabilità lineare della memoria ($O(N)$), rendendo inefficiente l'esecuzione di più agenti su contesti lunghi.

Le soluzioni esistenti affrontano questo problema in isolamento:

• Compressione KV: Riduce la memoria per cache tramite quantizzazione, ma opera su cache isolate, per richiesta.
• Condivisione Multi-Agente: Riutilizza i prefissi comuni tra gli agenti, ma utilizza tipicamente cache in precisione completa, non sfruttando la compressione.

PolyKV mira a risolvere l'intersezione di questi problemi: creare un singolo pool di cache KV compresso con perdita (lossy), condiviso, accessibile simultaneamente da più agenti concorrenti.

2. Metodologia

PolyKV introduce un'architettura innovativa che combina quantizzazione asimmetrica, compressione basata su rotazione e un'astrazione della memoria condivisa.

A. Architettura Principale

• SharedKVPool: Un'astrazione di scrittura singola, lettura multipla. Calcola lo stato KV completo per un contesto documentale condiviso una sola volta, lo comprime e lo memorizza come un singolo oggetto (complessità di memoria $O(1)$ rispetto al numero di agenti).
• PooledAgent: Agenti di inferenza individuali che ricevono un riferimento al SharedKVPool. Al momento dell'inferenza, ogni agente inietta i tensori KV decompressi nella propria nuova istanza DynamicCache, strato per strato, aggirando l'accumulo standard per ogni passaggio in avanti. Questo elimina l'overhead di copia per agente.

B. Schema di Compressione Asimmetrica (TurboQuant MSE)

PolyKV impiega una strategia di quantizzazione asimmetrica adattata alla sensibilità delle Chiavi (K) e dei Valori (V):

1. Chiavi (K): Quantizzate a int8 (q8_0) utilizzando la quantizzazione lineare per tensore. Questo preserva la stabilità dei punteggi di attenzione softmax, che sono altamente sensibili alla quantizzazione a basso numero di bit.
2. Valori (V): Compressi utilizzando TurboQuant MSE a precisione 3-bit. Questo comporta:
   • Rotazione: Applicazione di una Trasformata di Walsh-Hadamard Veloce (FWHT) per ruotare il vettore di input, distribuendo i valori anomali e facendo convergere la distribuzione verso una Gaussiana $N(0, 1)$.
   • Quantizzazione: Mappatura delle coordinate ruotate su centroidi Lloyd-Max precalcolati (ottimizzati per $N(0, 1)$).
   • Dequantizzazione: Recupero dei valori dei centroidi e applicazione della FWHT inversa.

Rapporto di Compressione:
Con le Chiavi a 8 bit e i Valori a 3 bit (rispetto a una linea di base bfloat16 a 16 bit), il rapporto di compressione teorico è:
$$\text{Rapporto} = \frac{16}{(8 + 3)/2} = \frac{16}{5.5} \approx 2.91\times$$

3. Contributi Chiave

• Architettura a Pool Condiviso: La prima implementazione di un modello di memoria KV compresso, di scrittura singola e lettura multipla per agenti concorrenti, eliminando la scalabilità di memoria $O(N)$ dei precedenti sistemi multi-agente.
• TurboQuant Asimmetrico: Una combinazione specifica di Chiavi int8 e Valori ruotati con FWHT a 3 bit che ottiene un rapporto di compressione stabile di 2,91× attraverso diverse architetture di modelli.
• Validazione Cross-Modello: Ha dimostrato stabilità sia su SmolLM2-1.7B che su Llama-3-8B (inclusa l'architettura GQA), dimostrando che il metodo è agnostico rispetto all'architettura.
• Scalabilità: Ha scalato con successo a 15 agenti concorrenti che condividono un singolo contesto, con risparmi di memoria che crescono linearmente con il numero di agenti (fino a una riduzione del 97,7%).
• Scoperta di Inversione PPL: Ha scoperto che per contesti coerenti e lunghi, la cache compressa può effettivamente produrre una perplessità inferiore rispetto alla linea di base in precisione completa, suggerendo che il rumore di quantizzazione agisce come regolarizzazione implicita.

4. Risultati Sperimentali

Il sistema è stato valutato su due modelli, tre lunghezze di contesto (da 600 a 7.194 token) e fino a 15 agenti.

Efficienza della Memoria

• Llama-3-8B (15 agenti, contesto 4K): Ha ridotto la memoria della cache KV da 19,8 GB (precisione completa) a 0,45 GB.
• Riduzione: 97,7% di riduzione della memoria mantenendo un'alta qualità di output.
• Scalabilità: L'utilizzo di memoria per il pool rimane costante ($O(1)$) indipendentemente dal numero di agenti (es. 0,116 GB per 3 agenti vs 10 agenti a contesto 2K).

Metriche di Qualità

• Perplessità (PPL):
  • Stabilità: Il delta PPL è invariante rispetto al numero di agenti.
  • Tendenza della Lunghezza del Contesto: Il delta PPL migliora all'aumentare della lunghezza del contesto.
    • A 2K token: +1,59% di degradazione.
    • A 4K token: +0,57% di degradazione.
    • Inversione: A 1.851 token coerenti su SmolLM2-1.7B, la cache compressa ha ottenuto un delta PPL di -0,26% (migliore della linea di base).
• Similarità Semantica (BERTScore):
  • I punteggi F1 medi sono variati da 0,928 a 0,970 su tutte le configurazioni.
  • La sovrapposizione dei token era perfetta (1,000) nei test SmolLM2.
  • La leggera degradazione PPL non si è tradotta in perdita semantica; le variazioni di formulazione hanno preservato l'equivalenza fattuale.

5. Significato e Ipotesi

Significato:
PolyKV è il primo sistema a combinare un singolo pool KV compresso con perdita, condiviso, con accesso multi-lettore concorrente. Cambia fondamentalmente l'economia dell'inferenza LLM multi-agente, consentendo una concorrenza massiccia su contesti condivisi con un'impronta di memoria minima, rendendo fattibile l'esecuzione di sciami di agenti complessi su hardware di livello consumer (es. GPU T4 di Kaggle).

L'Ipotesi di Regolarizzazione:
Gli autori ipotizzano che il rumore TurboQuant FWHT agisca come regolarizzazione implicita. In documenti lunghi e coerenti, le teste di attenzione spesso si focalizzano su correlazioni ridondanti e spurie (pozzi di attenzione). Il rumore di quantizzazione introdotto dalla rotazione FWHT interrompe questi schemi spuri, "denoizzando" efficacemente il meccanismo di attenzione. Questo spiega perché la PPL migliora (o si degrada meno) all'aumentare della lunghezza del contesto e della coerenza lessicale.

6. Limitazioni e Lavori Futuri

• Scalabilità: Gli esperimenti sono stati limitati a modelli da 1,7B e 8B; il comportamento a 70B+ è sconosciuto.
• Metriche: Il tempo fino al primo token (TTFT) e il throughput non sono stati riportati.
• Validazione: Il meccanismo di inversione PPL richiede studi di ablazione controllati per confermare la causalità.
• Futuro: I piani includono test su modelli più grandi (Qwen2.5-7B, 70B), la standardizzazione dei protocolli di valutazione WikiText-2 e la misurazione della latenza end-to-end.