LRAgent: Efficient KV Cache Sharing for Multi-LoRA LLM Agents

Sintesi Tecnica: LRAgent

Definizione del Problema

Nei sistemi di agenti multi-LLM, la specializzazione dei ruoli è frequentemente realizzata attraverso architetture Multi-LoRA, dove molteplici agenti condividono un unico backbone pre-addestrato ma utilizzano distinti e leggeri adattatori LoRA per compiti specifici (ad esempio, pianificazione, azione, riflessione). Sebbene questo design migliori l'efficienza dei parametri, introduce una significativa ridondanza durante l'inferenza.

Quando gli agenti collaborano su traiettorie lunghe e aumentate da strumenti (tool-augmented), spesso elaborano contesti sovrapposti. Tuttavia, i sistemi esistenti tipicamente mantengono cache KV indipendenti per ogni agente. Ciò risulta in:

1. Overhead di Memoria: Ogni agente memorizza una cache KV completa per lo stesso contesto, portando a una scalabilità lineare della memoria con il numero di agenti.
2. Overhead Computazionale: Gli agenti ricalcolano ridondantemente gli stati nascosti e le proiezioni KV per segmenti di contesto già elaborati da altri agenti.

I metodi di condivisione della cache KV esistenti (ad esempio, ICaRus, Cache2Cache, KVLink) si concentrano generalmente sul disallineamento posizionale o richiedono modifiche architettoniche e addestramento aggiuntivo. Fondamentalmente, essi non sfruttano esplicitamente la struttura Multi-LoRA, fallendo spesso nell'eliminare la computazione ridondante per contesti che sono già stati elaborati da altri agenti nello stesso sistema.

Metodologia

Gli autori propongono LRAgent, un framework di condivisione della cache KV progettato specificamente per sistemi di agenti Multi-LoRA. L'intuizione centrale che guida LRAgent è che, per lo stesso contesto, la cache base (attivazioni dal backbone pre-addestrato condiviso) rimane altamente simile tra gli agenti, mentre l'output dell'adattatore (attivazioni dai pesi LoRA) agisce come una piccola perturbazione, ampiamente decorrelata, contenente informazioni specifiche dell'agente.

1. Decomposizione della Cache

LRAgent decompone la cache dei valori ($Y$) in due componenti:

• Cache Base ($Y_{base}$): Derivata dai pesi pre-addestrati congelati ($W_0$). Questa componente è condivisa tra tutti gli agenti.
• Output dell'Adattatore ($\Delta Y$): Derivata dai pesi LoRA ($A, B$). Questa componente è specifica per ogni agente.

Gli autori osservano che la cache base presenta un'alta similarità del coseno (>0,97) tra gli agenti, mentre l'output dell'adattatore ha una bassa similarità. Di conseguenza, la condivisione della cache base preserva la maggior parte dell'informazione semantica, mentre il contributo dell'adattatore può essere memorizzato in modo efficiente.

2. Schemi di Condivisione

LRAgent introduce due strategie di condivisione specifiche:

• BaseShared:

  • Meccanismo: Condivide una singola cache base tra tutti gli agenti. Ogni agente mantiene la propria cache LR (Low-Rank cache), che memorizza le attivazioni intermedie dopo la proiezione verso il basso del LoRA ($X_i A_i$) in una dimensione di basso rango $r$.
  • Ricostruzione: Al runtime, il contributo completo dell'adattatore viene ricostruito moltiplicando la cache LR per la matrice di proiezione verso l'alto del LoRA ($B_i$) solo quando necessario.
  • Vantaggio: Riduce significativamente l'uso della memoria (circa $1/N$ rispetto al baseline non condiviso) evitando l'archiviazione a dimensione intera degli output dell'adattatore. Tuttavia, richiede ancora la ricalcolazione della cache LR per l'intero contesto accumulato durante il passaggio tra gli agenti ("LR prefill").

• BaseLRShared:

  • Meccanismo: Si basa su un'architettura Multi-LoRA Shared-A, dove la matrice di proiezione verso il basso $A$ è condivisa tra tutti gli agenti (mentre $B$ rimane specifica per il compito).
  • Innovazione: Poiché $A$ è condivisa, anche la cache LR ($X_i A$) è altamente simile tra gli agenti. Questo schema condivide una singola cache LR per l'intero sistema, oltre alla cache base.
  • Vantaggio: Elimina la necessità di ricalcolare il "LR prefill". Gli agenti possono riutilizzare sia la cache base che la cache LR per i token precedentemente visti, riducendo la complessità computazionale in modo comparabile a una cache completamente condivisa.

