AutoKernel: sistema automático de pesquisa de kernels de GPU

• Aplica o Autoresearch de Karpathy à otimização de kernels de GPU
• Ferramenta automática de pesquisa de kernels de GPU que, ao receber um modelo PyTorch como entrada, otimiza automaticamente kernels em Triton ou CUDA C++
• Fornece um pipeline totalmente automatizado para perfilar, extrair, otimizar e validar os kernels gargalo do modelo
• Ajusta prioridades com orquestração baseada na lei de Amdahl, e cada experimento leva cerca de 90 segundos, permitindo centenas de iterações durante a noite
• Integrado ao KernelBench, realiza 50 a 300 execuções iterativas em mais de 250 problemas, possibilitando uma exploração sistemática em vez de geração pontual
• Compatível com backend duplo Triton e CUDA C++, validação com prioridade na corretude e estrutura de modificação de arquivo único, dando suporte a uma otimização de kernels eficiente e reproduzível

---

Como funciona

• O AutoKernel recebe um modelo PyTorch como entrada e executa as seguintes etapas
  • Profiling para identificar os kernels gargalo na GPU
  • Extração para separá-los em kernels independentes em Triton ou CUDA C++
  • Loop automático de otimização para repetir modificações, benchmark e manter/reverter mudanças
  • Validação de corretude e, depois, relatório do ganho total de velocidade
• Com base nas instruções incluídas em program.md, modifica kernel.py e usa bench.py para realizar uma verificação de corretude em 5 etapas e análise roofline
• Cada experimento leva cerca de 90 segundos, permitindo cerca de 40 execuções por hora e aproximadamente 320 durante uma noite

Execução e componentes

• Ambiente necessário: GPU NVIDIA (H100/A100/RTX 4090), Python 3.10+, pacote uv
• Principais scripts
  • profile.py: classifica kernels com base no tempo de GPU
  • extract.py: extrai os principais kernels gargalo
  • bench.py: benchmark de desempenho e corretude em 5 etapas
  • orchestrate.py: agendamento de múltiplos kernels com base na lei de Amdahl
  • verify.py: validação do modelo completo e relatório de ganho de velocidade
• program.md inclui uma estratégia de otimização em 6 etapas, tratamento de conflitos e um framework de tomada de decisão, permitindo execução automática de longa duração

Kernels suportados e modelos de exemplo

• Suporta 9 tipos de kernel: matmul, softmax, layernorm, rmsnorm, flash_attention, fused_mlp, cross_entropy, rotary_embedding, reduce
• Cada kernel inclui implementação de referência em PyTorch (reference.py) e versões iniciais em Triton/CUDA (kernels/, kernels/cuda/)
• Modelos de exemplo: GPT-2, LLaMA (incluindo 7B), BERT-base e modelos personalizados
• Modelos do HuggingFace também podem ser integrados com uv sync --extra models

Integração com o KernelBench

• Integrado ao KernelBench (Stanford Scaling Intelligence Lab) para executar benchmark padrão de kernels de GPU gerados por IA
• O AutoKernel faz 50 a 300 execuções iterativas por problema para explorar sistematicamente o espaço de otimização
• Principais ferramentas
  • bridge.py: carrega problemas e gera kernels iniciais
  • bench_kb.py: avalia corretude e desempenho
  • scorer.py: calcula a pontuação total por nível
  • program_kb.md: instruções do agente dedicadas ao KernelBench

Exportação para o HuggingFace Hub

• Os kernels otimizados podem ser exportados para o HuggingFace Hub e carregados facilmente com get_kernel()
• export_hf.py oferece suporte para exportação e upload de kernels CUDA

Princípios de design

