Flash-dLLM: caching KV consapevole dell’I/O e decodifica parallela per modelli linguistici di grandi dimensioni a diffusione veloci ed efficienti in termini di memoria

Abstract degli autori

I modelli linguistici di grandi dimensioni a diffusione (dLLM) sono emersi di recente come una promettente alternativa ai modelli linguistici di grandi dimensioni autoregressivi, grazie alla generazione di testo non autoregressiva. Tuttavia, il loro impiego pratico è ancora limitato dall’inferenza inefficiente, dovuta in larga misura alla mancanza di meccanismi efficaci di caching Key-Value (KV) e di decodifica parallela scalabili. I metodi di accelerazione esistenti studiano in genere il caching KV e la decodifica parallela separatamente, trascurando i colli di bottiglia di I/O che si verificano quando il riutilizzo della cache e la verifica parallela dei token vengono applicati congiuntamente. In questo lavoro presentiamo $\textbf{Flash-dLLM}$, un framework di accelerazione dell’inferenza per dLLM veloce ed efficiente in termini di memoria, che non richiede addestramento. Flash-dLLM individua innanzitutto nell’I/O della memoria della GPU un collo di bottiglia predominante nell’inferenza dei dLLM con cache KV e lo affronta con un kernel fuso per la cache KV, consapevole dell’I/O, che riduce gli spostamenti di memoria ridondanti. A partire da questo meccanismo di caching ottimizzato, Flash-dLLM propone inoltre una strategia efficiente di decodifica basata sulla cache KV, con generazione e verifica di proposte: il dLLM stesso funge sia da generatore di proposte sia da verificatore, senza bisogno di un modello ausiliario. Questo approccio unificato consente una decodifica più veloce, preservando la qualità della generazione e migliorando la scalabilità a sequenze più lunghe e batch di dimensioni maggiori. Esperimenti approfonditi su benchmark di ragionamento matematico e generazione di codice dimostrano che Flash-dLLM supera costantemente i metodi di accelerazione dei dLLM esistenti, all’avanguardia, sia in termini di velocità d’inferenza sia di efficienza della memoria. In particolare, raggiunge accelerazioni di $5.1\times$ e $11.0\times$ rispetto al precedente baseline più performante, Elastic-Cache, rispettivamente su GSM8K e HumanEval.