BeaconKV: compressione della cache key-value guidata da query faro per un’inferenza efficiente dei modelli di ragionamento di grandi dimensioni

Abstract degli autori

I modelli di ragionamento di grandi dimensioni (LRM) raggiungono capacità superiori di risoluzione dei problemi generando lunghe catene di pensiero (Chain-of-Thought, CoT), ma la cache key-value (KV) che ne deriva cresce linearmente con la lunghezza della sequenza e crea gravi limiti di memoria, superando spesso la capacità delle GPU per le tracce di ragionamento più lunghe. I metodi esistenti di compressione della cache KV si basano sulle query recenti per stimare l’importanza dei token nei passaggi futuri, assumendo implicitamente che siano indicatori affidabili dei futuri schemi di attenzione. Dimostriamo che questa ipotesi non regge nei ragionamenti di lunga durata: alcuni passaggi della generazione producono token che tornano su parti precedenti del ragionamento (Thought Revisiting Tokens, TRT) e rivolgono nuovamente l’attenzione a contesti passati e lontani, come i piani per risolvere il compito formulati nelle prime fasi della traccia. Attraverso un’analisi sistematica, scopriamo che le query corrispondenti ai TRT si raggruppano in un numero ristretto di gruppi di similarità nello spazio degli embedding. Sulla base di questa osservazione, proponiamo BeaconKV, un metodo di compressione della cache KV che non richiede addestramento e mantiene query faro, rappresentanti compatti di ciascun gruppo globale di query, per prevedere quali coppie KV saranno consultate di nuovo senza conservare l’intera cronologia delle query. Su quattro LRM open source e diversi benchmark di ragionamento, BeaconKV supera in generale i metodi di compressione esistenti, ottenendo una riduzione dell’uso di memoria fino a $5.8\times$, mantenendo un’accuratezza quasi pari a quella della cache completa e aumentando il throughput di oltre $4.3\times$.