Accelerazione dell'inferenza AI per la creazione 3D su Roblox

• Roblox ha implementato CUDA Graphs e il caching KV per velocizzare la generazione delle mesh 3D e ottenere iterazioni più reattive.
• Al momento del lancio, il modello Cube 3D era in grado di generare token in 7,8 millisecondi (rispetto ai precedenti 60,5 millisecondi) e oggetti completi in 4 secondi (rispetto ai precedenti 31 secondi). 

All'inizio di quest'anno, Roblox ha condiviso la prima funzionalità del nostro modello di base Cube 3D. Con Cube 3D, i creatori possono generare modelli e ambienti 3D direttamente da prompt di testo. Fin dall'inizio, abbiamo dato priorità all'ottimizzazione della latenza, riconoscendo che i tempi di generazione lenti interrompono quello che è intrinsecamente un processo iterativo. Prima del lancio di Cube 3D a marzo, avevamo già reso la fase di inferenza 7,8 volte più veloce e reattiva sia per gli sviluppatori che per gli utenti. 

Dal lancio, sono stati generati più di 578.000 oggetti in diverse esperienze di rilievo. Gli sviluppatori hanno inoltre espresso interesse nel consentire agli utenti di generare oggetti 3D all'interno delle esperienze tramite prompt di testo come "gatti", "hamburger", ecc. In particolare, Mic Up, un popolare gioco di ritrovo che utilizza la chat vocale, ha sfruttato Cube 3D per offrire ai giocatori un modo divertente e interattivo di generare oggetti. Nella loro implementazione, i giocatori possono aprire un menu a sinistra con funzionalità aggiuntive, tra cui un'icona AI. Dopo aver cliccato su quell'icona, i giocatori possono inserire un prompt di testo per generare un oggetto 3D. Per gli utenti, tempi di generazione più lunghi creano attrito, privandoli della magia di vedere le loro idee trasformate in 3D in tempo reale. 

Volevamo trasformare l'esperienza di generazione 3D da un'interazione "stop-and-wait" a qualcosa che risultasse reattivo e naturale, consentendo una rapida sperimentazione. La possibilità di aggiungere rapidamente oggetti a una scena è fondamentale per gli sviluppatori. Per accelerare Cube 3D, abbiamo prima profilato la pipeline di inferenza per identificare i colli di bottiglia delle prestazioni. Nonostante l'utilizzo di potenti GPU, abbiamo riscontrato un significativo tempo di inattività tra le operazioni.

I moderni framework di deep learning si affidano alla CPU per pianificare e avviare operazioni (o kernel) sulla GPU. La CPU prepara ogni operazione, la invia alla GPU e attende conferma prima di preparare l'operazione successiva. Questa attesa crea un collo di bottiglia nella pianificazione, in cui la GPU potrebbe rimanere inattiva mentre la CPU prepara il batch di lavoro successivo. Idealmente, vorremmo che la CPU anticipasse la GPU, preparando e mettendo in coda le operazioni in modo che la GPU abbia sempre del lavoro da svolgere.

Ciò è particolarmente problematico per i decodificatori autoregressivi nei modelli di tipo transformer come Cube 3D, che devono elaborare l'input e generare token in modo sequenziale. Questi modelli richiedono migliaia di operazioni individuali per una singola generazione e il sovraccarico computazionale si accumula ad ogni passo della sequenza.

"Volevamo costruire qualcosa che consentisse un'interazione quadridimensionale", ha affermato il vicepresidente dell'ingegneria Anupam Singh, spiegando perché Roblox ha scelto un approccio autoregressivo. "Non vogliamo solo costruire l'auto; vogliamo anche poter aprire la portiera dell'auto ed entrarci".

Ogni operazione ha comportato:

• Tempo di CPU per preparare ogni kernel
• Overhead derivante dall'avvio del kernel
• Tempo di esecuzione della GPU (il calcolo effettivo)
• Overhead di sincronizzazione durante il controllo del completamento

Nel caso di operazioni di piccole dimensioni che vengono eseguite rapidamente sulla GPU, questo overhead può incidere in modo determinante sul tempo di inferenza. La GPU potrebbe essere attivamente impegnata nel calcolo solo per una piccola frazione del tempo totale di inferenza.

Per risolvere questo collo di bottiglia, abbiamo sfruttato CUDA Graphs, una funzionalità che permette di registrare e riprodurre sequenze di operazioni GPU senza l'intervento della CPU. Il componente decodificatore autoregressivo dell'architettura di Cube 3D elabora i prompt di testo e genera token di forma attraverso un vettore a lunghezza fissa. 

