CUDA Agent treina modelos de linguagem para escrever kernels de GPU mais rápidos que o torch.compile

Introdução

A equipe Seed da ByteDance, em parceria com o Instituto de Pesquisa em Inteligência Artificial (AIR) da Universidade de Tsinghua, lançou o CUDA Agent, um sistema de aprendizado por reforço em escala de agente projetado para ensinar modelos de linguagem a gerar kernels CUDA de alto desempenho.

Este trabalho aborda uma fraqueza de longa data na geração de código de baixo nível por IA. Modelos de linguagem de ponta frequentemente conseguem gerar código CUDA que compila e executa corretamente, mas, para cargas de trabalho de GPU de nível de produção, a correção não é suficiente. Os kernels gerados também precisam competir com implementações geradas por compiladores altamente otimizados.

O CUDA Agent é construído exatamente em torno desse segundo problema: não apenas gerar código CUDA válido, mas também aprender a iterativamente analisar, validar e otimizar os kernels até alcançar melhorias significativas de desempenho em tempo de execução.

Por que apenas gerar código CUDA correto não é suficiente

Antes da introdução do aprendizado por reforço com agente, o modelo base Seed1.6 já demonstrava capacidade considerável de programação CUDA no KernelBench.

Nos 250 tarefas de KernelBench Nível 1–3 utilizadas pelos pesquisadores, o Seed1.6 alcançou uma taxa de aprovação de 74,0%. Em outras palavras, ele conseguia gerar soluções funcionalmente corretas para a maior parte das tarefas do benchmark.

Mas o desempenho era outra história.

Apenas 27,2% dos kernels gerados pelo modelo base eram mais rápidos que o torch.compile, com uma velocidade média geométrica de apenas 0,69× em relação ao baseline do compilador.

Modelo	Taxa de aprovação	Mais rápido que torch.compile	Velocidade média geométrica em relação ao torch.compile
Seed1.6 modelo base	74,0%	27,2%	0,69×
CUDA Agent	98,8%	96,8%	2,11×

Essa lacuna destaca o principal desafio na geração automática de kernels de GPU.

Os modelos de linguagem conseguem entender a sintaxe do CUDA, mas ainda podem produzir acessos de memória ineficientes, lançamentos excessivos de kernels, escolhas ruins de tiling, e interrupções desnecessárias.

Perfil de desempenho

• Orientação de otimização específica para CUDA
• Ambiente de sandbox restrito
• Sinais de recompensa vinculados aos resultados reais de execução

O arquivo SKILL.md também restringe atalhos que poderiam distorcer os resultados do benchmark. Por exemplo, o agente não pode simplesmente recorrer a operações arbitrárias do PyTorch em implementações de kernels personalizadas.

Isso é importante porque sistemas de aprendizado por reforço podem encontrar maneiras de maximizar a recompensa sem realmente resolver o problema de otimização pretendido.

Treinamento PPO expande o agente para trajetórias longas de otimização

Os pesquisadores treinaram o CUDA Agent usando otimização de política proximal (PPO).

Um dos principais desafios é que a otimização de GPU é, por natureza, uma tarefa de horizonte longo. O modelo pode precisar passar por múltiplas rodadas de edição de código, compilação, depuração, análise de desempenho e otimização adicional até alcançar o resultado desejado.

Portanto, o artigo adota diferentes comprimentos de contexto e estágios de treinamento para estabilizar o aprendizado.

O modelo base é o Seed1.6, um modelo de mistura de especialistas com 230 bilhões de parâmetros totais e 23 bilhões de parâmetros ativos.

Para o aprendizado por reforço com agente:

• A janela de contexto é de 131.072 tokens.
• O replay de treinamento permite até 150 rodadas de interação do agente.
• A avaliação permite até 200 rodadas de interação do agente.
• O modelo é treinado com 150 passos de RL.
• Os modelos actor e critic usam uma estratégia de aquecimento em múltiplos estágios antes da otimização longa completa.

