Sintesi Tecnica: RKSC – Condivisione della KV Cache Consapevole del Ragionamento e Uscita Anticipata Convincente

1. Definizione del Problema

I moderni modelli linguistici di grandi dimensioni (LLM) progettati per il ragionamento complesso (ad esempio, modelli come DeepSeek-R1 o di tipo o1), spesso impiegano strategie di inferenza multi-branch, come la "Tree of Thoughts" o la self-consistency, dove vengono generate molteplici traiettorie di ragionamento per un singolo problema. Il documento identifica due ridondanze strutturali negli attuali sistemi di inferenza in produzione (come vLLM e SGLang) che degradano l'efficienza in queste pipeline:

1. Spreco di KV Cross-Branch: Le branche parallele condividono tipicamente una parte significativa del loro prefisso (la dichiarazione del problema e il contesto iniziale). Tuttavia, i sistemi esistenti riutilizzano la cache Key-Value (KV) solo quando le sequenze di token sono byte-identiche. Se le branche divergono semanticamente ma condividono gli stessi token del prefisso, o se condividono un prefisso ma utilizzano una formulazione diversa per gli indizi (hints), il sistema ricalcola l'estesa attenzione $O(n^2)$ per il prefisso condiviso indipendentemente per ogni branca.
2. Verifica Over-Deep (Troppo Profonda): Nelle pipeline di verifica con reward di processo, il modello esegue un passaggio forward a profondità completa per valutare ogni branca generata, anche quando la confidenza della generazione è già decisiva o quando lo stato interno del modello (distribuzione dei logit) si è stabilizzato in uno strato intermedio. Gli attuali sistemi mancano di meccanismi per saltare l'intero passaggio di verifica o per interromperlo anticipatamente sulla base dell'entropia per singolo strato.

2. Metodologia: Il Framework RKSC

Gli autori propongono RKSC (Reasoning-Aware KV Cache Sharing), un framework di inferenza privo di addestramento che elimina queste ridondanze attraverso tre meccanismi complementari e commutabili. RKSC non richiede fine-tuning, modifiche all'architettura o kernel CUDA personalizzati.

A. ASKS (Attention-Similarity KV Sharing)

ASKS generalizza il caching del prefisso basato sull'identità dei token a un regime semantico.

• Meccanismo: Invece di richiedere token byte-identici, ASKS calcola una similitudine del coseno pesata tra gli stati nascosti (hidden states) della suffissa di una branca e quelli del radice.
• Pesatura: Il punteggio di similitudine ($\sigma_b$) applica pesi in crescita esponenziale agli strati successivi del transformer ($w_l = \exp(\alpha l/L)$), poiché questi strati trasportano informazioni più specifiche sul compito e sono migliori nel rilevare la divergenza semantica.
• Azione: Se $\sigma_b \geq \tau$, la branca è considerata semanticamente coerente con il prefisso e riceve la cache KV condivisa tramite un'operazione di repeat_interleave a costo zero. Le branche al di sotto della soglia tornano a una computazione indipendente.
• Generalizzazione: Questo generalizza strettamente il caching basato sull'identità dei token (dove $\sigma_b=1$ per token identici) e consente il riutilizzo per branche con formulazioni superficiali diverse ma rappresentazioni semantiche simili.

B. CGEE (Confidence-Gated Early Exit)

CGEE riduce i costi di verifica attraverso una strategia di uscita a due livelli:

1. Verify-Skip (Livello 1): Prima di eseguire il passaggio di verifica, il sistema controlla la confidenza della generazione. Se la branca vincente ha un'alta confidenza assoluta ($p(b) \geq \tau_{conf}$) e domina i competitor con un margine netto, l'intero passaggio forward di verifica viene saltato.
2. Uscita per Entropia a Livello di Strato (Livello 2): Se la verifica procede, hook leggeri monitorano l'entropia dei logit ad ogni strato. Se l'entropia scende al di sotto di una soglia ($\theta$) e si stabilizza (variazione $< \epsilon$) in uno strato intermedio $l^*$, il passaggio forward viene terminato anticipatamente. L'output parziale viene utilizzato come risultato della verifica, saltando i restanti $L - l^*$ strati.

C. RSBCM (Reasoning-Selective Block Cache Manager)

Per prevenire una crescita illimitata della memoria durante le ricerche profonde in albero, RSBCm gestisce la cache condivisa.

