Beschleunigung der KI-Inferenz für die 3D-Erstellung auf Roblox

• Roblox hat CUDA-Graphs und KV-Caching implementiert, um die Erzeugung von 3D-Meshes zu beschleunigen und so eine reaktionsschnellere Iteration zu ermöglichen.
• Zum Start konnte das Cube-3D-Modell Token in 7,8 Millisekunden (statt 60,5 Millisekunden) und vollständige Objekte in 4 Sekunden (statt 31 Sekunden) generieren. 

Anfang dieses Jahres stellte Roblox die ersten Funktionen unseres Cube-3D-Grundlagemodells vor. Mit Cube 3D können Entwickler 3D-Modelle und Umgebungen direkt aus Textbefehlen generieren. Von Anfang an haben wir der Latenzoptimierung Priorität eingeräumt, da wir erkannt haben, dass lange Generierungszeiten einen von Natur aus iterativen Prozess stören. Bereits vor dem Start von Cube 3D im März hatten wir den Inferenzschritt um das 7,8-Fache beschleunigt und sowohl für Entwickler als auch für Nutzer reaktionsschneller gemacht. 

Seit der Einführung wurden mehr als 578.000 Objekte in verschiedenen namhaften Erlebnissen generiert. Entwickler zeigten zudem Interesse daran, Nutzern die Generierung von 3D-Objekten innerhalb von Erlebnissen anhand von Textbefehlen wie „Katzen“, „Burger“ usw. zu ermöglichen. Vor allem „Mic Up“, ein beliebtes Hangout-Spiel mit Voice-Chat, nutzte Cube 3D, um Spielern eine unterhaltsame und interaktive Möglichkeit zur Objektgenerierung zu bieten. In ihrer Implementierung können Spieler ein Menü auf der linken Seite mit zusätzlichen Funktionen öffnen, darunter ein KI-Symbol. Nach dem Klicken auf dieses Symbol können Spieler eine Textanweisung eingeben, um ein 3D-Objekt zu generieren. Für Nutzer verursachen längere Generierungszeiten Reibungsverluste und nehmen ihnen die Magie, ihre Ideen in Echtzeit in 3D umgesetzt zu sehen. 

Wir wollten das Erlebnis der 3D-Generierung von einer Interaktion, bei der man warten muss, in etwas verwandeln, das reaktionsschnell und natürlich wirkt und schnelles Experimentieren ermöglicht. Die Möglichkeit, Objekte schnell zu einer Szene hinzuzufügen, ist für Entwickler entscheidend. Um Cube 3D zu beschleunigen, haben wir zunächst die Inferenz-Pipeline profiliert, um Leistungsengpässe zu identifizieren. Trotz der Verwendung leistungsstarker GPUs stellten wir erhebliche Leerlaufzeiten zwischen den Operationen fest.

Moderne Deep-Learning-Frameworks sind auf die CPU angewiesen, um Operationen (oder Kernel) auf der GPU zu planen und zu starten. Die CPU bereitet jede Operation vor, sendet sie an die GPU und wartet auf eine Bestätigung, bevor sie die nächste Operation vorbereitet. Dieses Warten führt zu einem Planungsengpass, bei dem die GPU untätig bleiben könnte, während die CPU den nächsten Arbeitsstapel vorbereitet. Im Idealfall soll die CPU der GPU voraus sein und Operationen vorbereiten und in die Warteschlange stellen, sodass die GPU immer Arbeit hat.

Dies ist besonders problematisch für autoregressive Decoder in Transformer-Modellen wie Cube 3D, die Eingaben sequenziell verarbeiten und Token generieren müssen. Diese Modelle erfordern Tausende einzelner Operationen für eine einzige Generierung, und der Rechenaufwand summiert sich mit jedem Schritt in der Sequenz.

„Wir wollten etwas entwickeln, das eine vierdimensionale Interaktion ermöglicht“, sagte Anupam Singh, Vice President of Engineering, und erklärte, warum Roblox sich für einen autoregressiven Ansatz entschieden hat. „Wir wollen nicht nur das Auto bauen; wir wollen auch in der Lage sein, die Autotür zu öffnen und einzusteigen.“

Jede Operation verursachte:

• CPU-Zeit für die Vorbereitung jedes Kernels
• Overhead durch das Starten des Kernels
• GPU-Ausführungszeit (die eigentliche Berechnung)
• Synchronisations-Overhead bei der Überprüfung auf Fertigstellung

Bei kleinen Operationen, die auf der GPU schnell ausgeführt werden, kann dieser Overhead die Inferenzzeit dominieren. Die GPU könnte nur einen kleinen Bruchteil der gesamten Inferenzzeit aktiv mit Berechnungen verbringen.

Um diesen Engpass zu beheben, nutzten wir CUDA Graphs – eine Funktion, die das Aufzeichnen und Wiedergeben von Sequenzen von GPU-Operationen ohne Eingriff der CPU ermöglicht. Die autoregressive Decoder-Komponente der Architektur von Cube 3D verarbeitet Text-Prompts und generiert Form-Token über einen Vektor fester Länge. 

