Sintesi Tecnica: KVzap - Pruning della KV Cache Veloce, Adattivo e Fedele

1. Definizione del Problema

Con l'espansione della lunghezza del contesto nei modelli linguistici basati su transformer verso decine o centinaia di migliaia di token, la cache Key-Value (KV) è diventata un collo di bottiglia dominante per l'efficienza dell'inferenza dei Large Language Model (LLM). La cache KV, che memorizza gli stati di attenzione intermedi, consuma una quantità significativa di memoria GPU (ad esempio, 335 GB per un modello da 65B con una lunghezza di sequenza di 128k) e aumenta il tempo di primo token (time-to-first-token), riducendo al contempo il throughput di decodifica.

Sebbene le modifiche architettoniche (ad esempio, Grouped Query Attention, Multi-head Latent Attention) e i meccanismi di attenzione sparsa abbiano affrontato la compressione lungo gli assi della testa ($H$), della dimensione ($D$) e dello strato ($L$), l'asse della lunghezza della sequenza ($T$) rimane in gran parte non affrontato dai metodi ampiamente adottati. Gli attuali metodi di pruning ad hoc spesso non riescono a soddisfare quattro criteri critici per l'implementazione in produzione:

1. Veloce e leggero: Overhead trascurabile.
2. Indipendente dalla fase (Phase-agnostic): Applicabile sia al prefilling (contesto lungo) che alla decodifica (ragionamento).
3. Compatibile con l'ottimizzazione: Compatibile con kernel come FlashAttention2 e PagedAttention.
4. Fedele: Degradazione minima dell'accuratezza.

Lo stato dell'arte attuale, come KVzip, raggiunge un'alta compressione ma richiede un costoso doppio passaggio forward su prompt estesi, rendendolo inadatto per l'inferenza online e i compiti di decodifica.

2. Metodologia

Gli autori introducono KVzap, un'approssimazione veloce e adattiva all'input di KVzip che opera sia durante il prefilling che durante la decodifica. La metodologia consiste in tre componenti principali:

2.1. Politica di Scoring Potenziata (KVzip+)

Gli autori perfezionano innanzitutto il meccanismo di scoring di KVzip. KVzip assegna punteggi ai token basandosi sulla "ricostruzione del contesto", misurando il peso massimo di attenzione che un token riceve quando il prompt viene ripetuto. Gli autori migliorano questo processo incorporando un termine di normalizzazione ispirato a Devoto et al. (2025), che tiene conto del contributo del vettore di valore di un token allo stream residuo.
Lo score KVzip+ ($s^+_i$) è definito come:
$$s^+_i = \max_{j \in \text{<prompt>}} \frac{a_{ji}}{\|W_O v_i\| / \|h_j\|}$$
dove $a_{ji}$ è il peso di attenzione, $W_O$ è la matrice di proiezione di output, $v_i$ è il vettore di valore e $h_j$ è lo stato nascosto.

2.2. Modellazione Surrogata

Per evitare il costo proibitivo del doppio passaggio forward di KVzip, KVzap addestra modelli surrogati leggeri (uno strato lineare o un MLP a due strati) per predire direttamente i log-score ($\log(s^+)$) a partire dagli stati nascosti di input ($h$) del modello.

• Input: Stati nascosti $h_t \in \mathbb{R}^{D_h}$.
• Output: Punteggi per ogni testa KV in $\mathbb{R}^H$.
• Addestramento: I modelli sono addestrati su 1,2 milioni di coppie (stato nascosto, log-score) campionate dal dataset Nemotron-Pretraining-Dataset. La variante MLP utilizza una larghezza nascosta di $D_h/8$ con attivazione GELU.
• Performance: I surrogati raggiungono una correlazione $R^2$ media di 0,60–0,80 con l'oracolo KVzip+, dimostrando che punteggi costosi possono essere approssimati accuratamente a partire dagli stati nascosti.

2.3. Pruning Adattivo ed Eviction

A differenza di KVzip, che impone un budget fisso (ad esempio, mantenere esattamente il 50% dei token), KVzap utilizza il pruning basato su soglia:

