Kernel-Smith: A Unified Recipe for Evolutionary Kernel Optimization

Resumo Técnico: Kernel-Smith

1. O Problema

A geração de kernels de GPU de alto desempenho é crítica para o throughput de sistemas de modelos grandes (LLMs) e cargas de trabalho científicas. Embora os Modelos de Linguagem (LLMs) tenham demonstrado capacidade em programação, a geração autônoma e confiável de kernels otimizados permanece um desafio significativo. As abordagens atuais baseadas em LLMs enfrentam duas barreiras principais:

• Limitação da Geração "One-Shot": A maioria dos sistemas tenta gerar um kernel correto e rápido em uma única tentativa, falhando em explorar o vasto espaço de soluções que requer múltiplas iterações e refinamentos.
• Instabilidade de Avaliação e Diversidade: A otimização de kernels é um processo não convexo. Sistemas de agentes multi-turno existentes tendem a ancorar-se em decisões iniciais, limitando a diversidade de exploração. Além disso, o ruído na medição de desempenho (tempo de execução) em GPUs pode levar a seleção de soluções subótimas ou eliminação de promissoras.
• Correção vs. Desempenho: Garantir a correção funcional (output numérico) e o alto desempenho (aceleração) são capacidades distintas que raramente são otimizadas simultaneamente de forma iterativa e robusta.

2. Metodologia

O Kernel-Smith propõe um framework unificado que combina um agente evolutivo guiado por avaliação estável com uma receita de pós-treinamento orientada à evolução.

• Framework de Agente Evolutivo (OpenEvolve Adaptado):

  • Em vez de refinar um único rastro de diálogo, o sistema mantém uma população de candidatos executáveis.
  • Utiliza um mecanismo de feedback estruturado que avalia três dimensões: compilação, correção (consistência numérica com a referência PyTorch) e aceleração (speedup).
  • Implementa ilhas evolutivas e um grid MAP-Elites para preservar a diversidade de soluções de alta qualidade, evitando o colapso para uma única solução.
  • Redução de Ruído: Para garantir estabilidade na avaliação, o sistema utiliza gráficos de execução (CUDAGraph), múltiplas medições com remoção de outliers e execuções de warm-up, restringindo a flutuação do tempo de execução a menos de 1%.
  • Detecção de "Hacking": O sistema impede que modelos contornem a geração de kernels reais chamando operadores PyTorch nativos (falsos positivos de aceleração) ou aplicando otimizações triviais sem valor prático.

• Receita de Treinamento (Post-Training):

  • O objetivo não é treinar o modelo para gerar kernels perfeitos de primeira, mas para atuar como um otimizador local forte dentro do loop evolutivo.
  • Síntese de Dados: Coleta de módulos PyTorch de repositórios GitHub reais, limpeza automática de dependências e geração de casos de teste.
  • Filtragem de Trajetórias: As trajetórias de evolução de longo prazo são convertidas em sinais de supervisão centrados em passos. Apenas as revisões que preservam a correção e trazem ganhos significativos de desempenho (high-gain revisions) são retidas.
  • RL (Reinforcement Learning): Utiliza o algoritmo GRPO (Group Relative Policy Optimization). O modelo é treinado para gerar variações que superam um código "pai" (baseline), usando a razão de aceleração como recompensa. Isso ensina o modelo a compor melhorias atômicas ao longo de múltiplas rodadas.

• Backends de Avaliação:

  • O framework é agnóstico ao hardware, com implementações específicas para NVIDIA GPUs (Triton) e MetaX GPUs (Maca), permitindo adaptação a plataformas heterogêneas.

3. Principais Contribuições

1. Framework Unificado de Otimização Evolutiva: Integração de um agente de busca populacional com um backend de avaliação robusto e estável, superando as limitações de abordagens puramente baseadas em diálogo ou geração única.
2. Receita de Treinamento Orientada à Evolução: Uma nova metodologia de pós-treinamento (SFT + RL) que foca em "passos de melhoria" dentro de trajetórias evolutivas, transformando o LLM em um otimizador iterativo eficaz em vez de um gerador estático.
3. Validação em Múltiplas Plataformas: Sucesso demonstrado tanto no backend NVIDIA (Triton) quanto no MetaX (MACA), provando a adaptabilidade do método.
4. Contribuições para o Mundo Real: O framework não se limita a benchmarks; gerou Pull Requests (PRs) aceitos e mesclados em sistemas de produção de alto impacto, como SGLang e LMDeploy, e otimizou módulos de pesquisa recentes (Engram do DeepSeek).

4. Resultados

• KernelBench (NVIDIA Triton):

  • O modelo Kernel-Smith-235B-RL alcançou o estado da arte (SOTA) no KernelBench, obtendo a melhor razão média de aceleração (3.70x) entre todos os níveis de dificuldade.
  • Superou modelos proprietários de ponta como Gemini-3.0-pro e Claude-4.6-opus, bem como modelos open-weights como DeepSeek-v3.2-Speciale e Qwen3.
  • No nível de dificuldade 2, alcançou uma aceleração média impressionante de 7.77x, superando o Claude-4.6-opus (5.83x).
  • A curva de crescimento de desempenho ao longo das gerações evolutivas mostrou que o modelo se beneficia mais efetivamente de compute adicional em tempo de teste (test-time compute) do que os concorrentes.

• Backend MetaX (MACA):

  • O Kernel-Smith-MACA-30B superou contrapartes de grande escala como DeepSeek-V3.2-think e Qwen3-235B-2507-think, demonstrando eficácia em arquiteturas de hardware alternativas.

• Aplicações no Mundo Real:

  • SGLang: Geração de um kernel fundido para normal_decode_set_metadata, resultando em 4.78x de aceleração no nível do operador e melhorias mensuráveis de latência em cenários de serviço completo.
  • LMDeploy: Otimização do módulo de roteamento MoE, com 1.36x de aceleração no operador e ganhos de throughput de 1.85% a 3.00% em inferência completa.
  • DeepSeek Engram: Otimização de um módulo de pesquisa recente, alcançando 14.59x de aceleração, demonstrando a capacidade de generalização para código não visto em benchmarks tradicionais.

5. Significado e Impacto

O Kernel-Smith representa um avanço fundamental na automação da engenharia de sistemas de alto desempenho. Ele demonstra que:

• A otimização de kernels não deve ser tratada como um problema de geração de código único, mas como um processo de busca evolutiva que requer feedback estável e treinamento específico para iteração.
• É possível transferir o sucesso de ambientes controlados (benchmarks) para sistemas de produção reais, contribuindo diretamente para a eficiência de infraestrutura de IA.
• A combinação de avaliação estável (para garantir a dinâmica de busca confiável) e treinamento centrado em passos evolutivos (para melhorar a qualidade de cada iteração) é a chave para desbloquear o potencial de otimização dos LLMs em tarefas de sistemas complexos.

Este trabalho estabelece um novo paradigma onde LLMs atuam não apenas como geradores de código, mas como engenheiros de otimização iterativos, capazes de explorar espaços de implementação complexos e entregar ganhos de desempenho mensuráveis em hardware heterogêneo.