Isso é mais complexo do que simplesmente ajustar o modelo em pares de prompts e respostas CUDA.

O objetivo é ensinar o modelo a melhorar seus próprios kernels por meio de interação repetida com o feedback de execução.

A sandbox separa compilação de análise de GPU

A medição confiável de desempenho é crucial, pois a velocidade de execução faz parte da recompensa.

Por isso, os pesquisadores construíram uma arquitetura de sandbox com recursos CPU–GPU desacoplados.

Uma sandbox de terminal baseada em Docker lida com tarefas orientadas à CPU (como compilação), enquanto a validação e a análise de desempenho são despachadas para um pool dedicado de sandboxes de GPU contendo 128 GPUs NVIDIA H20.

Essa separação tem um propósito duplo.

Primeiro, isola a compilação e a interação do agente dos benchmarks de GPU. Segundo, a alocação dedicada de GPUs reduz a interferência entre cargas de trabalho concorrentes, tornando as medições de latência mais estáveis.

Isso é importante para o aprendizado por reforço: se as medições de tempo forem afetadas por ruído, o modelo recebe sinais de recompensa não confiáveis e pode aprender comportamentos de otimização incorretos.

CUDA Agent atinge 98,8% de taxa de aprovação e supera o compilador em 96,8% das tarefas

O sistema final foi avaliado em 250 tarefas dos níveis 1 a 3 do KernelBench.

O CUDA Agent alcançou:

• 98,8% de taxa de aprovação geral
• 98,4% mais rápido que a execução imediata do PyTorch
• 96,8% mais rápido que o torch.compile
• Aceleração média geométrica de 2,60× em relação à execução imediata
• Aceleração média geométrica de 2,11× em relação ao torch.compile

Esses resultados representam uma melhoria substancial em relação ao ponto de partida do Seed1.6.

O modelo não apenas se tornou melhor em gerar código que passa nos testes de correção, mas os kernels que ele gera também têm maior probabilidade de superar os baselines do compilador.

Resultados por nível de dificuldade do KernelBench

O ganho de desempenho se mantém forte nos três níveis de dificuldade do KernelBench.

Dificuldade KernelBench	Taxa de aprovação	Proporção mais rápida que torch.compile	Aceleração média geométrica vs. torch.compile
Nível 1	100,0%	97,0%	1,87×
Nível 2	100,0%	100,0%	2,80×
Nível 3	94,0%	90,0%	1,52×

O desempenho no Nível 2 é particularmente notável.

Essas tarefas envolvem sequências de operadores, onde as oportunidades de otimização trazidas por fusão e movimentação de memória não são necessariamente aproveitadas de forma eficaz pelas heurísticas estáticas dos compiladores.

Os autores do artigo acreditam que o otimizador aprendido iterativamente consegue explorar um espaço mais amplo de estratégias de implementação, incluindo padrões de acesso à memória específicos de hardware, escolhas de tiling e decisões de fusão.

Dados de treinamento serão disponibilizados publicamente

Atualmente, os pesos completos do modelo treinado não estão incluídos nos materiais publicamente liberados do projeto.

No entanto, os pesquisadores já disponibilizaram vários componentes importantes do pipeline de treinamento.

CUDA-Agent-Ops-6K

O conjunto de dados CUDA-Agent-Ops-6K contém 6.000 tarefas sintéticas de treinamento em nível de operador.

O processo de construção inclui:

1. Coleta de operadores semente de bibliotecas como torch, transformers, etc.
2. Uso de LLMs para combinar múltiplos operadores em tarefas de fusão mais complexas.
3. Filtragem das tarefas geradas por meio de verificações baseadas em resultados de execução.

A fase de filtragem remove casos com aleatoriedade, simplicidade excessiva, invalidez ou similaridade excessiva com as tarefas de avaliação.

