La quantizzazione della KV cache nei modelli linguistici non dipende soltanto dal numero di bit utilizzati, ma anche dall’asse lungo il quale vengono raggruppati i valori per calcolare scale factor e zero point. Nei test condotti su Llama-2-13B con quantizzazione a 2 bit e gruppi di 32 elementi, mantenendo invariati modello, precisione e benchmark, la sola scelta dell’asse porta il punteggio CoQA da 2,88 a 63,53, rispetto a 66,37 in piena precisione. La configurazione efficace utilizza una quantizzazione per canale per le chiavi e per token per i valori; invertire una delle due scelte riduce sensibilmente la qualità, mentre invertire entrambe può rendere il modello praticamente inutilizzabile.
La KV cache conserva durante la generazione le proiezioni key e value dei token già elaborati, evitando di ricalcolarle a ogni nuovo passo. La sua dimensione cresce linearmente con la lunghezza del contesto e con il batch size, fino a diventare una delle principali voci di consumo della memoria GPU. Nell’analisi KVQuant su LLaMA-7B, con sequenze da 512 token i pesi rappresentano circa il 98% della memoria e la cache soltanto il 2%, mentre con un contesto da 128.000 token il rapporto si ribalta: circa il 16% è occupato dai pesi e l’84% dalla KV cache. Per OPT-175B, con batch size 512 e prompt da 512 token, la cache può raggiungere circa 1,2 TB. La riduzione della KV cache non interviene quindi soltanto sulla capacità disponibile, ma anche sulla velocità di generazione, perché per ogni nuovo token la GPU deve leggere dalla memoria l’intera cache e durante questo trasferimento le unità di calcolo possono rimanere inutilizzate.
Nella quantizzazione uniforme intera, il range dei valori appartenenti a ciascun gruppo viene suddiviso nel numero di livelli rappresentabili dal formato scelto. A 2 bit sono disponibili soltanto quattro livelli, quindi la presenza nello stesso gruppo di un valore estremamente superiore agli altri allarga lo step di quantizzazione e riduce contemporaneamente la precisione di tutti gli elementi associati alla stessa scala. La struttura dei gruppi diventa quindi determinante: scegliere l’asse significa stabilire quali elementi condivideranno lo stesso errore di quantizzazione. Nel caso delle chiavi, alcuni canali contengono sistematicamente attivazioni di ampiezza molto elevata. In Mixtral 8x7B sono state osservate magnitudini massime vicine a 7.000 rispetto a un valore mediano delle feature nell’ordine di 0,3. Questi outlier tendono a concentrarsi in dimensioni specifiche e relativamente stabili rispetto all’input e sono collegati anche al comportamento degli attention sink.
Per questo motivo le chiavi beneficiano della quantizzazione per canale. Raggruppando lungo la dimensione dei token, i canali contenenti outlier finiscono per ampliare il range di numerosi gruppi e deteriorare la precisione delle componenti ordinarie; lavorando invece per canale, le dimensioni anomale rimangono maggiormente isolate. Nei risultati KIVI su Llama-2-13B, mediati tra layer e teste di attenzione, l’errore di ricostruzione delle chiavi passa da 13,67 con quantizzazione per token a 4,55 con quella per canale. La differenza è ancora più marcata sull’errore dei punteggi di attenzione, che scende da 47,00 a 9,60. La quantizzazione delle chiavi lungo il token produce quindi un errore sui punteggi di attenzione circa cinque volte superiore.
I valori richiedono invece il trattamento opposto. La value cache non presenta lo stesso schema di outlier concentrati in canali fissi e, osservando soltanto la distribuzione dei tensori memorizzati, la quantizzazione per canale potrebbe apparire persino preferibile. L’errore di ricostruzione misurato direttamente sul tensore è infatti pari a 3,73 con quantizzazione per canale contro 4,57 per token. Il risultato cambia però completamente quando si misura l’errore nel punto in cui i valori vengono effettivamente utilizzati. L’output dell’attenzione è una somma pesata dei value vector dei diversi token, dove i coefficienti sono determinati dai punteggi softmax. In questa operazione KIVI registra un errore relativo di 3,55 con quantizzazione per token e di 49,89 con quella per canale, oltre quattordici volte superiore nonostante quest’ultima risultasse migliore secondo il semplice errore di ricostruzione.
La differenza deriva dalla forte sparsità dell’attenzione, misurata nell’esperimento al 84,3%. Una parte consistente dell’output è determinata da pochi token rilevanti, mentre molti altri ricevono pesi prossimi allo zero. Quantizzare ciascun token separatamente confina quindi l’errore all’interno della relativa rappresentazione: eventuali imprecisioni associate ai token poco importanti vengono quasi annullate dai bassi coefficienti di attenzione. La quantizzazione per canale distribuisce invece lo stesso scale factor tra token differenti e può propagare gli errori anche alle rappresentazioni che ricevono maggiore attenzione. Di conseguenza, per valutare correttamente una tecnica di compressione della KV cache non è sufficiente confrontare il tensore quantizzato con quello originale, ma occorre misurare l’errore dopo l’operazione che consuma quel tensore oppure direttamente sulle prestazioni finali del modello.
Un ulteriore problema riguarda le Rotary Position Embeddings, RoPE, che ruotano coppie di canali in funzione della posizione relativa dei token. Questa trasformazione può disperdere nei canali adiacenti gli outlier che rendevano efficace la quantizzazione per canale delle chiavi. KVQuant affronta il problema quantizzando le chiavi prima dell’applicazione di RoPE e applicando la rotazione dopo la dequantizzazione. Combinando questo ordine delle operazioni con quantizzazione per canale, tipi numerici non uniformi e isolamento di una piccola quota di outlier, viene riportato un degrado della perplexity inferiore a 0,1 a 3 bit e la possibilità di gestire con LLaMA-7B contesti fino a un milione di token su una singola GPU A100 da 80 GB. Altri risultati mostrano inoltre un aumento del 145% dell’errore di quantizzazione dopo l’applicazione di RoPE e indicano che gli outlier possono trovarsi in canali differenti tra le varie attention head, rendendo più efficaci rotazioni adattate alla singola testa rispetto a una matrice condivisa.
La quantizzazione per canale introduce però anche un problema operativo durante il decoding, perché le statistiche di un canale dipendono da token che non sono ancora stati generati e non è quindi possibile determinarne immediatamente la scala completa. KIVI utilizza per questo un residual buffer che mantiene fino a 128 token recenti in piena precisione e procede alla quantizzazione in gruppi una volta raccolto un numero sufficiente di elementi. Il buffer incide direttamente sulla qualità: su GSM8K con Llama-2-7B, il modello in piena precisione ottiene 13,50 punti; la quantizzazione completa a 2 bit con gli assi corretti scende a 5,76, mentre mantenendo gli stessi assi e la stessa precisione ma aggiungendo la finestra dei token recenti in piena precisione il risultato risale a 12,74. Questa soluzione consente di conservare con precisione proprio la parte più recente del contesto, frequentemente utilizzata nei passaggi successivi dei problemi multi-step.
Con questa architettura KIVI riporta una riduzione fino a 2,6 volte del picco di memoria per Llama-2-7B, batch size fino a quattro volte superiori e un incremento del throughput compreso tra 2,35 e 3,47 volte in scenari di servizio reali. L’evoluzione della ricerca comprende tecniche come InnerQ, che incorpora la normalizzazione delle chiavi per canale direttamente nei pesi di key e query durante il prefill per evitare overhead aggiuntivo a runtime, mantenendo inoltre finestre ad alta precisione sia per i token recenti sia per gli attention sink. Altri approcci studiano invece la possibilità di non conservare integralmente key e value, memorizzando rappresentazioni più compatte dalle quali ricostruirli quando necessario.
La scelta della configurazione deve infine considerare anche parametri diversi dall’accuratezza dei benchmark. Sono stati osservati fenomeni di degradazione dell’allineamento in seguito alla quantizzazione della KV cache anche in ambienti di serving vLLM con cache FP8; specifici protocolli di recupero senza ulteriore training hanno permesso di ripristinare fino al 97% dell’allineamento perso. Nel complesso, la configurazione della KV cache richiede quindi quantizzatori distinti per chiavi e valori, quantizzazione delle chiavi prima di RoPE, una finestra recente mantenuta ad alta precisione e procedure di validazione basate sull’output dell’attenzione o sulle prestazioni effettive del task, evitando di considerare il solo errore di ricostruzione del tensore come indicatore sufficiente della qualità.
Questo articolo è stato redatto con il supporto di strumenti di intelligenza artificiale (AI)