Accélération de l'inférence IA pour la création 3D sur Roblox

• Roblox a mis en œuvre CUDA Graphs et la mise en cache KV pour accélérer la génération de maillages 3D et obtenir des itérations plus réactives.
• Au moment du lancement, le modèle Cube 3D pouvait générer des jetons en 7,8 millisecondes (contre 60,5 millisecondes auparavant) et des objets complets en 4 secondes (contre 31 secondes auparavant). 

Plus tôt cette année, Roblox a présenté les premières fonctionnalités de son modèle de base Cube 3D. Avec Cube 3D, les créateurs peuvent générer des modèles et des environnements 3D directement à partir de invites textuelles. Dès le départ, nous avons donné la priorité à l’optimisation de la latence, conscients que des temps de génération lents perturbent ce qui est intrinsèquement un processus itératif. Avant le lancement de Cube 3D en mars, nous avions déjà rendu l’étape d’inférence 7,8 fois plus rapide et plus réactive, tant pour les développeurs que pour les utilisateurs. 

Depuis son lancement, plus de 578 000 objets ont été générés dans le cadre de plusieurs expériences notables. Les développeurs ont également exprimé leur intérêt pour permettre aux utilisateurs de générer des objets 3D au sein des expériences à l’aide de commandes textuelles telles que « chats », « burgers », etc. Notamment, Mic Up, un jeu de rencontre populaire utilisant le chat vocal, a exploité Cube 3D pour offrir aux joueurs un moyen amusant et interactif de générer des objets. Dans leur implémentation, les joueurs peuvent ouvrir un menu de gauche proposant des fonctionnalités supplémentaires, dont une icône IA. Après avoir cliqué sur cette icône, les joueurs peuvent saisir une commande textuelle pour générer un objet 3D. Pour les utilisateurs, des temps de génération plus longs créent une friction, les privant de la magie de voir leurs idées se transformer en 3D en temps réel. 

Nous voulions transformer l'expérience de génération 3D, passant d'une interaction de type « s'arrêter et attendre » à quelque chose qui semble réactif et naturel, permettant une expérimentation rapide. La capacité à ajouter rapidement des objets à une scène est essentielle pour les développeurs. Pour accélérer Cube 3D, nous avons d'abord profilé le pipeline d'inférence afin d'identifier les goulots d'étranglement en termes de performances. Malgré l'utilisation de GPU puissants, nous avons constaté un temps d'inactivité significatif entre les opérations.

Les frameworks modernes d'apprentissage profond s'appuient sur le CPU pour planifier et lancer des opérations (ou noyaux) sur le GPU. Le CPU prépare chaque opération, l'envoie au GPU et attend la confirmation avant de préparer l'opération suivante. Cette attente crée un goulot d'étranglement dans la planification, le GPU pouvant rester inactif pendant que le CPU prépare le lot de travail suivant. Idéalement, nous souhaitons que le CPU devance le GPU, en préparant et en mettant en file d'attente les opérations afin que le GPU ait toujours du travail à effectuer.

Ce problème est particulièrement aigu pour les décodeurs autorégressifs dans les modèles de type transformateur comme Cube 3D, qui doivent traiter les entrées et générer des tokens de manière séquentielle. Ces modèles nécessitent des milliers d’opérations individuelles pour une seule génération, et la charge de calcul s’accumule à chaque étape de la séquence.

« Nous voulions créer quelque chose qui permette une interaction en quatre dimensions », a déclaré Anupam Singh, vice-président de l'ingénierie, expliquant pourquoi Roblox a choisi une approche autorégressive. « Nous ne voulons pas seulement construire la voiture ; nous voulons aussi pouvoir ouvrir la portière et monter à l'intérieur. »

Chaque opération a nécessité :

• Temps CPU pour préparer chaque noyau
• La charge liée au lancement du noyau
• Temps d'exécution du GPU (le calcul proprement dit)
• Une surcharge de synchronisation lors de la vérification de l'achèvement

Dans le cas d'opérations de petite taille s'exécutant rapidement sur le GPU, cette surcharge peut prédominer sur le temps d'inférence. Le GPU pourrait ne calculer activement que pendant une fraction minime du temps total d'inférence.

Pour remédier à ce goulot d'étranglement, nous avons exploité CUDA Graphs, une fonctionnalité qui permet d'enregistrer et de rejouer des séquences d'opérations GPU sans intervention du CPU. Le composant décodeur autorégressif de l'architecture de Cube 3D traite les invites textuelles et génère des jetons de forme via un vecteur de longueur fixe. 