Obwohl sie funktional einem traditionellen großen Sprachmodell (LLM) ähnelt, weist unsere Dual-Stream-Decoder-Architektur einen wichtigen Unterschied auf: Sie nutzt zwei parallele Attention-Streams. Ein Stream ist für Bedingungstoken vorgesehen, der andere für Form-Token. Standardmäßige LLM-Inferenz-Engines waren für unsere Anforderungen nicht geeignet, und wir benötigten eine maßgeschneiderte Implementierung, die auf unsere spezifische Architektur zugeschnitten war.

Stellen Sie sich CUDA-Graphen wie die Aufzeichnung eines Makros für die GPU vor. Anstatt dass die CPU jeden Befehl einzeln ausgibt, zeichnet sie eine gesamte Abfolge von GPU-Operationen (den Graphen) auf und startet den gesamten Graphen mit einem einzigen CPU-Befehl. Dieser Ansatz reduziert den Overhead beim Kernel-Start drastisch, da die CPU während der Inferenz nicht mehr jede Operation einzeln planen muss. Sobald der Graph gestartet ist, führt die GPU die gesamte Abfolge autonom aus, ohne auf weitere Anweisungen zu warten.

CUDA-Graphen weisen jedoch einige Einschränkungen auf. Da die Graphstruktur im Voraus festgelegt werden muss, erfordern sie eine feste Batchgröße und feste Eingabedimensionen. Das bedeutet, dass für jede Batchgröße oder Eingabeform separate Graphen erstellt werden müssen. Für unseren Anwendungsfall mit Cube 3D war diese Einschränkung akzeptabel, da wir den Inferenzprozess auf gemeinsame Eingabedimensionen standardisieren konnten.

Wir mussten unseren Ansatz anpassen, um CUDA-Graphen für unser Cube-3D-Modell zu implementieren. In herkömmlichen LLMs werden Attention-Operationen immer mit derselben Sequenzlänge durchgeführt, was eine statische Form für die Verarbeitung bietet. In unserer benutzerdefinierten Dual-Stream-Architektur arbeiten einige Attention-Schichten jedoch allein mit der Sequenzlänge, während andere mit einer Kombination aus Sequenz- und Bedingungslänge arbeiten.

Trotz dieser Herausforderungen konnten wir nach der Implementierung von CUDA Graphs bemerkenswerte Ergebnisse erzielen. Wir verwenden die Zeit pro Ausgabetoken (TPOT), um die Generierungszeit für jedes Token während der Inferenz zu messen. Nach der Implementierung von CUDA Graphs verbesserte sich unser TPOT von 60,5 Millisekunden auf 20,5 Millisekunden, was einer 2,9-fachen Verbesserung entspricht. Die Gesamtgenerierungszeit sank um 66 % von 31 Sekunden auf 10,5 Sekunden.

Um die Latenz weiter zu verbessern, haben wir KV-Caching implementiert, eine Standardpraxis bei der LLM-Inferenz, die sich branchenweit als äußerst effektiv erwiesen hat. 

In Transformer-basierten Modellen wie Cube 3D erfordert jede Token-Generierung die Berechnung von Schlüssel- (K) und Wertmatrizen (V) auf der Grundlage aller zuvor generierten Token. Je länger die Sequenz wird, desto ineffizienter wird die Neuberechnung dieser Matrizen für jedes einzelne Token.

KV-Caching löst dieses Problem durch:

1. Speicherung der K- und V-Matrizen für alle zuvor generierten Token
2. Berechnung der K- und V-Matrizen nur für neue Token
3. Anhängen dieser neuen Matrizen an die zwischengespeicherten Werte

Dieser Ansatz eliminiert redundante Berechnungen und reduziert den Aufwand für jedes neue Token. Dies wirkt sich besonders deutlich aus, wenn die generierte Sequenz länger wird.

Unser Ansatz zur Integration von KV-Caching in die CUDA-Graph-Implementierung ähnelte der traditionellen LLM-Inferenz. Durch das Hinzufügen von KV-Caching konnten wir unseren TPOT auf nur 7,8 Millisekunden senken. Die Gesamtgenerierungszeit verringerte sich um 87 % – von ursprünglich 31 Sekunden auf nur noch 4 Sekunden. Diese erhebliche Zeitersparnis macht das Tool für Kreative deutlich effektiver.

Wir erforschen Techniken, die die Latenz weiter reduzieren und das Benutzererlebnis verbessern, darunter optimierte Kernel, Modellquantisierung für noch schnellere Inferenz, hardwarespezifische Optimierungen und parallele Token-Generierung.

Diese Arbeit wird komplexer, wenn wir auf die vollständige Szenenerstellung und -interpretation ausweiten, bei der viele 3D-Elemente innerhalb eines Layouts kontextbezogen zusammenwirken müssen. Außerdem möchten wir, dass die von uns erstellten 3D-Objekte und -Welten voll funktionsfähig sind, sodass sich Türen öffnen und schließen, Räder sich drehen usw. Um dies zu erreichen, benötigen wir schnelle Generierung und Iteration, um auf ganze Szenen, voll funktionsfähige Objekte und Avatare zu skalieren. Wir freuen uns darauf, weitere Verbesserungen und neue Funktionen vorzustellen, während wir unser Cube 3D-Grundmodell erweitern – und darauf, die immersiven Welten zu sehen, die unsere Creator-Community damit erschafft.