Agent Memory Below the Prompt: Persistent Q4 KV Cache for Multi-Agent LLM Inference on Edge Devices

Questo paper presenta un sistema per l'inferenza multi-agente LLM su dispositivi edge che risolve i vincoli di memoria persistendo le cache KV quantizzate a 4 bit su disco, permettendo il ripristino diretto dello stato di attenzione e riducendo il tempo di primo token fino a 136 volte rispetto al ricalcolo completo, con un impatto minimo sulla qualità del modello.

L'essenziale

🧠 Il "Salvadanaio" dei Ricordi: Come far lavorare più Agenti AI su un computer portatile

Immagina di avere un computer portatile potente (come un MacBook moderno) e di voler far lavorare insieme 10 assistenti AI diversi allo stesso tempo. Ognuno di questi assistenti sta leggendo una conversazione lunga, ricordando cosa è stato detto prima, e deve rispondere in tempo reale.

Il problema? Il computer ha una memoria limitata, come un piccolo appartamento. Se ogni assistente occupa una stanza piena di libri (i "ricordi" della conversazione), presto non ci sarà più spazio per nessuno.

🏠 Il Problema: La Stanza Piena di Libri

In passato, quando un assistente finiva il suo turno e un altro doveva iniziare, il computer doveva buttare via tutti i libri della stanza precedente e ricominciare a leggerli da capo per il nuovo assistente.

• Risultato: L'assistente si bloccava per 15 secondi (o più) solo per "riprendere il filo".
• Conseguenza: Se cambiavi assistente 10 volte, perdevi quasi 3 minuti solo ad aspettare. Era come se dovessi rileggere un intero libro ogni volta che cambiavi personaggio in una storia.

💡 La Soluzione: Il "Salvadanaio" sul Disco

Gli autori di questo studio hanno inventato un sistema geniale per risolvere questo problema. Immagina che invece di buttare i libri, li mettano in scatole compattate (usando una tecnica chiamata quantizzazione Q4, che riduce le dimensioni dei libri a un quarto della grandezza originale) e li ripongano in un salvadanaio sicuro sul disco rigido (il tuo hard drive).

Ecco come funziona il loro sistema, passo dopo passo:

1. Compressione Magica (Q4):
   Immagina di prendere un'enciclopedia di 1000 pagine e comprimerla in un quadernino di 250 pagine, mantenendo tutte le informazioni importanti. Questo permette di far entrare 4 volte più assistenti nella stessa memoria del computer.

2. Il Salvadanaio Persistente (Disk Persistence):
   Quando un assistente finisce di lavorare, invece di cancellare i suoi ricordi, il sistema li salva nel "salvadanaio" (il disco). Se il computer si riavvia o se l'assistente deve tornare dopo un'ora, non deve rileggere tutto da capo. Basta prendere la scatola dal salvadanaio e riaprirlo.

   • Tempo perso: Da 15 secondi a mezzo secondo. È come passare da dover rileggere un libro intero a dover solo aprire un cassetto.

3. La Coreografia Perfetta (Interleaving):
   Qui sta la vera magia. Immagina una catena di montaggio dove un operaio sta montando un'auto mentre il suo collega, dall'altra parte della stanza, sta prendendo i pezzi dal magazzino.

   • Mentre l'Assistente A sta "parlando" (generando la risposta), il sistema sta già caricando i ricordi dell'Assistente B dal salvadanaio.
   • Risultato: L'utente non si accorge mai dell'attesa. Il caricamento dei ricordi avviene "nascosto" mentre l'altro assistente lavora.

🚀 Perché è importante?

Questo sistema permette di fare cose che prima erano impossibili su un computer portatile:

• Privacy Totale: Tutto rimane sul tuo computer. Nessuna conversazione viene inviata a server lontani (ottimo per la privacy e le leggi sulla protezione dei dati).
• Nessun Costo: Non paghi per ogni parola generata su cloud.
• Velocità: Anche con conversazioni lunghissime (come un intero libro), il computer risponde quasi istantaneamente.

📊 I Risultati in Pillole