• Politica di Evizione: I blocchi vengono evitti in base a un punteggio di importanza $\omega = \text{branch\_score} / (\text{depth} + 1)$. Ciò dà priorità alla conservazione dei blocchi dalle branche con punteggio elevato a profondità minori, evidenziando ed evitando per primi i blocoli delle traiettorie profonde e a bassa confidenza.

3. Contributi Chiave

• ASKS: Introduce un gate di similitudine semantica che ottiene un riutilizzo della KV cache del 28,6% superiore rispetto al matching dei token esatti in scenari con formulazioni diverse, generalizzando strettamente l'attuale caching del prefisso.
• CGEE: Propone un meccanismo di uscita a due livelli che salta la verifica quando la confidenza è decisiva e tronca il passaggio quando l'entropia per strato si stabilizza, riducendo la latenza di verifica senza riaddestramento.
• RSBCM: Fornisce una strategia di evizione basata sulla priorità della profondità per mantenere la stabilità della cache.
• Design a Zero-Overhead: Il framework non aggiunge parametri apprendibili e non richiede modifiche all'architettura, agendo come un wrapper "drop-in" per le attuali pipeline di ragionamento.

4. Risultati Sperimentali

Il framework è stato valutato su cinque famiglie di modelli (7B–10B parametri, inclusi Qwen2.5, Mistral, Falcon3 e Llama-3) e quattro benchmark (GPQA Diamond, MMLU-STEM, ARC-Challenge, GSM8K) su 1.000 problemi.

• Accelerazione (Speedup): RKSC ottiene un'accelerazione media di 3,008× rispetto alla baseline "No-KV" (che ricalcola i prefissi) e un miglioramento di 1,66× rispetto al caching del prefisso basato su token esatti equivalente a vLLM. L'accelerazione massima osservata è stata di 3,990× (Llama-3-8B su MMLU-STEM).
• Accuratezza: Il framework mantiene un'alta accuratezza con un tasso di errore indotto da CGEE di solo lo 0,37% (6 errori su 1.616 chiamate di verifica). Gli errori si sono concentrati nei benchmark a scelta multipla dove la fluidità della generazione talvolta divergeva dai segnali di correttezza; i compiti di generazione a testo libero hanno mostrato un tasso di errore vicino allo zero.
• Scomposizione dei Componenti:
  • La sola condivisione KV ha fornito un'accelerazione di 1,343× (riduzione della latenza del 25,5%).
  • Il sistema combinato RKSC (KV + CGEE) ha fornito un'accelerazione di 1,622× nella decomposizione del throughput e di 3,008× nella valutazione estesa.
  • Il tasso di "verify-skip" di CGEE è stato adattivo alla difficoltà del compito, variando da circa il 42% sui problemi difficili di GPQA a circa il 92% sui problemi più semplici di MMLU-STEM.
• Agnosticismo dell'Architettura: L'accelerazione della condivisione KV è stata coerente in tutti i cinque modelli (deviazione standard 0,02×), confermando che i guadagni sono guidati dalla struttura del carico di lavoro dell'inferenza (lunghezza del prefisso e fattore di branching) piuttosto che dalle specifiche architetture dei modelli.

5. Significato e Rivendicazioni

Il paper sostiene che RKSC affronti inefficienze strutturali fondamentali nel ragionamento multi-step degli LLM che gli attuali sistemi di produzione ignorano. Decoupling (disaccoppiando) il riutilizzo della KV dall'identità lessicale (ASKS) e sfruttando il fenomeno del "collasso dell'entropia" negli strati successivi del transformer (CGEE), RKSC riduce significativamente i costi energetici e temporali dello scaling dell'inferenza.

Gli autori sottolineano che questi guadagni sono ottenuti senza l'overhead computazionale dell'addestramento e senza la complessità architettonica dei metodi di decoding speculativo. Il framework è presentato come un'ottimizzazione pratica e "drop-in" che rende il ragionamento allo stato dell'arte più accessibile su hardware limitati, riducendo la computazione ridondante nelle pipeline multi-branch. I risultati suggeriscono che, per i compiti di ragionamento multi-step, una parte significativa della verifica e della computazione di prefill è strutturalmente ridondante e può essere eliminata o accorciata in sicurezza in base alla similitudine semantica e alle metriche di confidenza.