Versnelling van AI-inferentie voor 3D-creatie op Roblox

• Roblox heeft CUDA Graphs en KV-caching geïmplementeerd om het genereren van 3D-meshes te versnellen voor een snellere iteratie.
• Bij de lancering kon het Cube 3D-model tokens genereren in 7,8 milliseconden (een daling van 60,5 milliseconden) en volledige objecten in 4 seconden (een daling van 31 seconden). 

Eerder dit jaar deelde Roblox de eerste mogelijkheden van ons Cube 3D-basismodel. Met Cube 3D kunnen makers 3D-modellen en -omgevingen rechtstreeks genereren op basis van tekstprompts. Vanaf het begin hebben we prioriteit gegeven aan het optimaliseren van de latentie, omdat we ons realiseerden dat trage generatietijden een inherent iteratief proces verstoren. Voordat Cube 3D in maart werd gelanceerd, hadden we de inferentiestap al 7,8 keer sneller en responsiever gemaakt voor zowel ontwikkelaars als gebruikers. 

Sinds de lancering zijn er meer dan 578.000 objecten gegenereerd in verschillende opvallende ervaringen. Ontwikkelaars toonden ook interesse in het mogelijk maken dat gebruikers 3D-objecten kunnen genereren binnen ervaringen met tekstprompts zoals "katten", "burgers", enz. Met name Mic Up, een populair hangout-spel dat gebruikmaakt van voicechat, maakte gebruik van Cube 3D om spelers een leuke en interactieve manier te bieden om objecten te genereren. In hun implementatie kunnen spelers een menu aan de linkerkant openen met extra mogelijkheden, waaronder een AI-pictogram. Na het klikken op dat pictogram kunnen spelers een tekstprompt invoeren om een 3D-object te genereren. Voor gebruikers zorgen langere generatietijden voor wrijving, waardoor ze de magie missen om hun ideeën in realtime in 3D te zien veranderen. 

We wilden de 3D-generatie-ervaring transformeren van een interactie waarbij je moet stoppen en wachten naar iets dat responsief en natuurlijk aanvoelt, waardoor snel geëxperimenteerd kan worden. De mogelijkheid om snel objecten aan een scène toe te voegen is cruciaal voor ontwikkelaars. Om Cube 3D te versnellen, hebben we eerst de inferentie-pijplijn geprofileerd om prestatieknelpunten te identificeren. Ondanks het gebruik van krachtige GPU's, ontdekten we aanzienlijke wachttijd tussen bewerkingen.

Moderne deep learning-frameworks vertrouwen op de CPU om bewerkingen (of kernels) op de GPU in te plannen en te starten. De CPU bereidt elke bewerking voor, stuurt deze naar de GPU en wacht op bevestiging voordat de volgende bewerking wordt voorbereid. Dit wachten zorgt voor een bottleneck in de planning, waarbij de GPU stil kan staan terwijl de CPU de volgende batch werk voorbereidt. Idealiter willen we dat de CPU voorloopt op de GPU en bewerkingen voorbereidt en in de wachtrij plaatst, zodat de GPU altijd werk te doen heeft.

Dit is vooral problematisch voor autoregressieve decoders in transformator-achtige modellen zoals Cube 3D, die invoer sequentieel moeten verwerken en tokens moeten genereren. Deze modellen vereisen duizenden afzonderlijke bewerkingen voor één enkele generatie, en de rekenoverhead stapelt zich op bij elke stap in de reeks.

"We wilden iets bouwen dat vierdimensionale interactie mogelijk maakt", aldus Anupam Singh, Vice President of Engineering, die uitlegt waarom Roblox voor een autoregressieve aanpak heeft gekozen. "We willen niet alleen de auto bouwen; we willen ook de deur van de auto kunnen openen en erin kunnen stappen."

Elke bewerking kostte:

• CPU-tijd om elke kernel voor te bereiden
• Overhead door het starten van de kernel
• GPU-uitvoeringstijd (de daadwerkelijke berekening)
• Synchronisatie-overhead bij het controleren op voltooiing

Bij kleine bewerkingen die snel op de GPU worden uitgevoerd, kan deze overhead de inferentietijd domineren. De GPU zou slechts een klein deel van de totale inferentietijd actief kunnen rekenen.

Om deze bottleneck aan te pakken, hebben we gebruikgemaakt van CUDA Graphs – een functie waarmee reeksen GPU-bewerkingen kunnen worden opgenomen en afgespeeld zonder tussenkomst van de CPU. De autoregressieve decodercomponent van de architectuur van Cube 3D verwerkt tekstprompts en genereert vormtokens via een vector met vaste lengte. 

