Riepilogo Tecnico: Quando la Quantizzazione è Gratuita

Enunciato del Problema

L'inferenza di transformer a lungo contesto è tipicamente vincolata dall'impronta di memoria e dai requisiti di larghezza di banda della cache Key-Value (KV). All'aumentare delle lunghezze del contesto (ad esempio, 128K token), i tensori KV accumulati in fp16 possono superare la dimensione dei pesi del modello stesso, creando un collo di bottiglia in cui la larghezza di banda della memoria, piuttosto che la potenza di calcolo, limita la velocità di decodifica.

La pratica industriale standard affronta questo problema tramite strategie "ruota-poi-quantizza" (ad esempio, QuaRot, SpinQuant), che applicano rotazioni ortogonali casuali per gaussianizzare le attivazioni prima di quantizzarle in pochi bit (int4/int8). Tuttavia, questo approccio è tradizionalmente inquadrato come un compromesso qualità-latenza: il sovraccarico computazionale della rotazione e della de-rotazione, combinato con gli errori di quantizzazione, comporta tipicamente un aumento del 5–20% della latenza di decodifica, anche se si risparmia memoria.

Gli autori ipotizzano che questo compromesso si inverta su Apple Silicon (architettura di memoria unificata). La loro ipotesi è che i risparmi di larghezza di banda derivanti da una cache compressa di 3 volte possano superare il costo computazionale del kernel di quantizzazione, a condizione che il kernel sia altamente ottimizzato e fuso.

Metodologia

1. Trasformata di Fourier Randomizzata per Segno (SRFT)

Il paper adotta la SRFT (Sign-Randomized FFT) come meccanismo di rotazione, definita come $SRFT(x) = \text{pack}(F \cdot \text{diag}(s) \cdot x)$, dove $F$ è la DFT unitaria e $s$ è un vettore di segni casuali fissi.

• Razionale: Sebbene le trasformate di Hadamard (SRHT) siano comuni, la SRFT offre flessibilità per dimensioni non potenze di due e si allinea con le operazioni di moltiplicazione di matrici 16×16 (AMX, Tensor Cores).
• Risultato: Gli autori dimostrano che SRFT e SRHT sono statisticamente indistinguibili in termini di qualità della perplessità (PPL) su cinque modelli (SmolLM2, Qwen2.5, Gemma-3) e vari budget di bit.

2. Kernel Metal Fuso

Per eliminare il sovraccarico di dispatch, gli autori implementano un singolo kernel di calcolo Metal fuso che esegue l'intera pipeline:

1. Caricamento e Moltiplicazione per Segno: Applicazione del diagonale di segni casuali nel registro.
2. FFT: Esecuzione di $\log_2(d/2)$ passaggi a farfalla radix-2 di Stockham.
3. Imballaggio (Packing): Scompattamento della simmetria hermitiana in una rappresentazione reale.
4. Ridimensionamento e Quantizzazione: Esecuzione del ridimensionamento $\lambda$ per canale (appreso o calibrato), riduzione del massimo assoluto per gruppo e arrotondamento a int4/int8 con segno.
5. Imballaggio a Nibble: Imballaggio di due valori int4 in un singolo byte utilizzando shuffle SIMD.

Questo kernel è esposto tramite torch.mps.compile e integrato come sottoclasse Cache di HuggingFace (SRFTInt4Cache).

3. Strategie di Ridimensionamento Avanzate

• Per Canale + Per Gruppo: Per modelli con dimensioni di testa piccole (ad esempio, Qwen2.5-1.5B, $d=128$), il ridimensionamento per token fallisce catastroficamente a causa di coordinate outlier. Gli autori combinano il ridimensionamento per canale (una scala per coordinata) con il ridimensionamento per gruppo (gruppi di 16) per recuperare la qualità.
• Rotazioni Apprese: Il paper indaga la calibrazione post-allenamento in cui vengono aggiunte una rotazione $R$ e una scala $\lambda$ apprendibili. Una scoperta chiave è che apprendere $R+\lambda$ senza la base SRFT fissa riduce significativamente di più l'MSE di calibrazione (84,9% vs 50,3%) ma produce una PPL a valle peggiore, suggerendo che la SRFT casuale fissa agisce come un regolarizzatore necessario.
• Rotazioni di Householder: A $d=256$, le rotazioni di Householder con $k=d/2$ riflettori raggiungono una quantizzazione praticamente senza perdita ($\Delta PPL \approx 0.000$) con la metà dei parametri delle rotazioni Cayley complete.

