Riepilogo Tecnico: Spherical KV

1. Definizione del Problema

L'inferenza a lungo contesto nei Large Language Models (LLM) è sempre più limitata dalla cache Key-Value (KV) piuttosto che dal throughput aritmetico. Man mano che la lunghezza del contesto ($T$) cresce, l'impronta di memoria residente scala linearmente con $T$, la dimensione del batch ($B$), i livelli ($L$) e le teste ($H$). Negli ambienti di serving in produzione che utilizzano la memoria paginata (es. PagedAttention), il collo di bottiglia dominante si sposta dai Floating Point Operations (FLOPs) al traffico di High Bandwidth Memory (HBM). Nello specifico, la fase di decoding diventa limitata dallo streaming ripetuto di pagine KV storiche dalla memoria del dispositivo alle unità di calcolo.

Le soluzioni esistenti tentano di mitigare questo problema attraverso tre leve:

1. Ritenzione/Evizione: Memorizzare meno token (es. StreamingLLM, H2O).
2. Quantizzazione/Compressione: Memorizzare meno bit (es. KVQuant, KIVI).
3. Offloading: Memorizzare la KV altrove (es. CPU/Storage).

Tuttavia, questi metodi spesso non riescono ad affrontare pienamente il collo di bottiglia del percorso critico. Molte tecniche di compressione richiedono la ricostruzione di vettori densi da rappresentazioni compresse prima di calcolare i logit dell'attenzione (un percorso "decomprimi-e-poi-prodotto-scalare"). Questa ricostruzione reintroduce il traffico HBM e impedisce la fusione dei kernel, annullando i risparmi teorici di memoria in substrati di serving paged realistici. Inoltre, le semplici politiche di evizione spesso mancano della granularità necessaria per preservare informazioni critiche (es. evidenze recuperate) scartando al contempo stati a bassa utilità.

2. Metodologia: Spherical KV

Il paper propone Spherical KV, un metodo di inferenza a lungo contesto che tratta l'allocazione KV come un problema di rate-distortion (rapporto tasso-distorsione) basato sulla geometria dell'attenzione. Consiste in due componenti ortogonali progettati per operare all'interno di un substrato di serving paged e ragged: Angle-Domain Attention (ADA) e Rate–Distortion Retention (RDR).

2.1 Angle-Domain Attention (ADA)

ADA sostituisce la primitiva standard "decomprimi-e-poi-prodotto-scalare" con un approoccio nativo del calcolo che opera direttamente su codici sferici compatti.

• Parametrizzazione Sferica: Le chiavi vengono decomposte in un raggio scalare ($r = \|k\|$) e angoli direzionali ($\phi = \Phi(\hat{k})$).
• Similarità nel Dominio Angolare: Inveve di memorizzare vettori densi, ADA memorizza raggi quantizzati ($\tilde{r}$) e codici angolari impacchettati ($c_\theta$). Durante il loop caldo del decode, il kernel di attenzione calcola i logit direttamente da questi codici:
  $$\hat{\ell}(q, k) \approx \|q\| \cdot \tilde{r} \cdot f_b(c_\theta^q, c_\theta^k)$$
  dove $f_b$ è un'operazione di kernel fuso che calcola la similarità del coseno dai codici angolari.
• Invariante di Non Ricostruzione: Il vincolo di design critico è che il kernel di decode non materializza mai un vettore chiave denso ($\tilde{k} \in \mathbb{R}^d$) in HBM. Ciò garantisce che le riduzioni nei byte KV memorizzati si traducano direttamente in una riduzione del traffico HBM per token.

2.2 Rate–Distortion Retention (RDR)

RDR è un meccanismo di controllo che ottimizza congiuntamente quali stati trattenere e a quale precisione (tier) memorizzarli, soggetto a un budget di memoria rigoroso.

• Ottimizzazione Congiunta: A differenza dei metodi che trattano l'evizione e la quantizzazione separatamente, RDR risolve un problema di allocazione discreta rate-distortion. Per ogni stato token/testa/livello $i$, seleziona una decisione di residenza binaria ($z_i \in \{0, 1\}$) e un livello di precisione ($b_i \in \mathcal{B}$).
• Surrogato della Distorsione: Il controller utilizza un proxy di distorsione meccanicistico derivato dall'analisi dell'errore di ADA. Stima che l'errore del logit sia dominato dagli errori di approssimazione angolare ($\epsilon_\theta$) scalati dalle magnitudo della query e della chiave.
• Vincoli di Protezione: Per prevenire fallimenti catastrofici su informazioni critiche (es. evidenze recuperate, span di sicurezza, outlier), RDR impone vincoli di protezione rigidi ($I_{prot}(i)=1 \implies z_i=1, b_i=b_{max}$).
• Page Omogenee per Tier: Il controller raggruppa gli stati trattenuti in pagine dove tutti gli elementi condividono lo stesso livello di precisione. Ciò consente letture streaming a singolo passaggio guidate dall'header, evitando gather irregolari.