Hoewel onze dual-stream decoderarchitectuur functioneel vergelijkbaar is met een traditioneel groot taalmodel (LLM), is er een belangrijk verschil: deze maakt gebruik van twee parallelle aandachtsstromen. De ene stroom is gewijd aan conditietokens en de andere aan vormtokens. Standaard LLM-inferentie-engines waren niet geschikt voor onze behoeften en we hadden een aangepaste implementatie nodig die was afgestemd op onze specifieke architectuur.

Zie CUDA Graphs als het opnemen van een macro voor de GPU. In plaats van dat de CPU elke opdracht afzonderlijk geeft, neemt deze een volledige reeks GPU-bewerkingen (de grafiek) op en start de gehele grafiek met één enkele CPU-instructie. Deze aanpak vermindert de overhead bij het starten van de kernel drastisch, doordat de CPU niet langer elke bewerking afzonderlijk hoeft in te plannen tijdens de inferentie. Zodra de grafiek is gestart, voert de GPU de gehele reeks autonoom uit, zonder te wachten op verdere instructies.

CUDA Graphs hebben wel enkele beperkingen. Omdat de grafiekstructuur van tevoren moet worden bepaald, vereisen ze een vaste batchgrootte en invoerdimensies. Dit betekent dat er voor elke batchgrootte of invoervorm afzonderlijke grafieken moeten worden gemaakt. Voor onze use case met Cube 3D was deze beperking acceptabel, omdat we het inferentieproces konden standaardiseren rond gangbare invoerdimensies.

We moesten onze aanpak aanpassen om CUDA Graphs te implementeren voor ons Cube 3D-model. In traditionele LLM's worden attention-bewerkingen altijd uitgevoerd met dezelfde sequentielengte, wat een statische vorm oplevert om mee te werken. In onze aangepaste dual-stream-architectuur werken sommige attention-lagen echter alleen op de sequentielengte, terwijl andere werken op een combinatie van sequentie- en conditielengte.

Ondanks deze uitdagingen zagen we opmerkelijke resultaten na de implementatie van CUDA Graphs. We gebruiken tijd per outputtoken (TPOT) om de generatietijd voor elk token tijdens de inferentie te meten. Na de implementatie van CUDA Graphs verbeterde onze TPOT van 60,5 milliseconden naar 20,5 milliseconden, een verbetering van 2,9x. De totale generatietijd daalde met 66%, van 31 seconden naar 10,5 seconden.

Om de latentie verder te verbeteren, hebben we KV-caching geïmplementeerd, een standaardpraktijk bij LLM-inferentie die in de hele sector zeer effectief is gebleken. 

In op transformatoren gebaseerde modellen zoals Cube 3D vereist elke token-generatie het berekenen van sleutel- (K) en waardematrices (V) op basis van alle eerder gegenereerde tokens. Naarmate de reeks langer wordt, wordt het opnieuw berekenen van deze matrices voor elk token steeds inefficiënter.

KV-caching lost dit op door:

1. De K- en V-matrices voor alle eerder gegenereerde tokens op te slaan
2. De K- en V-matrices alleen voor nieuwe tokens te berekenen
3. deze nieuwe matrices toe te voegen aan de in de cache opgeslagen waarden

Deze aanpak elimineert overbodige berekeningen, waardoor de benodigde rekenkracht voor elk nieuw token wordt verminderd. Dit heeft vooral een groot effect naarmate de gegenereerde reeks langer wordt.

Onze aanpak voor het integreren van KV-caching met de CUDA Graph-implementatie was vergelijkbaar met traditionele LLM-inferentie. Door de toevoeging van KV-caching daalde onze TPOT tot slechts 7,8 milliseconden. De totale generatietijd daalde met 87%, van de oorspronkelijke 31 seconden tot slechts 4 seconden. Deze aanzienlijke tijdwinst maakt het veel effectiever voor makers die deze tool gebruiken.

We onderzoeken technieken die de latentie verder verminderen en de gebruikerservaring verbeteren, waaronder geoptimaliseerde kernels, modelkwantisering voor nog snellere inferentie, hardwarespecifieke optimalisaties en parallelle token-generatie.

Dit werk wordt complexer wanneer we uitbreiden naar het genereren en begrijpen van volledige scènes, waarbij veel 3D-elementen binnen een lay-out in context met elkaar moeten samenwerken. We willen ook dat de 3D-objecten en werelden die we creëren volledig functioneel zijn, zodat deuren openen en sluiten, wielen draaien, enzovoort. Om dat te bereiken, hebben we snelle generatie en iteratie nodig om op te schalen naar volledige scènes, volledig functionele objecten en avatars. We kijken ernaar uit om verdere verbeteringen en nieuwe functionaliteit te delen naarmate we ons Cube 3D-basismodel uitbreiden – en om de meeslepende werelden te zien die onze community van makers hiermee bouwt.