4. Percorso CPU (AMX)

Un percorso CPU complementare utilizza il coprocessore matriciale AMX di Apple tramite cblas sgemm. Raggruppando DFT di lunghezza 8 come moltiplicazioni di matrici reali 16×16, gli autori raggiungono fino al 91% del picco AMX single-core, offrendo un'alternativa a bassa latenza per regimi con batch piccoli dove il sovraccarico di dispatch GPU è dominante.

Risultati Chiave

Prestazioni (Latenza)

Su Apple M1 (16 GB di memoria unificata), la cache KV int4 fusa supera l'fp16 nella latenza di decodifica end-to-end:

• Gemma-3 1B: Esegue 3–8% più velocemente (sovraccarico di latenza negativo) su prefissi da 256 a 4096 token.
• Qwen2.5-1.5B: Esegue 0,7–2,6% più velocemente su contesti brevi (≤1K token).
• Meccanismo: Il sovraccarico del kernel (~25 ns/vec) è inferiore ai risparmi di larghezza di banda della memoria ottenuti trasferendo 3 volte meno dati per ogni passo di decodifica.

Qualità (Perplessità)

• SRFT vs. SRHT: Nessuna differenza statisticamente significativa in $\Delta PPL$ su tutti i modelli testati.
• Qwen2.5-1.5B: Il kernel fuso con ridimensionamento per canale + per gruppo riduce il $\Delta PPL$ a 4 bit da un catastrofico +7975 (baseline per token) a +638,6 (una riduzione di 12,5 volte), corrispondendo all'implementazione di riferimento Python.
• Gemma-3 1B: Il $\Delta PPL$ a 4 bit è migliorato da +18,7 (per token) a +3,6 con il kernel scalato fuso.
• Regime Senza Perdita: A $d=256$, le rotazioni di Householder apprese raggiungono $\Delta PPL = 0.000$.

Efficienza di Memoria

• Compressione: Raggiunge una compressione di memoria persistente da 3,0× a 3,56× sui livelli di attenzione completa.
• Impatto Reale: Per un contesto di 128K su Llama-3-70B, la cache int4 riduce la memoria da ~41 GB (fp16) a ~10,9 GB.

Significato e Affermazioni

Il paper afferma che la quantizzazione è "gratuita" (positiva per il throughput) sull'hardware a memoria unificata, invertendo l'assunzione industriale standard secondo cui la quantizzazione a pochi bit comporta sempre una penalità di latenza.

1. Inversione del Compromesso: Il contributo principale è dimostrare che su Apple Silicon, i risparmi di larghezza di banda derivanti da una cache compressa di 3 volte superano il costo computazionale del kernel di quantizzazione. Questo rende la cache KV int4 strettamente superiore all'fp16 per inferenze a lungo contesto sensibili alla latenza su laptop e dispositivi edge.
2. Design Consapevole dell'Hardware: Il lavoro evidenzia che la scelta della rotazione (FFT vs. Hadamard) riguarda meno la superiorità matematica e più l'allineamento hardware (ad esempio, la struttura radix-8 della FFT che corrisponde alle tessere 16×16 di AMX).
3. Separazione tra Calibrazione e PPL: Il paper isola un fenomeno nuovo in cui la minimizzazione dell'MSE di calibrazione non garantisce una migliore PPL a valle, suggerendo che la rotazione casuale fissa fornisce un prior di regolarizzazione necessario.
4. Fattibilità di Deployment: Fondendo l'intera pipeline in un singolo kernel Metal e integrandolo con HuggingFace, gli autori dimostrano che una quantizzazione aggressiva a 4 bit può essere deployata senza sacrificare la qualità o la velocità del modello, a condizione che l'architettura hardware (memoria unificata) e l'ottimizzazione del kernel (fusione) siano allineate.

Gli autori concludono che, mentre il compromesso standard di quantizzazione vale per le GPU con HBM separata e alti sovraccarichi di dispatch, esso si inverte per il regime a memoria unificata in cui operano l'inferenza su laptop e dispositivi edge.