WUSH-KV: quantizzazione della cache KV con trasformazioni adattive ai dati

Abstract degli autori

I costi di memoria e di larghezza di banda della cache KV aumentano con la lunghezza del contesto e la dimensione del batch, limitando l’efficienza dell’inferenza su contesti lunghi. Per affrontare questo ostacolo, presentiamo WUSH-KV per la quantizzazione a pochi bit della cache KV. WUSH-KV adatta WUSH, che costruisce una trasformazione basata sui dati a partire dalle statistiche del secondo ordine di entrambi i fattori di un prodotto matriciale, per ridurre l’errore di quantizzazione. WUSH-KV usa dati di calibrazione per costruire trasformazioni separate per le chiavi e i valori: la trasformazione dei valori viene incorporata nei pesi del modello, mentre quella delle chiavi viene applicata dopo RoPE. Le trasformazioni possono essere abbinate a quantizzatori con clipping. Per uno di questi quantizzatori, QuEST INT, mostriamo che, sotto ipotesi poco restrittive, la trasformazione WUSH è quasi ottimale. Con questo quantizzatore, WUSH-KV riduce l’errore di ricostruzione per livello e ottiene la perplexity end-to-end più bassa tra le trasformazioni testate. Per la valutazione end-to-end, integriamo WUSH-KV in SGLang usando una quantizzazione affine con clipping basato sui percentili, nello stile di OSCAR. A 2 bit, WUSH-KV ottiene prestazioni paragonabili o superiori a quelle della trasformazione OSCAR in tutti i modelli e i task downstream valutati.