• Sogliatura (Thresholding): I token con punteggi predetti inferiori a una soglia fissa $\tau$ vengono scartati. Ciò consente al rapporto di compressione di adattarsi dinamicamente alla densità informativa del prompt di input.
• Finestra Scorrevole (Sliding Window): Per preservare il contesto locale, i più recenti $w=128$ token sono sempre mantenuti, seguendo l'approccio di StreamingLLM.
• Implementazione: Il pruning avviene dopo l'operazione di attenzione. Il processo è computazionalmente efficiente, aggiungendo un overhead trascurabile (limitato all'1,1% per i modelli MLP e allo 0,02% per i modelli lineari).

3. Contributi Chiave

1. KVzip+: Una metrica di scoring migliorata per l'importanza della cache KV che incorpora la normalizzazione dello stato nascosto.
2. Approssimazione Surrogata: Una dimostrazione del fatto che i punteggi KVzip+ possono essere efficacemente approssimati da modelli leggeri addestrati sugli stati nascosti, eliminando la necessità di doppi passaggi forward.
3. KVzap: Una nuova tecnica di pruning che applica questi surrogati con l'eviction basata su soglia e finestre scorrevoli, raggiungendo prestazioni allo stato dell'arte pur soddisfacendo tutti e quattro i criteri pratici (velocità, indipendenza dalla fase, compatibilità con l'ottimizzazione e fedeltà).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su Qwen3-8B, Llama-3.1-8B-Instruct e Qwen3-32B attraverso tre benchmark: RULER (recupero/aggregazione a lungo contesto), LongBench (compiti multi-task a lungo contesto) e AIME25 (ragionamento matematico).

• Compressione vs Accuratezza: KVzap raggiunge una compressione della cache KV di 2–4× con una perdita di accuratezza trascurabile.
  • Su RULER 4k, KVzap eguaglia o supera le prestazioni dell'oracolo KVzip+ e supera altri 15 metodi di pruning (inclusi Expected Attention e Duo Attention).
  • Su LongBench, KVzap mantiene un'accuratezza quasi perfetta fino a una compressione di 2–3×.
  • Su AIME25, KVzap-MLP preserva l'accuratezza del ragionamento anche scartando oltre il 50% della cache KV.
• Adattività: L'approccio basato su soglia regola automaticamente i rapporti di compressione in base alla difficoltà del compito. Ad esempio, ottiene una compressione maggiore sui dati sintetici di RULER rispetto ai dati più densi e reali di LongBench.
• Efficienza: Il metodo aggiunge un overhead computazionale e di memoria trascurabile. Durante la decodifica, i FLOP aggiuntivi utilizzano cicli GPU che altrimenti sarebbero bloccati dalle limitazioni della larghezza di banda della memoria.

5. Significato e Rivendicazioni

Il paper sostiene che KVzap colmi il divario tra la ricerca accademica sul pruning e i motori di inferenza reali. La sua importanza risiede in:

• Praticità: È il primo metodo che riesce a ottenere una compressione allo stato dell'arte senza sacrificare la velocità o richiedere cambiamenti architettonici, rendendolo compatibile con i principali motori di inferenza (vLLM, SGLang, TRT-LLM).
• Fedeltà: Dimostra che gli LLM non sfruttano completamente la cache KV e che le coppie inutilizzate possono essere identificate tramite gli stati nascosti senza riaddestrare il modello base end-to-end.
• Scalabilità: L'approccio è validato su modelli che vanno da 8B a 32B parametri ed è progettato per essere compatibile con le future architetture di attenzione sparsa.

Gli autori riconoscono i limiti, osservando che KVzap è un'aggiunta post-hoc piuttosto che un obiettivo addestrato end-to-end, e che è necessaria ulteriore ingegneria per tradurre i rapporti di compressione in miglioramenti del tempo di esecuzione (wall-clock speedups) tramite kernel ottimizzati (ad esempio, PagedAttention). Tuttavia, pongono che la combinazione di semplicità, alta compressione e prestazioni robuste rende KVzap un candidato primario per l'implementazione in produzione.