Bien que fonctionnellement similaire à un grand modèle linguistique (LLM) traditionnel, notre architecture de décodeur à double flux présente une différence importante : elle utilise deux flux d’attention parallèles. Un flux est dédié aux jetons de condition et l’autre aux jetons de forme. Les moteurs d’inférence LLM prêts à l’emploi ne répondaient pas à nos besoins, et nous avions besoin d’une implémentation sur mesure adaptée à notre architecture spécifique.

Considérez les graphes CUDA comme l'enregistrement d'une macro pour le GPU. Au lieu que le CPU émette chaque commande individuellement, il enregistre une séquence complète d'opérations GPU (le graphe) et lance l'ensemble du graphe à l'aide d'une seule instruction CPU. Cette approche réduit considérablement la surcharge liée au lancement des noyaux en éliminant la nécessité pour le CPU de planifier individuellement chaque opération pendant l'inférence. Une fois le graphe lancé, le GPU exécute l'ensemble de la séquence de manière autonome, sans attendre d'autres instructions.

Les graphes CUDA présentent toutefois certaines limites. La structure du graphe devant être déterminée à l’avance, ils nécessitent une taille de lot et des dimensions d’entrée fixes. Cela signifie qu’il faut créer des graphes distincts pour chaque taille de lot ou forme d’entrée. Pour notre cas d’utilisation avec Cube 3D, cette limitation était acceptable, car nous pouvions standardiser le processus d’inférence autour de dimensions d’entrée communes.

Nous avons dû adapter notre approche pour implémenter les graphes CUDA dans notre modèle Cube 3D. Dans les LLM traditionnels, les opérations d’attention sont toujours effectuées avec la même longueur de séquence, ce qui fournit une forme statique avec laquelle travailler. Cependant, dans notre architecture double flux personnalisée, certaines couches d’attention opèrent uniquement sur la longueur de la séquence, tandis que d’autres opèrent sur une combinaison de la longueur de la séquence et de la condition.

Malgré ces défis, nous avons obtenu des résultats remarquables après la mise en œuvre de CUDA Graphs. Nous utilisons le temps par token de sortie (TPOT) pour mesurer le temps de génération de chaque token pendant l'inférence. Après la mise en œuvre de CUDA Graphs, notre TPOT est passé de 60,5 millisecondes à 20,5 millisecondes, soit une amélioration de 2,9 fois. Le temps de génération global a baissé de 66 %, passant de 31 secondes à 10,5 secondes.

Afin d'améliorer encore davantage la latence, nous avons mis en place la mise en cache KV, une pratique courante dans l'inférence LLM qui s'est avérée très efficace dans l'ensemble du secteur. 

Dans les modèles basés sur des transformateurs comme Cube 3D, chaque génération de token nécessite le calcul de matrices de clés (K) et de valeurs (V) basées sur tous les tokens générés précédemment. À mesure que la séquence s'allonge, recalculer ces matrices pour chaque token devient de plus en plus inefficace.

La mise en cache KV résout ce problème en :

1. Stockant les matrices K et V pour tous les tokens générés précédemment
2. Calculant les matrices K et V uniquement pour les nouveaux tokens
3. Ajoutant ces nouvelles matrices aux valeurs mises en cache

Cette approche élimine les calculs redondants, réduisant ainsi le travail requis pour chaque nouveau jeton. Cela devient particulièrement significatif à mesure que la séquence générée s'allonge.

Notre approche pour intégrer la mise en cache KV à l'implémentation CUDA Graph était similaire à celle de l'inférence LLM traditionnelle. L'ajout de la mise en cache KV a réduit notre TPOT à seulement 7,8 millisecondes. Le temps de génération global a diminué de 87 %, passant de 31 secondes à seulement 4 secondes. Cette réduction significative du temps rend cet outil beaucoup plus efficace pour les créateurs qui l'utilisent.

Nous explorons des techniques permettant de réduire davantage la latence et d'améliorer l'expérience utilisateur, notamment des noyaux optimisés, la quantification des modèles pour une inférence encore plus rapide, des optimisations spécifiques au matériel et la génération parallèle de tokens.

Ce travail se complexifie lorsque nous passons à la génération et à la compréhension de scènes complètes, où de nombreux éléments 3D doivent interagir les uns avec les autres dans un contexte donné au sein d’une mise en page. Nous souhaitons également que les objets et les mondes 3D que nous créons soient pleinement fonctionnels, afin que les portes s’ouvrent et se ferment, que les roues tournent, etc. Pour y parvenir, nous avons besoin d’une génération et d’une itération rapides afin de pouvoir traiter des scènes entières, des objets pleinement fonctionnels et des avatars. Nous sommes ravis de partager de nouvelles améliorations et fonctionnalités à mesure que nous développons notre modèle de base Cube 3D, et de découvrir les mondes immersifs que notre communauté de créateurs construira grâce à lui.