O conjunto de dados está publicado no Hugging Face sob a licença CC BY 4.0.

SKILL.md e ambiente do agente

O repositório GitHub também inclui o arquivo de instruções SKILL.md voltado para CUDA, além do ambiente de trabalho do agente.

Esses materiais documentam o fluxo de trabalho de otimização, as ferramentas disponíveis, as restrições e a configuração do ambiente de execução usada para orientar o agente.

Mecanismo de recompensa e regime de treinamento

O artigo e o repositório descrevem o design de recompensa, os estágios de aquecimento, a inicialização do crítico e o regime de treinamento baseado em PPO para estabilizar o processo de otimização de longo horizonte.

Isso é particularmente útil para pesquisadores interessados no treinamento de agentes para programação de sistemas, e não apenas na reprodução das pontuações finais do benchmark.

Por que isso importa além do KernelBench

A otimização de kernels CUDA sustenta grande parte da infraestrutura moderna de IA.

Kernels mais rápidos podem melhorar:

• Throughput de treinamento de modelos
• Latência de inferência de LLMs
• Utilização de GPU
• Eficiência de custo de serviços
• Pipelines de IA em tempo real
• Cargas de trabalho de computação numérica de alto desempenho

O artigo original da AIBase aponta sua importância para a IA

Aplicações potenciais em áreas como infraestrutura, serviços de inferência de modelos de grande porte, direção autônoma e negociação quantitativa — domínios em que diferenças mínimas de latência podem ser cruciais.

O resultado do CUDA Agent não significa que os compiladores tradicionais estejam obsoletos.

O torch.compile continua sendo uma linha de base automática poderosa, e a própria avaliação do CUDA Agent depende de ambientes de benchmark controlados e de um ambiente de GPU específico.

As conclusões demonstradas por esta pesquisa, embora mais restritas, permanecem importantes: com feedback de execução suficiente e aprendizado por reforço dedicado, os modelos de linguagem conseguem aprender estratégias de otimização que superam a linha de base de compiladores estáticos na grande maioria das tarefas de kernels de GPU testadas.

Isso significa transformar a engenharia de desempenho de baixo nível em um domínio viável para aprendizado baseado em agente, e não apenas em geração de código em uma única etapa.

Perguntas Frequentes

O que é o CUDA Agent?

O CUDA Agent é um sistema de aprendizado por reforço em larga escala para agentes, desenvolvido pela ByteDance Seed, pelo Instituto de Pesquisa da Indústria de Inteligência Artificial da Universidade Tsinghua (AIR) e pelo seu SIA-Lab conjunto. Ele treina modelos de linguagem para escrever, validar, analisar e otimizar kernels CUDA de forma iterativa.

Em qual modelo o CUDA Agent é baseado?

A pesquisa utiliza o Seed1.6 como modelo base. O artigo o descreve como um modelo de especialistas mistos (MoE), com um total de 230 bilhões de parâmetros, dos quais 23 bilhões são ativados.

O que é o KernelBench?

O KernelBench é um benchmark aberto usado para avaliar se modelos de linguagem conseguem gerar kernels de GPU corretos e eficientes para programas PyTorch. O CUDA Agent foi avaliado em 250 tarefas que abrangem os níveis 1 a 3 do KernelBench.

Quão mais rápido o CUDA Agent é em comparação com o torch.compile?

Em todo o conjunto de benchmarks, o CUDA Agent alcançou uma aceleração média geométrica de 2,11 vezes em relação ao torch.compile. Seus kernels gerados foram mais rápidos que a linha de base do compilador em 96,8% das tarefas avaliadas.

O CUDA Agent usa aprendizado por reforço?

Sim. O sistema utiliza PPO, combinado com aquecimento em múltiplos estágios, pré-treinamento do modelo de valor, design robusto de recompensas e trajetórias longas de agentes com múltiplas rodadas.

Quanto tempo pode durar uma trajetória de otimização do CUDA Agent?

