SplitZip: Ultra Fast Lossless KV Compression for Disaggregated LLM Serving

Sintesi Tecnica: SplitZip

Definizione del Problema
I sistemi contemporanei di serving di Large Language Models (LLM) adottano sempre più la disaggregazione prefill-decode (PD) per ottimizzare l'utilizzo delle risorse, separando la fase di prefill (vincolata dalla computazione) dalla fase di decode (vincolata dalla memoria). Tuttavia, questa architettura introduce un collo di bottiglia critico: il trasferimento della cache Key-Value (KV) dai worker di prefill ai worker di decode attraverso nodi fisici distinti. In scenari di produzione reali, dove prefill e decode sono separati da link inter-nodo o inter-cluster (spesso con larghezze di banda di 50–100 GB/s, significativamente inferiori all'RDMA intra-rack), questo trasferimento diventa un fattore di latenza dominante, specialmente per carichi di lavoro con input lunghi o agentici.

Le soluzioni esistenti affrontano un compromesso:

• I metodi lossy (pruning, quantizzazione, ritenzione selettiva) offrono una compressione maggiore ma alterano gli output del modello, richiedendo una validazione per ogni deployment e rischiando l'accumulo di errori in contesti lunghi.
• I metodi lossless (ad esempio, codifica Huffman, librerie GPU generiche come nvCOMP) sono spesso inadatti a questo scenario. Essi mirano alla compressione dei pesi offline, operano su CPU o utilizzano una codifica a lunghezza variabile che crea pattern di accesso alla memoria irregolari e dipendenze sequenziali, risultando in throughput di compressione/decompressione troppo lenti per tenere il passo con la generazione KV durante il prefill.

Metodologia: SplitZip
Gli autori propongono SplitZip, uno schema di compressione lossless, compatibile con le GPU, progettato specificamente per il trasferimento della cache KV nel serving disaggregato di LLM. SplitZip preserva i tensori KV bit a bit senza alterare l'esecuzione o la semantica del modello.

• Insight Centrale: Il documento osserva che nelle attivazioni KV in BF16 (e FP8), la ridondanza è concentrata nel campo dell'esponente, mentre il segno e la mantissa presentano un'alta entropia e scarsa comprimibilità. L'analisi statistica su diverse famiglie di modelli (Qwen, Llama, Phi) rivela che l'entropia dell'esponente è tipicamente di 3–4 bit, e i 16 valori di esponente più frequenti coprono oltre il 99% di tutte le voci dell'esponente.
• Strategia di Codifica:
  • Codifica a Lunghezza Fissa: SplitZip utilizza un codebook Top-16 calibrato offline per mappare i valori di esponente più frequenti in codici a lunghezza fissa di 4 bit. Due codici da 4 bit vengono impacchettati in un singolo byte.
  • Flusso di Escape: Gli esponenti che cadono al di fuori del set Top-16 sono trattati come "escape". Per questi valori rari, il sistema registra la posizione dell'elemento (16 bit) e il valore grezzo dell'esponente (8 bit) in un flusso di correzione sparsa.
  • Preservazione di Segno/Mantissa: I bit di segno e mantissa sono memorizzati senza perdite come un flusso a 8 bit senza compressione.
  • Ottimizzazione: Il processo di codifica utilizza un kernel vettorializzato "Quad64" fuso per massimizzare il throughput della GPU, evitando l'overhead dell'istogramma online tramite l'uso di un codebook pre-calibrato.
• Strategia di Decodifica:
  • Parallelismo: La decodifica spacchetta i byte in codici a 4 bit e li mappa tramite una tabella di lookup in parallelo.
  • Correzione: Il percorso denso ricostruisce la maggior parte dei valori. Il flusso di escape sparso viene quindi applicato per sovrascrivere i pochi esponenti errati con i loro valori grezzi esatti, garantendo una ricostruzione bit-exact.
  • Scelta di Design: Gli autori scelgono esplicitamente un design Top-16 (utilizzando tutti i 16 possibili codici a 4 bit per i valori comuni) rispetto a un design Top-15-più-sentinella. Sebbene l'approccio con sentinella risparmi metadati, introduce un controllo di flusso irregolare che degrada severamente il throughput di decodifica della GPU. SplitZip privilegia un'esecuzione GPU uniforme e priva di divergenza.

