Acelerando a inferência de IA para criação 3D no Roblox

• A Roblox implementou CUDA Graphs e cache KV para acelerar a geração de malhas 3D, proporcionando iterações mais responsivas.
• No lançamento, o modelo Cube 3D conseguia gerar tokens em 7,8 milissegundos (uma redução de 60,5 milissegundos) e objetos completos em 4 segundos (uma redução de 31 segundos). 

No início deste ano, a Roblox compartilhou a primeira funcionalidade do nosso modelo base Cube 3D. Com o Cube 3D, os criadores podem gerar modelos e ambientes 3D diretamente a partir de prompts de texto. Desde o início, priorizamos a otimização da latência, reconhecendo que tempos de geração lentos atrapalham o que é, por natureza, um processo iterativo. Antes do lançamento do Cube 3D em março, já tínhamos tornado a etapa de inferência 7,8 vezes mais rápida e mais responsiva tanto para desenvolvedores quanto para usuários. 

Desde o lançamento, mais de 578.000 objetos foram gerados em várias experiências notáveis. Os desenvolvedores também demonstraram interesse em permitir que os usuários gerassem objetos 3D dentro das experiências por meio de comandos de texto como “gatos”, “hambúrgueres” etc. Mais notavelmente, o Mic Up, um popular jogo social que usa chat de voz, aproveitou o Cube 3D para oferecer aos jogadores uma maneira divertida e interativa de gerar objetos. Na implementação deles, os jogadores podem abrir um menu à esquerda com recursos adicionais, incluindo um ícone de IA. Após clicar nesse ícone, os jogadores podem inserir um prompt de texto para gerar um objeto 3D. Para os usuários, tempos de geração mais longos criam atrito, privando-os da magia de ver suas ideias transformadas em 3D em tempo real. 

Queríamos transformar a experiência de geração 3D de uma interação do tipo “parar e esperar” em algo que parecesse responsivo e natural, permitindo experimentação rápida. A capacidade de adicionar objetos rapidamente a uma cena é fundamental para os desenvolvedores. Para acelerar o Cube 3D, primeiro analisamos o pipeline de inferência para identificar gargalos de desempenho. Apesar de usarmos GPUs potentes, encontramos um tempo de inatividade significativo entre as operações.

As estruturas modernas de aprendizado profundo dependem da CPU para agendar e iniciar operações (ou kernels) na GPU. A CPU prepara cada operação, envia-a para a GPU e aguarda a confirmação antes de preparar a próxima operação. Essa espera cria um gargalo de agendamento, em que a GPU pode ficar ociosa enquanto a CPU prepara o próximo lote de trabalho. Idealmente, queremos que a CPU trabalhe à frente da GPU, preparando e enfileirando operações para que a GPU sempre tenha trabalho a fazer.

Isso é especialmente problemático para decodificadores autorregressivos em modelos do tipo transformador, como o Cube 3D, que precisam processar entradas e gerar tokens sequencialmente. Esses modelos exigem milhares de operações individuais para uma única geração, e a sobrecarga computacional se acumula a cada etapa da sequência.

“Queríamos construir algo que permitisse uma interação quadridimensional”, disse o vice-presidente de engenharia Anupam Singh, explicando por que a Roblox escolheu uma abordagem autorregressiva. “Não queremos apenas construir o carro; também queremos poder abrir a porta do carro e entrar nele.”

Cada operação incorria em:

• Tempo de CPU para preparar cada kernel
• Sobrecarga do lançamento do kernel
• Tempo de execução da GPU (o cálculo propriamente dito)
• Sobrecarga de sincronização ao verificar a conclusão

No caso de pequenas operações que são executadas rapidamente na GPU, essa sobrecarga pode dominar o tempo de inferência. A GPU pode estar computando ativamente por apenas uma pequena fração do tempo total de inferência.

Para resolver esse gargalo, utilizamos o CUDA Graphs — um recurso que permite gravar e reproduzir sequências de operações da GPU sem a intervenção da CPU. O componente decodificador autorregressivo da arquitetura do Cube 3D processa prompts de texto e gera tokens de forma por meio de um vetor de comprimento fixo. 

