Aceleración de la inferencia de IA para la creación 3D en Roblox
Generación de objetos 3D 7,8 veces más rápida y con mayor capacidad de respuesta
- Roblox implementó CUDA Graphs y el almacenamiento en caché KV para acelerar la generación de mallas 3D y lograr una iteración más ágil.
- En el momento del lanzamiento, el modelo 3D Cube podía generar tokens en 7,8 milisegundos (frente a los 60,5 milisegundos anteriores) y objetos completos en 4 segundos (frente a los 31 segundos anteriores).
A principios de este año, Roblox compartió la primera capacidad de nuestro modelo base Cube 3D. Con Cube 3D, los creadores pueden generar modelos y entornos 3D directamente a partir de indicaciones de texto. Desde el principio, dimos prioridad a la optimización de la latencia, reconociendo que los tiempos de generación lentos interrumpen lo que es, en esencia, un proceso iterativo. Antes del lanzamiento de Cube 3D en marzo, ya habíamos logrado que el paso de inferencia fuera 7,8 veces más rápido y más receptivo tanto para los desarrolladores como para los usuarios.
Desde su lanzamiento, se han generado más de 578 000 objetos en varias experiencias destacadas. Los desarrolladores también mostraron interés en permitir a los usuarios generar objetos 3D dentro de las experiencias mediante indicaciones de texto como «gatos», «hamburguesas», etc. En particular, Mic Up, un popular juego social que utiliza el chat de voz, aprovechó Cube 3D para ofrecer a los jugadores una forma divertida e interactiva de generar objetos. En su implementación, los jugadores pueden abrir un menú a la izquierda con funciones adicionales, incluido un icono de IA. Tras hacer clic en ese icono, los jugadores pueden introducir un comando de texto para generar un objeto 3D. Para los usuarios, los tiempos de generación más largos crean fricción, privándoles de la magia de ver sus ideas transformadas en 3D en tiempo real.
Queríamos transformar la experiencia de generación 3D de una interacción de «parar y esperar» a algo que se sintiera receptivo y natural, permitiendo una experimentación rápida. La capacidad de añadir rápidamente objetos a una escena es fundamental para los desarrolladores. Para acelerar Cube 3D, primero analizamos el proceso de inferencia para identificar los cuellos de botella en el rendimiento. A pesar de utilizar potentes GPU, encontramos un tiempo de inactividad significativo entre operaciones.
Resolución del cuello de botella en la programación de CPU y GPU
Los marcos modernos de aprendizaje profundo dependen de la CPU para programar y ejecutar operaciones (o kernels) en la GPU. La CPU prepara cada operación, la envía a la GPU y espera la confirmación antes de preparar la siguiente operación. Esta espera crea un cuello de botella en la programación, en el que la GPU puede permanecer inactiva mientras la CPU prepara el siguiente lote de trabajo. Lo ideal sería que la CPU se adelantara a la GPU, preparando y poniendo en cola las operaciones para que la GPU siempre tenga trabajo que hacer.
Esto resulta especialmente problemático para los decodificadores autorregresivos en modelos de tipo transformador como Cube 3D, que necesitan procesar la entrada y generar tokens de forma secuencial. Estos modelos requieren miles de operaciones individuales para una sola generación, y la sobrecarga computacional se acumula con cada paso de la secuencia.
«Queríamos crear algo que permitiera una interacción en cuatro dimensiones», afirmó el vicepresidente de Ingeniería, Anupam Singh, al explicar por qué Roblox eligió un enfoque autorregresivo. «No solo queremos construir el coche; también queremos poder abrir la puerta del coche y subirnos a él».
Cada operación requirió:
- Tiempo de CPU para preparar cada kernel
- Sobrecarga derivada del lanzamiento del núcleo
- Tiempo de ejecución de la GPU (el cálculo real)
- Sobrecarga de sincronización al comprobar la finalización
En el caso de operaciones pequeñas que se ejecutan rápidamente en la GPU, esta sobrecarga puede dominar el tiempo de inferencia. La GPU podría estar calculando activamente solo durante una pequeña fracción del tiempo total de inferencia.
Implementación de CUDA Graphs: eliminando al intermediario
Para abordar este cuello de botella, aprovechamos CUDA Graphs, una función que permite grabar y reproducir secuencias de operaciones de la GPU sin intervención de la CPU. El componente decodificador autorregresivo de la arquitectura de Cube 3D procesa las indicaciones de texto y genera tokens de forma a través de un vector de longitud fija.
Aunque funcionalmente similar a un modelo de lenguaje grande (LLM) tradicional, nuestra arquitectura de decodificador de doble flujo presenta una diferencia importante: utiliza dos flujos de atención paralelos. Un flujo se dedica a los tokens de condición y el otro a los tokens de forma. Los motores de inferencia LLM estándar no se adaptaban a nuestras necesidades, por lo que necesitábamos una implementación personalizada a medida de nuestra arquitectura específica.
Piensa en los gráficos CUDA como si se tratara de grabar una macro para la GPU. En lugar de que la CPU emita cada comando individualmente, graba una secuencia completa de operaciones de la GPU (el gráfico) y ejecuta todo el gráfico con una sola instrucción de la CPU. Este enfoque reduce drásticamente la sobrecarga de la ejecución del kernel al eliminar la necesidad de que la CPU programe individualmente cada operación durante la inferencia. Una vez ejecutado el gráfico, la GPU ejecuta toda la secuencia de forma autónoma, sin esperar más instrucciones.
Los gráficos CUDA tienen algunas limitaciones. Dado que la estructura del gráfico debe determinarse de antemano, requieren un tamaño de lote y unas dimensiones de entrada fijas. Esto significa que hay que crear gráficos separados para cada tamaño de lote o forma de entrada. Para nuestro caso de uso con Cube 3D, esta limitación era aceptable, ya que podíamos estandarizar el proceso de inferencia en torno a dimensiones de entrada comunes.
Tuvimos que adaptar nuestro enfoque para implementar los gráficos CUDA en nuestro modelo Cube 3D. En los LLM tradicionales, las operaciones de atención siempre se realizan con la misma longitud de secuencia, lo que proporciona una forma estática con la que trabajar. Sin embargo, en nuestra arquitectura personalizada de doble flujo, algunas capas de atención operan solo sobre la longitud de la secuencia, mientras que otras operan sobre una combinación de la longitud de la secuencia y la condición.
A pesar de estos retos, observamos resultados notables tras implementar CUDA Graphs. Utilizamos el tiempo por token de salida (TPOT) para medir el tiempo de generación de cada token durante la inferencia. Tras implementar CUDA Graphs, nuestro TPOT mejoró de 60,5 milisegundos a 20,5 milisegundos, lo que supone una mejora de 2,9 veces. El tiempo total de generación se redujo en un 66 %, pasando de 31 segundos a 10,5 segundos.
KV Caching: Aprovechando nuestro éxito
Para mejorar aún más la latencia, implementamos el almacenamiento en caché KV, una práctica estándar en la inferencia de LLM que ha demostrado ser muy eficaz en todo el sector.
En modelos basados en transformadores como Cube 3D, cada generación de tokens requiere el cálculo de matrices de clave (K) y valor (V) basadas en todos los tokens generados anteriormente. A medida que la secuencia se alarga, volver a calcular estas matrices para cada token se vuelve cada vez más ineficiente.
El almacenamiento en caché KV resuelve esto mediante:
- Almacenar las matrices K y V de todos los tokens generados anteriormente
- Calculando las matrices K y V solo para los nuevos tokens
- Añadiendo estas nuevas matrices a los valores almacenados en caché
Este enfoque elimina el cálculo redundante, reduciendo el trabajo necesario para cada nuevo token. Esto tiene un impacto especialmente significativo a medida que la secuencia generada se alarga.
Nuestro enfoque para integrar el almacenamiento en caché KV con la implementación de CUDA Graph fue similar al de la inferencia tradicional de LLM. La incorporación del almacenamiento en caché KV redujo nuestro TPOT a tan solo 7,8 milisegundos. El tiempo total de generación disminuyó un 87 %, pasando de los 31 segundos originales a tan solo 4 segundos. Esta significativa reducción de tiempo hace que la herramienta resulte mucho más eficaz para los creadores que la utilizan.
Evaluación del impacto real en desarrolladores y usuarios
Estamos explorando técnicas que reduzcan aún más la latencia y mejoren la experiencia del usuario, incluyendo núcleos optimizados, cuantificación de modelos para una inferencia aún más rápida, optimizaciones específicas para el hardware y generación paralela de tokens.
Este trabajo se vuelve más complejo cuando nos expandimos a la generación y comprensión de escenas completas, donde muchos elementos 3D deben funcionar juntos en contexto unos con otros dentro de un diseño. También queremos que los objetos y mundos 3D que creamos sean totalmente funcionales, de modo que las puertas se abran y se cierren, las ruedas giren, etc. Para lograrlo, necesitamos una generación y una iteración rápidas que se adapten a escenas completas, objetos totalmente funcionales y avatares. Estamos deseando compartir nuevas mejoras y funcionalidades a medida que ampliamos nuestro modelo base Cube 3D, y ver los mundos inmersivos que nuestra comunidad de creadores construye con ellos.