2.3 Integrazione di Sistema

Spherical KV opera su un normale substrato KV paged/ragged (compatibile con i layout in stile vLLM). La pipeline mappa le chiavi dense in codici sferici, applica RDR per filtrarli e categorizzarli per tier, e li memorizza in un formato compresso. Il kernel di decode consuma questi codici direttamente, mantenendo un invariante di "nessuna ricostruzione densa".

3. Contributi Chiave

1. Angle-Domain Attention (ADA): Una primitiva a livello di kernel che calcola i logit dell'attenzione direttamente da codici sferici compatti, eliminando la penalità di traffico HBM associata alla ricostruzione di chiavi dense.
2. Rate–Distortion Retention (RDR): Un controller unificato che alloca congiuntamente la residenza e i livelli di precisione sotto un budget fisso, dando priorità agli stati ad alta utilità e proteggendo quelli fragili.
3. Valutazione Orientata al Deployment: Il metodo è valutato non su cache contigue giocattolo, ma su un realistico substrato di serving paged/ragged, misurando throughput end-to-end, traffico HBM e stabilità comportamentale.
4. Spostamento della Frontiera di Pareto: Il paper dimostra che Spherical KV sposta la frontiera memoria-qualità-throughput, ottenendo un throughput più elevato a parità di footprint di memoria rispetto a KV densi e baseline forti.

4. Risultati Sperimentali

Gli autori hanno valutato Spherical KV su Llama-3.1-8B-Instruct, Qwen2.5-14B-Instruct e gpt-oss-20b attraverso lunghezze di contesto di 8K, 32K e 128K. I carichi di lavoro includevano modellazione del linguaggio (PG-19), QA ad alto recupero (LongBench: HotpotQA, 2WikiMultiHopQA) e rollout di agenti (AgentBench, ToolBench).

• Efficienza di Throughput e Memoria: Sotto vincoli di qualità corrispondenti (iso-quality), Spherical KV ha migliorato il throughput di decode di 1.55×–1.72× riducendo al contempo i byte KV residenti per token del 24–42%.
• Scalabilità del Contesto: I guadagni sono stati più pronunciati nel regime a 128K, dove la residenza KV e il traffico HBM dominano le prestazioni del sistema.
• Efficienza Realizzata dal Kernel: Il metodo ha ridotto con successo i byte HBM per token generato, confermando che i risparmi di memoria non sono stati compensati da overhead di ricostruzione.
• Stabilità: Il sistema ha dimostrato una degradazione fluida sotto budget più stretti e ha preservato le traiettorie di ragionamento e le sequenze di chiamata a strumenti nei carichi di lavoro agentici, evitando i fallimenti bruschi osservati nelle baseline di evizione rigida.

5. Significato e Rivendicazioni

Il paper afferma che Spherical KV affronta il collo di bottiglia fondamentale dell'inferenza a lungo contesto: il conflitto tra capacità di memoria e larghezza di banda della memoria. Trattando la compressione KV come un problema nativo del calcolo (dove la rappresentazione compressa è direttamente eseguibile dal kernel) piuttosto che solo come un problema di memorizzazione, il metodo ottiene un vero spostamento della frontiera operativa.

Gli autori sottolineano che:

• La compressione deve essere nativa del calcolo: Memorizzare semplicemente meno bit è insufficiente se il percorso di decode richiede una ricostruzione densa.
• La memoria deve essere allocata dove conta: Il controllo rate-distortion permette al sistema di investire precisione sugli stati ad alto impatto risparmiando su quelli a bassa utilità.
• La valutazione deve essere fedele al deployment: I guadagni devono sopravvivere a layout paged realistici, overhead di metadati e vincoli di kernel per essere considerati reali.

Spherical KV è presentato come un meccanismo per ridurre la residenza KV e il traffico HBM preservando l'efficienza di decode e la stabilità comportamentale, mirando specificamente ai regimi limitati dalla memoria dei moderni sistemi di serving di LLM.