QuantMLA: quantizzazione a doppio percorso allineata alla funzione per la cache KV di MLA a pochi bit

Abstract degli autori

La Multi-Head Latent Attention (MLA) consente un’attenzione multi-testa espressiva con cache compatte per il percorso del contenuto e quello RoPE disaccoppiato, ma la memoria occupata dalla cache cresce comunque in modo lineare con la lunghezza del contesto e la dimensione del batch. In questo lavoro, definiamo un modello sistematico degli errori di quantizzazione dei due percorsi di MLA, caratterizzandone i diversi effetti sulla distorsione dell’output dell’attenzione e spiegando la marcata amplificazione degli errori nel percorso RoPE. Sulla base di questa analisi, presentiamo QuantMLA, un framework allineato alla funzione per la quantizzazione a pochi bit dei due percorsi. Deriviamo spazi di trasformazione specifici per ciascun percorso che preservano il calcolo a precisione piena e possono essere incorporati integralmente nei parametri del modello offline, eliminando il costo delle trasformazioni online. All’interno di questi spazi, QuantMLA apprende trasformazioni specifiche per ciascun percorso con obiettivi allineati alla funzione: la ricostruzione dell’output dell’attenzione coglie gli errori accoppiati di corrispondenza e aggregazione del percorso del contenuto, mentre la ricostruzione posizionale QK preserva la componente dei logit dell’attenzione indotta da RoPE e consente di stabilire un limite teorico alla distorsione dell’output. Su quattro famiglie di modelli MLA, QuantMLA consente, a quanto ci risulta, la prima memorizzazione congiunta in INT4 delle cache del contenuto e RoPE di cui sia stata data notizia, con una perdita minima di accuratezza. Comprimendo ulteriormente la cache del contenuto a INT2 e mantenendo la cache delle chiavi RoPE a INT4, le prestazioni restano competitive su benchmark impegnativi di ragionamento e codice. Sviluppiamo un kernel nativo per l’attenzione MLA a pochi bit che integra il disimballaggio e la dequantizzazione direttamente nel calcolo dell’attenzione. La disposizione fisica della cache offre una compressione di 3,59x con un contesto di 128K, mentre un carico di lavoro di serving che mette sotto pressione la cache raggiunge un throughput di output sull’intero job 5,168x superiore a quello di BF16. Il codice sarà reso disponibile dopo l’accettazione.