Contributi Chiave

1. Caratterizzazione della Ridondanza KV: Il documento dimostra che la ridondanza dell'esponente in BF16, precedentemente notata nei pesi dei modelli, è prevalente e costante anche nelle attivazioni KV in architetture dense, MoE e ibride.
2. Codec ad Alto Throughput: SplitZip raggiunge un throughput di codifica di 613,3 GB/s e di decodifica di 2181,8 GB/s su attivazioni BF16 reali. Questi tassi superano considerevolmente i precedenti compressori lossless (ad esempio, 68,9× più veloce nella codifica rispetto a Falcon, 14,1× più veloce nella decodifica rispetto a nvCOMP Cascaded).
3. Integrazione End-to-End: Il sistema è integrato nel framework di serving SGLang (utilizzando Mooncake per il trasporto). Fornisce fino a un incremento di velocità di 1,32× nel trasferimento della cache KV in BF16, un incremento di 1,30× nel Time-To-First-Token (TTFT) e un aumento di 1,23× del throughput delle richieste per carichi di lavoro a contesto lungo.
4. Estensibilità: L'approccio si estende alle cache KV in FP8 (E5M2), fornendo un incremento di compressione di ulteriori 1,14× rispetto alla memorizzazione nativa FP8.

Risultati

• Performance del Codec: SplitZip mantiene un rapporto di compressione di 1,324× per BF16 (avvicinandosi al limite teorico di 1,33× quando i tassi di escape scendono allo 0,16%). Supera significativamente i baseline in termini di throughput mantenendo la fedeltà lossless.
• Latenza di Trasferimento: Negli esperimenti di trasferimento Mooncake, SplitZip riduce il tempo di trasferimento di 1,22× fino a 1,32× a seconda della lunghezza della sequenza e della dimensione del batch. Il beneficio scala con la lunghezza della sequenza poiché il trasferimento diventa limitato dalla larghezza di banda; per sequenze lunghe (es. 64K), l'overhead del codec diventa trascurabile (solo circa il 7% del tempo totale) e la riduzione della dimensione del payload abbassa direttamente la latenza.
• Studi di Ablazione:
  • Top-16 vs. Top-8: Il Top-16 è preferito; il Top-8 riduce la copertura (92% contro 99,8%), aumenta i tassi di escape e rallenta la decodifica di 3× a causa di un impacchettamento meno efficiente.
  • Calibrazione: Un singolo codebook calibrato offline generalizza bene attraverso diversi dataset (WikiText, HumanEval, MMLU) e layer del modello, evitando la massiccia penalità di throughput derivante da una calibrazione dinamica, per-tensore o per-token.
  • Compromesso sui Metadati: La memorizzazione esplicita delle posizioni di escape (Top-16) è preferita rispetto alla rilevazione tramite sentinella (Top-15) perché preserva percorsi di decodifica uniformi, risultando in un throughput di decodifica 3,5× superiore.

Significatività
Il documento afferma che SplitZip affronta un collo di bottiglia specifico e crescente nel serving disaggregato di LLM: l'overhead di comunicazione del trasferimento della cache KV. Sfruttando le proprietà statistiche specifiche degli esponenti in virgola mobile e progettando un codec che mappa nativamente il parallelismo delle GPU, SplitZip consente un trasferimento lossless efficiente senza sacrificare l'accuratezza del modello. Permette al serving PD di scalare efficacementmente su standard link di rete (non solo su fabric RDMA ad alta larghezza di banda) e supporta carichi di lavoro a contesto lungo e ragionamento multi-step dove la compressione lossy è spesso inaccettabile. Gli autori pongono SplitZip come una tecnica complementare che può essere stratificata sopra altre ottimizzazioni (come la quantizzazione) per ridurre ulteriormente i costi di banda.