CONF-KV: Confidence-Aware KV Cache Eviction with Mixed-Precision Storage for Long-Horizon LLM

Riepilogo Tecnico: CONF-KV

Enunciato del Problema
L'inferenza di modelli linguistici su larga scala (LLM) con orizzonti temporali lunghi incontra un collo di bottiglia critico in cui la cache Key-Value (KV) diventa il principale consumatore di memoria GPU, causando una latenza di attenzione per token che scala linearmente con la lunghezza della sequenza. Le politiche di espulsione esistenti, come le finestre scorrevoli o i metodi basati sull'attenzione storica (ad esempio H2O, SCISSORHANDS), si affidano alla recenza statica o a segnali passati per determinare quali token mantenere. Questi approcci spesso non riescono a sfruttare un segnale disponibile ad ogni passo di decodifica: l'incertezza corrente del modello. Di conseguenza, le politiche statiche possono espellere prematuramente il contesto durante fasi di generazione difficili o trattenere token non necessari quando il modello è sicuro, portando a un compromesso subottimale tra utilizzo della memoria e qualità della generazione.

Metodologia: CONF-KV
Gli autori propongono CONF-KV, un gestore di cache KV consapevole della confidenza che adatta dinamicamente il budget della cache in base alla distribuzione corrente del prossimo token del modello. Il sistema opera attraverso i seguenti meccanismi:

1. Stima della Confidenza: Ad ogni passo di decodifica, il metodo converte i logit in un punteggio di confidenza scalare ($c$). Questo punteggio è una combinazione ponderata di entropia normalizzata, il margine di log-probabilità tra i primi due token e la probabilità del primo token.
   $$c = w_H(1 - \hat{H}) + w_m\sigma(m) + w_p p_{(1)}$$
   dove $\hat{H}$ è l'entropia normalizzata, $m$ è il margine logaritmico e i pesi sono impostati su $(0.4, 0.3, 0.3)$.

2. Selezione Adattiva del Budget: Una soglia $\tau$ determina il budget della cache ($N$) per il passo corrente.

   • Se $c \ge \tau$ (alta confidenza), viene selezionato un budget ristretto ($N_{high}$), permettendo un'espulsione aggressiva.
   • Se $c < \tau$ (bassa confidenza/incertezza), viene selezionato un budget più ampio ($N_{low}$) per trattenere più contesto.
     Questo crea un profilo di memoria "a dente di sega" in cui la cache si espande durante fasi di incertezza e si contrae durante fasi di alta confidenza.

3. Classificazione ed Espulsione dei Token: All'interno del budget selezionato, i token sono classificati in base a un punteggio composito che combina la massa di attenzione a media mobile esponenziale (EMA) e la recenza. Una finestra protetta rigida ($P$) garantisce che i token più recenti non vengano mai espulsi per preservare la coerenza locale.

4. Integrazione di Sistema:

   • Archiviazione in Precisione Mista: I $W$ token più recenti sono archiviati in FP16, mentre i token più vecchi mantenuti sono quantizzati in INT8 per testa e canale. Questo riduce l'impronta di memoria minimizzando il degrado della perplessità rispetto alla quantizzazione a bit inferiori.
   • Budget Piramidali (CONF-KV-L): Una variante opzionale assegna i budget in modo non uniforme tra i livelli, assegnando budget più piccoli ai livelli più profondi basandosi sull'osservazione che i livelli più profondi spesso concentrano le informazioni in un numero minore di token.
   • Implementazione: Il gestore utilizza la compattazione contigua per mantenere un layout denso per i kernel di attenzione, evitando tabelle di indirizzamento, e fonde la dequantizzazione nelle letture di attenzione a blocchi.

Contributi Chiave

• Politica Consapevole della Confidenza: Una strategia di espulsione innovativa che utilizza la distribuzione corrente dell'output per guidare le decisioni sul budget della cache per passo, invece di affidarsi esclusivamente a segnali storici.
• Progettazione di Sistema: Un'implementazione completa che combina espulsione adattiva con attenzione softmax online a blocchi, archiviazione mista FP16/INT8 e budget piramidali opzionali per i livelli.
• Validazione Meccanicistica: Un test che dimostra come il punteggio di confidenza sia fortemente anticorrelato (Pearson $r = -0.77$) con lo spostamento KL indotto dall'ablazione del contesto recente, validando l'assunzione che una bassa confidenza indichi la necessità di trattenere più contesto.
• Valutazione Completa: Ampie valutazioni su quattro famiglie di modelli (GPT-2, Qwen-14B/32B, gpt-oss-20b) e carichi di lavoro diversificati, inclusa la perplessità, il recupero in contesti lunghi e compiti di agenti web.

Risultati Sperimentali

• Compromesso Memoria-Qualità: Su GPT-2 con una lunghezza di generazione di 2048 token, CONF-KV+INT8 corrisponde all'impronta di memoria di una finestra scorrevole fissa da 512 token (~38,7 MB) ma riduce la perplessità di 3,11 punti rispetto alla finestra scorrevole. CONF-KV-L (variante piramidale) riduce ulteriormente la memoria a 34,2 MB migliorando la perplessità di 3,89 punti, colmando il 74% del divario di qualità tra una finestra scorrevole e la KV completa.
• Recupero in Contesti Lunghi: Nei test "Ago nel Pagliaio" (NIAH) fino a 32K token, CONF-KV raggiunge un'accuratezza di recupero del 91,4%, superando significativamente le finestre scorrevoli (53,8%) e H2O (80,6%).
• Compiti di Agente: Su 75 compiti VisualWebArena, CONF-KV mantiene il 95,3% del tasso di successo della KV completa riducendo la memoria di picco di 2,8 volte.
• Throughput e Latenza: CONF-KV riduce la latenza per passo di 1,8 volte rispetto alla KV completa su GPT-2 e Qwen-32B. Il throughput scala con la dimensione del batch; a batch size 8, CONF-KV raggiunge 2,06 volte il throughput della KV completa.
• Ablazione: Isolare il segnale di confidenza mostra che un'espulsione casuale allo stesso tasso degrada le prestazioni a 36,54 PPL, mentre CONF-KV completo raggiunge 30,92 PPL, confermando che sia il gating della confidenza che il meccanismo di classificazione sono essenziali.

Significato e Affermazioni
Il paper afferma che CONF-KV dimostra che l'incertezza corrente è un metadato di sistema utile per migliorare il fronte di Pareto memoria-qualità. Il metodo è distinto perché è prospettico: non può recuperare i token espulsi ma previene l'espulsione prematura durante fasi difficili. L'approccio è progettato per essere senza addestramento, non altera i pesi del modello ed è ortogonale alle tecniche esistenti come l'allocazione per testa, l'allocazione per livello e la selezione della precisione, permettendone la composizione con esse.

Gli autori notano limitazioni, tra cui il fatto che il metodo è un "no-op" su contesti brevi che non superano mai la soglia di espulsione e che la compattazione contigua comporta un costo di raccolta (sebbene ciò sia compensato dai risparmi di attenzione). Riconoscono inoltre che il segnale di confidenza diventa meno informativo sotto campionamento ad alta temperatura. Infine, il paper evidenzia una considerazione di doppio uso: mentre costi di inferenza più bassi migliorano l'efficienza, potrebbero anche abbassare le barriere per automazioni indesiderate a lungo termine, suggerendo che i controlli di sicurezza per i modelli sottostanti rimangano necessari.