• Backend duplo Triton + CUDA C++: Triton para iteração rápida, CUDA para desempenho máximo
• Corretude em primeiro lugar: se o resultado não corresponder ao do PyTorch, a mudança é revertida imediatamente
• Orquestração baseada na lei de Amdahl para definir prioridades de acordo com a contribuição ao desempenho total
• Estrutura de modificação de arquivo único (kernel.py) para facilitar rastreamento e recuperação de mudanças
• Logs em TSV (results.tsv) registram os resultados dos experimentos de forma simples e legível

Formato de registro dos resultados

• Em results.tsv, são registrados o número de cada experimento, tipo de kernel, throughput (TFLOPS), latência, proporção em relação ao pico da GPU, ganho de velocidade em relação ao PyTorch, corretude, uso de VRAM e descrição

Contexto do projeto

• Inspirado no conceito de autoresearch de Andrej Karpathy, aplicando a estrutura de um agente autônomo de IA para pesquisa com LLM à otimização de kernels de GPU
• A integração com o KernelBench é baseada em pesquisas do Stanford Scaling Intelligence Lab, e o AutoKernel realiza otimização iterativa em vez de geração pontual
• O projeto foi desenvolvido pela equipe Forge da RightNow AI e é disponibilizado sob a licença MIT

[GN⁺]

Comentários do Hacker News

• Projeto realmente muito legal
  Eu também estava criando algo parecido nas últimas semanas, mas era uma versão bem mais genérica e excessivamente projetada
  Essa abordagem, focada só no Triton e em kernels específicos, é simples e eficiente
  Mas o gráfico de progresso é confuso. Parece ser um benchmark de matmul fp16 4096x4096x4096, e embora diga haver uma melhoria de 1,31x sobre o cuBLAS, mostra apenas 187 TFLOPS, ou seja, 18,9% de utilização de pico
  O cuBLAS entrega um desempenho muito mais próximo do pico, então provavelmente há overhead de CPU ou outro gargalo
  Benchmark é realmente difícil, mas acho que essa área é extremamente promissora nos próximos 6 meses

• Se isso funcionar em algo como o llama.cpp, o ganho pode ser bem maior
  Existem vários kernels de quantização e diferentes configurações de hardware, e também há muitos usuários individuais, então a margem de ganho de eficiência deve ser grande
  Seria ótimo se este projeto contribuísse com isso

  • Faz sentido
    O llama.cpp tem vários kernels CUDA ajustados manualmente, como Q4_K_M, Q5_K_S e Q8_0, e cada um mira perfis de hardware diferentes
    Se fosse possível fazer otimização automática para cada GPU, isso seria uma mudança enorme
    Hoje, mesmo com o mesmo formato de quantização, a diferença de desempenho entre uma RTX 3090 e uma 5070 Ti é grande
    Ambientes com grande diversidade de hardware, como o llama.cpp, são exatamente onde a busca automática de kernels mais pode brilhar

• Muito legal!
  Eu estou adicionando a mesma funcionalidade para Apple Silicon
  No meu projeto autoresearch-everywhere, estou trabalhando para transformar o autoresearch em uma ferramenta séria

• Tem algo estranho aí
  Considerando GEMM fp16 4kx4kx4k, o cutlass é cerca de 3 vezes mais rápido

• Fiquei curioso se já fizeram benchmark comparando com escalonamento automático como o Ansor do TVM

• Isso é só o começo
  Como referência, o Google já fez algo parecido duas gerações de modelos atrás
  No post de blog de maio de 2025, AlphaEvolve, a empresa anunciou que acelerou em 23% um kernel central da arquitetura Gemini ao dividir uma grande multiplicação de matrizes em subproblemas menores, reduzindo o tempo de treinamento em 1%
  Agora estamos entrando numa era em que esse tipo de tecnologia também pode ser feito “em casa”
  Especialmente com o aumento recente do treinamento baseado em RL, melhorias na velocidade de inferência também devem se traduzir em melhorias na velocidade de treinamento

• Fico me perguntando quando otimizações desse tipo vão chegar também a runtimes de linguagens open source como Swift ou Rust, para extrair até a última gota de desempenho