Sebbene sia funzionalmente simile a un modello linguistico di grandi dimensioni (LLM) tradizionale, la nostra architettura di decodifica a doppio flusso presenta una differenza importante: utilizza due flussi di attenzione paralleli. Un flusso è dedicato ai token di condizione e l'altro ai token di forma. I motori di inferenza LLM disponibili in commercio non erano adatti alle nostre esigenze e avevamo bisogno di un'implementazione personalizzata su misura per la nostra specifica architettura.

Pensate ai grafici CUDA come alla registrazione di una macro per la GPU. Invece di emettere ogni comando singolarmente, la CPU registra un'intera sequenza di operazioni della GPU (il grafico) e avvia l'intero grafico con una singola istruzione della CPU. Questo approccio riduce drasticamente l'overhead di avvio del kernel, eliminando la necessità per la CPU di pianificare individualmente ogni operazione durante l'inferenza. Una volta avviato il grafico, la GPU esegue l'intera sequenza in modo autonomo, senza attendere ulteriori istruzioni.

I grafici CUDA presentano alcune limitazioni. Poiché la struttura del grafico deve essere determinata in anticipo, richiedono una dimensione di batch e dimensioni di input fisse. Ciò significa che è necessario creare grafici separati per ogni dimensione di batch o forma di input. Per il nostro caso d'uso con Cube 3D, questa limitazione era accettabile, poiché potevamo standardizzare il processo di inferenza attorno a dimensioni di input comuni.

Abbiamo dovuto adattare il nostro approccio per implementare i grafici CUDA per il nostro modello Cube 3D. Nei modelli di linguaggio di grandi dimensioni (LLM) tradizionali, le operazioni di attenzione vengono sempre eseguite con la stessa lunghezza di sequenza, fornendo una forma statica su cui lavorare. Tuttavia, nella nostra architettura personalizzata a doppio flusso, alcuni livelli di attenzione operano solo sulla lunghezza della sequenza, mentre altri operano su una combinazione di lunghezza della sequenza e della condizione.

Nonostante queste sfide, abbiamo ottenuto risultati notevoli dopo l'implementazione di CUDA Graphs. Utilizziamo il tempo per token di output (TPOT) per misurare il tempo di generazione di ciascun token durante l'inferenza. Dopo l'implementazione di CUDA Graphs, il nostro TPOT è migliorato da 60,5 millisecondi a 20,5 millisecondi, con un miglioramento di 2,9 volte. Il tempo di generazione complessivo è sceso del 66%, passando da 31 secondi a 10,5 secondi.

Per migliorare ulteriormente la latenza, abbiamo implementato il caching KV, una pratica standard nell'inferenza LLM che si è dimostrata altamente efficace in tutto il settore. 

Nei modelli basati su Transformer come Cube 3D, ogni generazione di token richiede il calcolo delle matrici chiave (K) e valore (V) basate su tutti i token generati in precedenza. Man mano che la sequenza si allunga, ricalcolare queste matrici per ogni token diventa sempre più inefficiente.

Il caching KV risolve questo problema:

1. Memorizzando le matrici K e V per tutti i token generati in precedenza
2. Calcolando le matrici K e V solo per i nuovi token
3. Aggiungendo queste nuove matrici ai valori memorizzati nella cache

Questo approccio elimina i calcoli ridondanti, riducendo il lavoro richiesto per ogni nuovo token. Ciò diventa particolarmente significativo man mano che la sequenza generata si allunga.

Il nostro approccio all'integrazione del caching KV con l'implementazione di CUDA Graph è stato simile all'inferenza LLM tradizionale. L'aggiunta del caching KV ha ridotto il nostro TPOT a soli 7,8 millisecondi. Il tempo di generazione complessivo è diminuito dell'87%, passando dai 31 secondi originali a soli 4 secondi. Questa significativa riduzione dei tempi rende lo strumento molto più efficace per i creatori che lo utilizzano.

Stiamo esplorando tecniche che riducano ulteriormente la latenza e migliorino l'esperienza utente, tra cui kernel ottimizzati, quantizzazione dei modelli per un'inferenza ancora più veloce, ottimizzazioni specifiche per l'hardware e generazione parallela di token.

Questo lavoro diventa più complesso quando ci espandiamo alla generazione e alla comprensione di scene complete, dove molti elementi 3D devono interagire tra loro all'interno di un layout. Vogliamo inoltre che gli oggetti e i mondi 3D che creiamo siano pienamente funzionali, in modo che le porte si aprano e si chiudano, le ruote girino, ecc. Per raggiungere questo obiettivo, abbiamo bisogno di una generazione e di un'iterazione rapide per scalare intere scene, oggetti pienamente funzionali e avatar. Siamo entusiasti di condividere ulteriori miglioramenti e nuove funzionalità man mano che espandiamo il nostro modello di base Cube 3D e di vedere i mondi immersivi che la nostra comunità di creatori costruisce con essi.