Resumo Técnico: Kernel Forge

Definição do Problema

Modelos de aprendizado de máquina dependem cada vez mais de um pequeno conjunto de kernels de computação (ex: multiplicação de matrizes, convolução, normalização) para a maior parte de seu tempo de execução. Embora a otimização desses kernels seja um caminho direto para reduzir a latência e o custo, os métodos tradicionais exigem que engenheiros especialistas escrevam código de baixo nível para GPU manualmente, um processo que é lento, caro e intensivo em termos de especialização.

Sistemas agênticos recentes baseados em Large Language Models (LLMs) têm mostrado promessa na geração e otimização de kernels com menos esforço humano. No entanto, as ferramentas existentes sofrem de quatro limitações críticas quando aplicadas a cenários do mundo real:

1. Lacuna de Integração: Elas frequentemente emitem código CUDA isolado, deixando aos desenvolvedores a tarefa de reintegrar manualmente os kernels otimizados em modelos PyTorch não modificados.
2. Estreiteza de Carga de Trabalho: Muitas ferramentas focam exclusivamente em cargas de trabalho de LLM ou benchmarks isolados, falhando em abordar a diversidade de formas, layouts e argumentos encontrados em modelos de visão, difusão e áudio.
3. Limitações de Busca: As abordagens atuais frequentemente dependem de refinamento linear ou busca do tipo beam search, que podem podar caminhos de otimização promissores que inicialmente parecem subótimos ou falham em explorar o espaço de soluções de forma eficaz.
4. Observabilidade: A maioria dos sistemas são protótipos de pesquisa ou scripts que carecem de interfaces gráficas integradas para monitorar o progresso, inspecionar candidatos ou depurar falhas.

Além disso, os benchmarks existentes frequentemente avaliam correção e latência usando tensores gerados aleatoriamente e invocações de kernels isoladas. Isso falha em capturar condições específicas de modelos do mundo real, como formas de tensores, distribuições de ativação e comportamentos de operadores vizinhos, levando a uma lacuna entre o desempenho do kernel isolado e a execução real do modelo.

Metodologia: Kernel Forge

O Kernel Forge é um harness agêntico de código aberto, de ponta a ponta, projetado para otimizar execuções reais de modelos PyTorch sem exigir que os usuários modifiquem o código do modelo. O sistema opera através de um pipeline de múltiplos estágios:

1. Ingestão de Modelo e Captura de Operadores

O Kernel Forge aceita qualquer modelo PyTorch padrão e entradas de exemplo fornecidas pelo usuário. Ele executa o modelo para capturar os operadores específicos que rodam durante a inferência. Em vez de tratar operadores como famílias genéricas, ele os agrupa em variantes compatíveis com o tempo de execução baseadas em padrões de invocação concretos (formas de tensores, dtypes, layouts, argumentos e posicionamento de dispositivo). Essas variantes são encapsuladas como "cartões de operador" (operator cards), que servem como alvos de otimização e preservam a execução eager do PyTorch como referência de correção.

2. Geração de Kernel Agêntica

O sistema emprega um agente baseado em LLM para gerar kernels CUDA especializados.

• Entrada: O agente recebe um cartão de operador e uma configuração de hardware/tempo de execução.
• Processo: Para o primeiro candidato, o LLM gera código fonte CUDA e um invólucro de lançamento (launch wrapper). Em iterações subsequentes, o agente recebe um kernel pai e sua linhagem do controlador de busca para propor revisões visando reduzir a latência enquanto preserva o comportamento.
• Restrição: O agente é instruído a produzir código CUDA customizado em vez de delegar a bibliotecas existentes (ex: cuDNN, cuBLAS), embora o sistema reconheça que o código gerado possa, por vezes, envolver efetivamente esses backends.

3. Validação e Reparo

Os candidatos gerados passam por um ciclo rigoroso de validação:

• Compilação: O código fonte é compilado com nvcc. Erros de compilação acionam um ciclo de reparo limitado.
• Execução: Candidatos bem-sucedidos são lançados com as entradas capturadas. Falhas de tempo de execução (travamentos, acesso ilegal à memória) são enviadas de volta ao ciclo de reparo.
• Correção Numérica: Os candidatos são comparados com as saídas eager do PyTorch capturadas durante a ingestão, usando tolerâncias numéricas específicas do operador. Apenas candidatos que passam por este filtro específico de carga de trabalho entram na árvore de busca.

4. Controlador de Monte Carlo Tree Search (MCTS)

Ao contrário do refinamento linear, o Kernel Forge usa MCTS para explorar múltiplos caminhos de otimização.

• Estrutura da Árvore: O controlador mantém uma árvore de revisão persistente onde os nós representam kernels candidatos, armazenando código fonte, linhagem, contagem de visitas e latência medida.
• Estratégia de Seleção: O controlador utiliza uma versão de minimização do Upper Confidence bounds applied to Trees (UCT) combinada com alargamento progressivo (progressive widening). Ele equilibra a exploração de ramos menos visitados com a explotação de caminhos promissores, permitindo que o sistema revisite candidatos que podem não mostrar ganhos imediatos, mas que poderiam levar a soluções melhores posteriormente.
• Feedback: As medições de latência são propagadas para cima na árvore para guiar a seleção do próximo kernel pai para refinamento.