3. Flash-LoRA-Attention

Per affrontare l'overhead computazionale della ricostruzione del contributo dell'adattatore dalla cache a basso rango, gli autori introducono Flash-LoRA-Attention.

• Problema: La ricostruzione ingenua ($V_{lr} B$) espande la cache alla dimensione intera $d_{out}$ prima dell'attenzione, incurrendo in una complessità $O(L \cdot r \cdot d_{out})$.
• Soluzione: Il kernel riordina le moltiplicazioni matriciali usando l'associatività: $O = P(V_{base} + V_{lr}B) \rightarrow O = PV_{base} + (PV_{lr})B$.
• Effetto: La moltiplicazione pesata dall'attenzione viene eseguita prima sulla cache a basso rango ($P V_{lr}$), scalando con $O(L \cdot r)$. La proiezione verso l'alto ($B$) viene applicata solo una volta per blocco di query. Ciò riduce l'overhead di espansione di un fattore $r/d_{out}$.

Risultati Chiave

Gli autori hanno valutato LRAgent su LLaMA-3.1-8B-Instruct e Ministral-8B-Instruct utilizzando i benchmark HotpotQA e ScienceQA con traiettorie di agenti multi-hop (ruoli Plan, Action, Reflect).

• Accuratezza:
  • BaseShared e BaseLRShared mantengono un'accuratezza molto vicina al baseline Non-Shared (calo medio $\le 0,7\%$ e $\le 1,5\%$, rispettivamente).
  • Al contrario, FullShared (condivisione ingenua della cache completa) e DroidSpeak (ricalcolo selettivo) soffrono cali di accuratezza maggiori (fino al $5,3\%$ e $2,6\%$, rispettivamente), evidenziando l'importanza di preservare i contributi specifici dell'adattatore.
• Efficienza della Memoria:
  • LRAgent riduce l'uso della memoria della cache KV di quasi 1/3 rispetto al baseline Non-Shared, in modo comparabile ad altri metodi di condivisione della cache e solo marginalmente superiore a FullShared.
  • Le cache LR sono di dimensioni trascurabili grazie alla dimensione di basso rango.
• Throughput e Latenza:
  • Con Flash-LoRA-Attention, BaseLRShared ottiene guadagni di throughput fino a 2,46$\times$ rispetto a Non-Shared, avvicinandosi alle prestazioni di FullShared.
  • Il Time-to-First-Token (TTFT) è ridotto fino a 4,44$\times$ per BaseLRShared, superando significativamente DroidSpeak.
  • BaseShared ottiene un guadagno di throughput fino a 1,42$\times$, bilanciando l'efficienza della memoria con il risparmio computazionale.

Significato e Rivendicazioni

Il paper afferma che LRAgent è il primo lavoro ad adattare esplicitamente la condivisione della cache KV agli ambienti di agenti Multi-LoRA. La sua importanza risiede in:

1. Decoupling (Disaccoppiamento) tra Condiviso e Specifico: Separando la cache in una componente base condivisa e una componente dell'adattatore a basso rango, LRAgent risolve il compromesso tra efficienza di memoria/calcolo e accuratezza specifica del ruolo.
2. Eliminazione Computazionale: A differenza dei metodi precedenti che si limitano a riutilizzare la memoria ma ricalcolano comunque gli stati nascosti (es. DroidSpeak), BaseLRShared (sotto Shared-A) elimina la computazione ridondante per i contesti precedentemente elaborati, raggiungendo un'efficienza vicina alla condivisione completa della cache senza la penalità di accuratezza.
3. Ottimizzazione del Kernel: Il kernel Flash-LoRA-Attention dimostra che l'espansione della cache a basso rango può essere resa computazionalmente economica attraverso il riordinamento, rendendola praticabile per il deployment.

Gli autori concludono che LRAgent consente ai sistemi di agenti Multi-LoRA di raggiungere un'efficienza di memoria e di inferenza vicina alla condivisione completa della cache, preservando al contempo l'accuratezza dei baseline non condivisi, rendendolo una soluzione robusta per applicazioni di agenti LLM scalabili.