• Prima: Con 10 assistenti, il computer si bloccava continuamente.
• Ora: Con la stessa memoria, il computer gestisce 12 assistenti (invece di 3) e risponde in meno di un secondo.
• Qualità: La "compressione" dei ricordi non ha reso gli assistenti più stupidi. Hanno mantenuto la loro intelligenza quasi intatta (con una differenza di qualità inferiore all'1% in alcuni casi).

🎯 In Sintesi

Questo articolo descrive un sistema che trasforma il computer portatile in un ufficio virtuale infinito. Invece di costringere gli assistenti AI a dimenticare tutto ogni volta che cambiano turno, il sistema li aiuta a conservare i loro "diari di bordo" in modo intelligente e compatto, permettendo loro di riprendere il lavoro istantaneamente, anche dopo aver spento e riaccinto il computer.

È come avere un segretario che non dimentica mai nulla, che lavora velocemente e che non ti fa mai aspettare, tutto direttamente sulla tua scrivania.

Sommario tecnico

1. Il Problema: Gestione della Memoria nei Sistemi Multi-Agente su Edge

Il paper affronta una sfida critica nell'esecuzione di modelli linguistici (LLM) multi-agente su dispositivi edge (es. MacBook Pro M4 Pro con 24 GB di RAM unificata):

• Vincoli di Memoria: I dispositivi edge hanno una RAM fissa e limitata. In un sistema multi-agente, ogni agente richiede il proprio contesto di attenzione (KV cache) indipendente per evitare bias posizionali e garantire l'isolamento dei dati.
• Inefficienza dello Swapping: Con un budget di cache di circa 10 GB (dopo aver sottratto il sistema operativo e i pesi del modello), un dispositivo come l'M4 Pro può ospitare solo 3 agenti con un contesto di 8K in formato FP16. Un workflow di 10 agenti richiede un continuo eviction (rimozione) e reload delle cache.
• Costo Computazionale: Senza persistenza, ogni volta che un agente "evicted" viene richiamato, il sistema deve eseguire un nuovo prefill completo (O(n)) attraverso il modello. Su hardware edge, questo processo è estremamente lento (es. 15,7 secondi per 4K token su Gemma 3), rendendo l'interattività impossibile.
• Limiti delle Soluzioni Attuali: Le soluzioni esistenti come vLLM o SGLang gestiscono la cache in memoria volatile o usano tecniche di prefix caching in FP16, che falliscono sotto pressione di memoria o dopo un riavvio del server, poiché la cache viene persa.

2. Metodologia e Architettura del Sistema

La soluzione proposta è un sistema che persiste le KV cache di ogni agente su disco in formato quantizzato a 4 bit (Q4), permettendo un ripristino diretto senza ricalcolo.

Componenti Chiave:

1. Block Pool con Isolamento per Agente:

   • La cache KV è suddivisa in blocchi fissi da 256 token.
   • Ogni agente ha una cache isolata e indirizzabile, salvata in formato safetensors. Questo garantisce che un riavvio del server o lo swap di modelli non corrompa lo stato degli altri agenti.
   • Il sistema supporta l'isolamento dei dati per conformità GDPR/HIPAA (ogni file è un record dati indipendente).

2. Pipeline di Quantizzazione Q4:

   • Le chiavi e i valori (KV) vengono quantizzati a 4 bit (formato uint32 per i dati pesati + bfloat16 per scale e bias) su disco, in memoria e durante l'operazione di attenzione.
   • Vantaggio: La quantizzazione Q4 riduce l'uso di memoria di circa il 72% rispetto al FP16 (rapporto 0.281), permettendo di ospitare 4 volte più agenti a parità di contesto.

3. BatchQuantizedKVCache:

   • Implementa l'inferenza batchizzata su cache quantizzate multiple, un'operazione non nativamente supportata dalle librerie MLX precedenti.
   • Utilizza un scheduler concorrente che alterna le fasi di prefill (a blocchi) e decode, nascondendo la latenza di I/O del disco dietro la generazione dei token degli altri agenti.

4. Iniezione di Contesto Cross-Phase:

   • Permette di riutilizzare lo stato di attenzione accumulato tra diverse fasi di una conversazione (es. negoziazione, dibattito) senza ricalcolo, trattando la cache come "memoria di lavoro" persistente.

5. Interfaccia API:

   • Espone un'API compatibile con OpenAI, rendendo il sistema trasparente per qualsiasi framework di agenti (es. AutoGen, CrewAI).

3. Risultati Sperimentali

Il sistema è stato valutato su tre architetture distinte: Gemma 3 12B (GQA), DeepSeek-Coder-V2-Lite 16B (MoE/MLA) e Llama 3.1 8B (GQA), su un MacBook Pro M4 Pro.

• Riduzione della Latenza (TTFT - Time To First Token):

  • Il ripristino della cache da disco (Warm) riduce il TTFT da secondi a millisecondi.
  • Gemma 3: Riduzione fino a 136x (da 172s a 1.3s a 32K contesto).
  • DeepSeek: Riduzione fino a 76x.
  • Llama: Riduzione fino a 111x.
  • Il ripristino da disco è sub-secondo (es. ~577 ms per 4K su Gemma), rientrando nella soglia di accettabilità per l'interattività umana (<1s).

• Capacità di Memoria:

  • A 8K contesto, il formato Q4 permette di ospitare 12 agenti su M4 Pro, contro soli 3 in FP16.
  • A 16K contesto, il FP16 non riesce nemmeno a ospitare un singolo workflow multi-agente, mentre Q4 ne supporta 6.

• Confronto con vllm-mlx:

  • vllm-mlx (che usa cache FP16 volatile) fallisce completamente sotto pressione di memoria multi-agente o dopo un riavvio.
  • Il sistema proposto mantiene prestazioni simili in termini di latenza quando la cache è calda, ma sopravvive a riavvii e swap di memoria grazie alla persistenza su disco.

• Qualità del Modello (Perplexity):

  • L'impatto della quantizzazione Q4 sulla qualità è minimo:
    • Gemma 3: -0.7% (entro il rumore di misura).
    • Llama 3.1: +2.8%.
    • DeepSeek: +3.0%.
  • Questi valori sono coerenti con la letteratura esistente sulla quantizzazione KV cache.

• Scenari Multi-Fase e Routing:

  • In scenari complessi (es. interrogatorio a 5 fasi), l'uso di cache persistenti riduce il tempo totale di esecuzione del 23% (Gemma) e del 17% (DeepSeek) grazie all'accumulo dello stato tra le fasi.
  • In benchmark di routing (10 esperti), il TTFT scende da ~20s (cold) a <1s (warm).

4. Contributi Principali

1. Persistenza della KV Cache: Un pool di blocchi persistente che sopravvive ai riavvii e allo swap, memorizzato in safetensors Q4.
2. Inferenza Batchizzata Quantizzata: Implementazione di BatchQuantizedKVCache per gestire inferenze concorrenti su cache Q4 multiple.
3. Memoria di Lavoro Cross-Phase: Meccanismo per accumulare e riutilizzare lo stato di attenzione tra diverse fasi di conversazione senza ricalcolo.
4. Validazione su Architetture Eterogenee: Dimostrazione empirica su modelli GQA, MoE e MLA, mostrando come l'astrazione ModelCacheSpec permetta di gestire differenze architetturali (es. dimensioni asimmetriche K/V in DeepSeek).

5. Significato e Implicazioni

Questo lavoro rappresenta un passo fondamentale per l'adozione di sistemi multi-agente autonomi su dispositivi edge.

• Sovranità dei Dati: Permette di eseguire workflow complessi e multi-agente interamente in locale, evitando la trasmissione di dati sensibili su cloud (cruciale per GDPR e HIPAA).
• Interattività Reale: Trasforma l'esperienza utente da tempi di attesa inaccettabili (decine di secondi) a risposte fluide, rendendo fattibili applicazioni come assistenti personali complessi o team di agenti collaborativi su hardware consumer.
• Efficienza Energetica e di Costo: Elimina la necessità di costose infrastrutture GPU cloud per carichi di lavoro multi-agente, sfruttando l'hardware esistente degli utenti finali.

Il sistema è open-source e disponibile su GitHub, fornendo un'infrastruttura pronta all'uso per framework di agenti che necessitano di gestione della memoria scalabile e persistente.