OSCAR: Come Together AI Comprime la KV Cache a 2-bit per LLM

L'inferenza a lungo contesto sta diventando il nuovo standard per i Large Language Models (LLM), ma porta con sé una sfida critica: la crescita esponenziale della KV cache. Per rispondere a questa esigenza, Together AI ha rilasciato OSCAR (Offline Spectral Covariance-Aware Rotation), un sistema open source che permette di comprimere la KV cache a soli 2 bit mantenendo un'accuratezza quasi identica ai modelli originali, risolvendo finalmente i limiti della quantizzazione INT2.

Il problema della quantizzazione a 2 bit nelle KV Cache

Durante il processo di decodifica autoregressiva, la KV cache cresce in base alla lunghezza del contesto, alla dimensione del batch e alla profondità del modello. Con contesti da 100.000 token e decine di richieste simultanee, questa cache può arrivare a consumare la quasi totalità della memoria GPU disponibile. Comprimerla è l'unico modo per aumentare la densità di serving.

Tuttavia, spingere la precisione verso il formato INT2 (che offre solo quattro livelli rappresentabili) è storicamente difficile. Il motivo principale risiede negli outlier per canale: un piccolo sottoinsieme di canali contiene valori estremamente elevati che dominano i fattori di scala. In una quantizzazione INT2 standard, questo porta a sprecare l'intero range dinamico su pochi picchi, appiattendo tutti gli altri valori e distruggendo l'accuratezza del modello.

OSCAR: Rotazione Consapevole dell'Attenzione

L'intuizione fondamentale di OSCAR è che la rotazione applicata prima della quantizzazione debba derivare dalle statistiche di attenzione del modello, non solo dalla distribuzione grezza delle attivazioni KV.

Ottimizzazione per Chiavi (Keys) e Valori (Values)

• Per le Keys: L'errore che conta non è la ricostruzione euclidea, ma l'impatto sui logit dell'attenzione. OSCAR utilizza la covarianza dei query (Q^TQ) per identificare le direzioni in cui i query hanno maggiore energia, poiché queste amplificano l'errore di quantizzazione.
• Per i Valori: L'errore rilevante riguarda l'output finale dell'attenzione. OSCAR definisce una covarianza dei valori pesata dai punteggi di attenzione (score-weighted value covariance), proteggendo le direzioni attraverso cui l'errore si propaga maggiormente.

I tre fattori della rotazione

La rotazione finale di OSCAR si compone di tre elementi distinti:

1. UQ / US: Allinea i canali con le direzioni di importanza dell'attenzione.
2. HHad (Walsh-Hadamard): Equalizza l'importanza dei canali, comprimendo lo spettro dei valori.
3. Pbr (Permuted Bit-Reversal): Riordina i canali affinché ogni gruppo di quantizzazione riceva una rappresentazione bilanciata della gerarchia di importanza.

Architettura del Sistema di Serving e Layout di Memoria

OSCAR non è solo un concetto teorico, ma è integrato nello stack di serving di SGLang con piena compatibilità per la paged attention. Il sistema utilizza un layout a precisione mista per massimizzare la robustezza:

• Sink tokens: I primi 64 token sono conservati in BF16 per fungere da ancore di attenzione.
• Recent tokens: Gli ultimi 256 token prima della posizione corrente rimangono in BF16.
• History tokens: Tutto il resto della cronologia viene convertito in INT2 dopo la rotazione e il clipping.

Le operazioni di scrittura e lettura utilizzano kernel Triton fusi. Durante la scrittura, ogni token viene ruotato, limitato (clipping) e quantizzato. Durante la lettura, il kernel decomprime, de-quantizza e inverte la rotazione in un unico passaggio, eliminando il traffico di memoria supplementare.

Analisi delle Prestazioni

Il team di ricerca ha testato OSCAR su diversi modelli, tra cui Qwen3 e GLM-4.7, ottenendo risultati sorprendenti rispetto ai metodi concorrenti.

Mentre la quantizzazione INT2 ingenua (senza rotazione) porta il punteggio dei modelli Qwen3 a zero, OSCAR mantiene prestazioni competitive. Su GLM-4.7-FP8, un modello da 358 miliardi di parametri, OSCAR ha persino mostrato un leggero miglioramento rispetto alla baseline BF16 (+0.27 punti).

Velocità di decodifica (H100, Context 100K)

Il risparmio di memoria si traduce direttamente in velocità, poiché la decodifica a lungo contesto è limitata dalla banda passante della memoria (bandwidth-bound):

• Qwen3-4B: Accelerazione di 3.08x.
• Qwen3-8B: Accelerazione di 2.88x.
• GLM-4.7-FP8: Accelerazione di 2.83x.

Guida Rapida all'Implementazione

Per utilizzare OSCAR, Together AI mette a disposizione il RotationZoo, una collezione di matrici di rotazione pre-calcolate.

1. Installazione

bashCopia
pip install sglang[all] --upgrade
pip install triton

2. Avvio del Server

È possibile lanciare il server SGLang abilitando la modalità INT2 e puntando al percorso delle rotazioni pre-calcolate:

bashCopia
python -m sglang.launch_server \
  --model-path Qwen/Qwen3-8B \
  --kv-cache-dtype int2 \
  --oscar-rotation-path ./oscar_rotations/ \
  --oscar-sink-size 64 \
  --oscar-recent-size 256 \
  --tp 1 \
  --port 30000