STEPQuant: quando e dove contano gli errori nella quantizzazione degli stati ricorrenti della regola delta

Abstract degli autori

L’attenzione lineare sostituisce le cache KV, che crescono nel tempo, con stati ricorrenti di dimensione fissa. Questi stati persistenti possono però occupare una quantità considerevole di memoria quando si gestiscono richieste simultanee. Quantizzare direttamente gli stati ricorrenti a bassa precisione provoca spesso un grave calo dell’accuratezza, perché gli errori di quantizzazione si propagano attraverso i successivi aggiornamenti dello stato. Scopriamo che l’impatto di questi errori dipende da due dimensioni complementari: sul piano temporale, gli errori nella memoria di lunga durata possono persistere per molti passaggi di decodifica; sul piano spaziale, gli errori nelle diverse righe delle chiavi influiscono in modo diverso sugli output del modello, mentre l’ampiezza dei valori dello stato varia notevolmente sia tra le righe sia tra le colonne. Sulla base di queste osservazioni, proponiamo STEPQuant, un framework di quantizzazione post-addestramento spazio-temporale per gli stati ricorrenti della regola delta. STEPQuant assegna la precisione in funzione dell’entità dell’errore e della durata della memoria, e adatta congiuntamente i fattori di scala delle righe delle chiavi e delle colonne dei valori in base alle distribuzioni degli stati e all’impatto delle righe delle chiavi sull’errore nell’output. Esperimenti su Qwen3.8-27B e Kimi-Linear-48B-A3B-Instruct, condotti su benchmark di generazione sia lunga sia breve, mostrano che STEPQuant raggiunge un’accuratezza molto vicina a quella degli stati FP32 con un budget nominale di 6 bit e supera la quantizzazione uniforme INT8 nella configurazione a 4 bit. Integrato in SGLang con kernel GPU ottimizzati, STEPQuant a 6 bit ottiene una compressione degli stati ricorrenti superiore a 5 volte e riduce la memoria totale necessaria per il serving fino al 68,7%. Il nostro codice è disponibile all’indirizzo https://github.com/Dreamer-Toby/STEPQuant.