5. Despacho Guardado e Empacotamento

A saída final é um pacote guardado (.cast) que integra os kernels otimizados de volta ao modelo.

• Lógica Guardada: Um kernel gerado é despachado apenas se for válido, compatível com o padrão de chamada capturado, numericamente correto e mais rápido que a latência medida do PyTorch eager para aquela variante específica.
• Fallback: Se um candidato falhar em qualquer verificação ou for mais lento, o sistema retorna automaticamente para o caminho original do PyTorch eager. Isso garante segurança e evita regressões de desempenho.
• Auditabilidade: O sistema distingue entre CUDA customizado gerado, wrappers de backend e caminhos de fallback, fornecendo uma trilha de auditoria clara do que foi realmente otimizado.

Interface do Usuário

O Kernel Forge fornece tanto uma Interface Gráfica (GUI) quanto uma Interface de Linha de Comando (CLI). Essas ferramentas permitem que os usuários monitorem o progresso da otimização, inspecionem kernels candidatos, visualizem a árvore de revisão MCTS, depurem falhas e comparem resultados entre diferentes provedores de modelos.

Principais Contribuições

1. Sistema Agêntico de Ponta a Ponta: A introdução do Kernel Forge, uma ferramenta que otimiza execuções reais de modelos PyTorch capturando operadores, gerando kernels CUDA especializados, explorando múltiplos caminhos via MCTS e integrando automaticamente os resultados de volta ao modelo.
2. Observabilidade e Usabilidade: A provisão de uma GUI e CLI para monitorar, inspecionar e depurar o processo de otimização, preenchendo a lacuna entre protótipos de pesquisa e ferramentas práticas.
3. Avaliação de Mundo Real: Uma avaliação através de quatro modelos PyTorch diversos (ResNet-50, Stable Diffusion 3.5 Medium, Gemma 4 E2B e Qwen 3.5 35B-A3B) representando cargas de trabalho de visão, difusão e LLM.
4. Lançamento Open-Source: O lançamento da ferramenta para apoiar a pesquisa reprodutível e a adoção prática da otimização de kernels agêntica.

Resultados

Avaliado em um NVIDIA DGX Spark com GPU GB10, o Kernel Forge foi testado com 50 iterações de otimização por kernel. Os resultados destacam tanto o potencial quanto os desafios da otimização impulsionada por LLM em cargas de trabalho reais:

• Ganhos de Desempenho: O Kernel Forge otimizou com sucesso 14 kernels para superar o modo PyTorch eager. Os ganhos de velocidade notáveis incluem:
  • 1.52× em adaptive_avgpool2d no ResNet-50.
  • 1.70× em group_norm no Stable Diffusion 3.5 Medium.
  • 2.83× em softmax no Gemma 4 E2B.
  • 1.54× em softmax no Qwen 3.5 35B-A3B.
• Especificidade de Operador: Ganhos de velocidade significativos foram encontrados predominantemente em caminhos nativos/open-source do PyTorch (ex: tensor_iadd, relu, softmax). Em contraste, operadores apoiados por implementações proprietárias ou de fornecedores maduros (ex: conv2d, linear, scaled_dot_product_attention) frequentemente permaneceram no nível ou abaixo do desempenho do PyTorch eager.
• Impacto no Tempo de Execução: Os operadores que alcançaram os maiores ganhos locais frequentemente representavam uma pequena fração do tempo de execução total capturado. Por outro vez, os contribuintes dominantes do tempo de execução eram frequentemente os caminhos de backend altamente otimizados que eram difíceis de melhorar.
• Custo de Busca: A avaliação observou que os custos de busca (uso de API de LLM) aumentaram significativamente em contagens de iteração mais altas (opt50), mas isso nem sempre correlacionou com ganhos proporcionais no desempenho agregado da região do operador, particularmente para operadores dominantes.

Significância e Alegações

O artigo afirma que o Kernel Forge demonstra que a otimização de kernel impulsionada por LLM pode ser aplicada efetivamente além de benchmarks isolados para operadores capturados de modelos PyTorch completos.

Os autores enfatizam a importância da avaliação consciente da implantação. Eles argumentam que, embora os LLMs possam gerar kernels rápidos isolados, o verdadeiro valor reside em como esses kernels se comportam dentro de execuções de modelos reais. O mecanismo de "despacho guardado" é apresentado como uma inovação crítica, permitindo que o sistema aplique kernels gerados apenas onde há evidências de benefício, preservando a confiabilidade das implementações existentes do framework.

O artigo conclui que a otimização de kernel agêntica deve ser avaliada não apenas por sua capacidade de gerar código rápido, mas por sua capacidade de navegar nos compromissos entre profundidade de otimização, custo de busca e as características específicas de cargas de trabalho do mundo real, particularmente a prevalência de implementações de backend maduras para operadores dominantes.