Embora funcionalmente semelhante a um modelo de linguagem grande (LLM) tradicional, nossa arquitetura de decodificador de fluxo duplo apresenta uma diferença importante: ela utiliza dois fluxos de atenção paralelos. Um fluxo é dedicado aos tokens de condição e o outro aos tokens de forma. Os mecanismos de inferência LLM prontos para uso não eram adequados às nossas necessidades, e precisávamos de uma implementação personalizada adaptada à nossa arquitetura específica.

Pense nos CUDA Graphs como a gravação de uma macro para a GPU. Em vez de a CPU emitir cada comando individualmente, ela grava uma sequência inteira de operações da GPU (o gráfico) e inicia o gráfico inteiro com uma única instrução da CPU. Essa abordagem reduz drasticamente a sobrecarga de inicialização do kernel, eliminando a necessidade de a CPU agendar individualmente cada operação durante a inferência. Uma vez que o gráfico é iniciado, a GPU executa a sequência inteira de forma autônoma, sem esperar por novas instruções.

Os CUDA Graphs vêm com algumas limitações. Como a estrutura do gráfico precisa ser determinada com antecedência, eles exigem um tamanho de lote e dimensões de entrada fixos. Isso significa que gráficos separados precisam ser criados para cada tamanho de lote ou formato de entrada. Para o nosso caso de uso com o Cube 3D, essa limitação era aceitável, pois pudemos padronizar o processo de inferência em torno de dimensões de entrada comuns.

Tivemos que adaptar nossa abordagem para implementar os gráficos CUDA em nosso modelo Cube 3D. Em LLMs tradicionais, as operações de atenção são sempre realizadas com o mesmo comprimento de sequência, fornecendo uma forma estática para trabalhar. No entanto, em nossa arquitetura personalizada de fluxo duplo, algumas camadas de atenção operam apenas com o comprimento da sequência, enquanto outras operam com uma combinação do comprimento da sequência e da condição.

Apesar desses desafios, observamos resultados notáveis após a implementação do CUDA Graphs. Usamos o tempo por token de saída (TPOT) para medir o tempo de geração de cada token durante a inferência. Após a implementação do CUDA Graphs, nosso TPOT melhorou de 60,5 milissegundos para 20,5 milissegundos, uma melhoria de 2,9 vezes. O tempo total de geração caiu 66%, passando de 31 segundos para 10,5 segundos.

Para melhorar ainda mais a latência, implementamos o cache KV, uma prática padrão na inferência de LLM que se mostrou altamente eficaz em todo o setor. 

Em modelos baseados em transformadores, como o Cube 3D, cada geração de token requer o cálculo de matrizes de chave (K) e valor (V) com base em todos os tokens gerados anteriormente. À medida que a sequência se torna mais longa, recalcular essas matrizes para cada token se torna cada vez mais ineficiente.

O cache KV resolve isso ao:

1. Armazenar as matrizes K e V para todos os tokens gerados anteriormente
2. Calcular as matrizes K e V apenas para novos tokens
3. Anexando essas novas matrizes aos valores armazenados em cache

Essa abordagem elimina cálculos redundantes, reduzindo o trabalho necessário para cada novo token. Isso se torna especialmente impactante à medida que a sequência gerada se torna mais longa.

Nossa abordagem para integrar o cache KV à implementação do CUDA Graph foi semelhante à inferência LLM tradicional. A adição do cache KV reduziu nosso TPOT para apenas 7,8 milissegundos. O tempo total de geração diminuiu 87%, passando dos 31 segundos originais para apenas 4 segundos. Essa redução significativa de tempo torna a ferramenta muito mais eficaz para os criadores que a utilizam.

Estamos explorando técnicas que reduzem ainda mais a latência e melhoram a experiência do usuário, incluindo kernels otimizados, quantização de modelos para inferência ainda mais rápida, otimizações específicas de hardware e geração paralela de tokens.

Esse trabalho se torna mais complexo quando expandimos para a geração e compreensão de cenas completas, onde muitos elementos 3D precisam interagir entre si dentro de um layout. Também queremos que os objetos e mundos 3D que criamos sejam totalmente funcionais, de modo que portas se abram e fechem, rodas girem, etc. Para chegar lá, precisamos de geração e iteração rápidas para escalar para cenas inteiras, objetos totalmente funcionais e avatares. Estamos animados para compartilhar mais melhorias e novas funcionalidades à medida que expandimos nosso modelo de base Cube 3D — e para ver os mundos imersivos que nossa comunidade de criadores constrói com eles.