O rollout de treinamento permite até 150 rodadas de agente, enquanto a avaliação permite até 200 rodadas. A janela de contexto de treinamento do agente é de 131.072 tokens.

O CUDA Agent é open source?

O projeto publicou seu repositório no GitHub, o ambiente do agente, o SKILL.md, a receita de treinamento e o conjunto de dados CUDA-Agent-Ops-6K. Os pesos completos do modelo treinado não estão incluídos na divulgação pública atual.

O que é o CUDA-Agent-Ops-6K?

O CUDA-Agent-Ops-6K é um conjunto de dados com 6.000 tarefas sintéticas de treinamento de CUDA em nível de operador. Ele foi projetado especificamente para aprendizado por reforço em larga escala com agentes e inclui etapas de filtragem para garantir que as tarefas sejam executáveis, determinísticas, não triviais e separadas do conjunto de avaliação do KernelBench.

Ferramentas Relacionadas

• CUDA Agent: Página oficial do projeto do sistema de geração de kernels CUDA da ByteDance Seed e do Instituto de Pesquisa da Indústria de Inteligência Artificial da Universidade Tsinghua.
• CUDA Agent no GitHub: Repositório oficial contendo o ambiente do agente, o SKILL.md e os materiais publicados do projeto.
• CUDA-Agent-Ops-6K: Conjunto de dados oficial de treinamento com 6.000 amostras, publicado junto ao projeto.
• KernelBench: Benchmark aberto para avaliar kernels de GPU gerados por modelos de linguagem de grande porte.
• PyTorch: Framework de aprendizado profundo cujo modo eager e a execução via torch.compile são linhas de base fundamentais na pesquisa.
• NVIDIA CUDA: Documentação oficial da NVIDIA sobre programação CUDA e desenvolvimento de kernels de GPU.

Links Relacionados

• Página do projeto CUDA Agent: Visão geral oficial, resultados de benchmark, gráficos do projeto, artigo, código e links para os conjuntos de dados.
• Artigo do CUDA Agent (arXiv): Artigo de pesquisa completo, "Aprendizado por Reforço em Larga Escala com Agentes para Geração de Kernels CUDA de Alto Desempenho".
• Repositório do CUDA Agent no GitHub: Repositório de código-fonte com o ambiente publicado e os arquivos do projeto.
• Conjunto de dados CUDA-Agent-Ops-6K: Conjunto de dados oficial no Hugging Face com 6.000 tarefas sintéticas de operadores.
• Repositório do KernelBench no GitHub: Ambiente de benchmark e avaliação para geração eficiente de kernels de GPU.
• Documentação do torch.compile no PyTorch: Documentação oficial do PyTorch sobre a linha de base do compilador utilizada na pesquisa.
• Guia de Programação CUDA C++ da NVIDIA: Guia oficial da NVIDIA para programação e otimização CUDA.

Resumo

O CUDA Agent resolve uma fraqueza específica no código de sistema gerado por modelos de linguagem de grande porte: o código CUDA gerado não precisa apenas ser correto, mas também genuinamente mais rápido que as alternativas produzidas por compiladores.

Partindo do Seed1.6, os pesquisadores utilizaram um ambiente de agente específico para CUDA, feedback de execução, design robusto de recompensas e treinamento PPO de horizonte longo para elevar a taxa de aprovação de 74,0% para 98,8%, e a proporção de kernels mais rápidos que o torch.compile de 27,2% para 96,8%.

O projeto também publicou um conjunto de dados de treinamento com 6.000 tarefas, o SKILL.md, materiais do ambiente do agente e a receita de treinamento, fornecendo uma base concreta para que pesquisadores avancem no campo de otimização de GPUs aprendida.

A principal contribuição do CUDA Agent é demonstrar que a engenharia de desempenho pode ser aprendida por meio de ciclos iterativos de agentes — e não apenas aproximada por geração de código em uma única etapa.