AutoKernel: A IA que Otimiza GPUs PyTorch Sozinha!

Chega de Sofrer com Otimização de GPU: Essa IA Faz o Trabalho Enquanto Você Dorme!

Olá, pessoal! Aqui é o Lucas Tech, e hoje a gente vai mergulhar em uma novidade que promete mudar a vida de muita gente no mundo da Inteligência Artificial e Machine Learning. Sabe aquele trabalho supercomplexo e exaustivo de escrever código rápido para GPU, que só os verdadeiros magos da computação conseguem fazer? Então, a equipe da RightNow AI está chegando com uma solução open-source que, pasmem, quer automatizar TUDO isso com uma IA! É sério, prepare-se para ficar chocado e talvez até um pouco mais esperançoso para aquela noite de sono tranquila!

Basicamente, eles lançaram o AutoKernel, um framework que usa uma IA autônoma para otimizar os "kernels" de GPU em qualquer modelo PyTorch. A ideia é simples: você entrega o seu modelo antes de dormir e, quando acorda, tem kernels Triton muito mais rápidos – sem precisar ser um guru em otimização de GPU. Incrível, né?

Mas Por Que Otimizar Kernels de GPU é Tão Difícil?

A gente sabe que as GPUs são a espinha dorsal de qualquer modelo de IA moderno. Um "kernel" de GPU é tipo uma função que executa milhares de vezes em paralelo nos núcleos da placa. Quando você roda um modelo transformer gigante, como o LLaMA ou o GPT-2, a maior parte do tempo de processamento acontece dentro desses kernels, em operações como multiplicação de matrizes (matmul), softmax, normalização de camadas e atenção.

O problema é que tirar o máximo de performance desses kernels é uma arte e uma ciência. Você precisa pensar em mil coisas ao mesmo tempo: intensidade aritmética, como a memória é acessada, a pressão dos registradores, o tamanho dos blocos de processamento, sincronização de "warps" e até a seleção de instruções específicas dos "tensor cores". Essa combinação de habilidades leva ANOS para ser desenvolvida! Um único kernel matmul de alta performance pode ter mais de 200 linhas de código CUDA ou Triton, com dezenas de parâmetros interligados. E o pior: essa especialidade é rara, e o processo de ajuste manual não escala bem conforme as arquiteturas dos modelos evoluem.

É nesse cenário que o AutoKernel aparece como um verdadeiro divisor de águas!

O Segredo do AutoKernel: Um Loop de Aprendizado Inteligente

A grande sacada do AutoKernel é que o fluxo de trabalho de um engenheiro especialista em kernels é, na verdade, um loop bem simples: escreve um candidato, testa, se melhorou, mantém, se piorou, descarta e repete. O framework simplesmente automatiza esse loop!

Um agente de LLM (Large Language Model) modifica um único arquivo – o kernel.py. Depois, um sistema de benchmark verifica a correção e mede a velocidade. O resultado decide se a mudança é mantida ou não. E o mais legal: cada experimento vira um commit no Git! Isso significa que você tem um histórico completo e navegável com ferramentas Git padrão, e os resultados são salvos em um arquivo simples, fácil de ler e de ser "entendido" pelo agente.

Cada iteração leva cerca de 90 segundos: 30s para checar a correção, 30s para o benchmark de performance e 30s para o agente "pensar" e modificar o código. Com isso, uma rodada de 10 horas durante a noite pode gerar de 300 a 400 experimentos em vários kernels! O agente é guiado por um "manual de instruções" de 909 linhas, que codifica o conhecimento de especialistas em um "livro de receitas" de otimização em seis níveis. É tipo um bootcamp intensivo para a IA!

Priorizando o que Realmente Importa: Otimização Estratégica

Diferente de outros trabalhos que otimizam kernels isoladamente, o AutoKernel começa com um modelo PyTorch completo. Ele usa o torch.profiler para medir o tempo de GPU de cada kernel e, em seguida, prioriza os alvos de otimização usando a Lei de Amdahl. Essa lei diz que o ganho de velocidade total que você consegue é limitado pela porcentagem do tempo total que aquela parte representa. Ou seja, otimizar algo que já é rápido tem menos impacto do que otimizar algo que consome a maior parte do tempo!

O perfilador ainda detecta a GPU que você está usando (NVIDIA H100, A100, AMD MI300X, etc.) e, para GPUs desconhecidas, estima o desempenho, tornando o sistema útil em uma vasta gama de hardwares. O orquestrador inteligente sabe quando passar para o próximo kernel, evitando que o agente perca tempo em otimizações com retornos decrescentes.

Performance Sim, Mas e a Confiança? O Crivo de 5 Estágios!

Performance sem correção não serve para nada, não é mesmo? O AutoKernel leva isso super a sério. Cada kernel candidato passa por CINCO estágios de validação antes que qualquer ganho de velocidade seja registrado:

Teste de fumaça: Um teste rápido para pegar erros de compilação ou de dimensão em menos de um segundo.
Varredura de configurações: Testa de 8 a 10 configurações de entrada e três tipos de dados (FP16, BF16 e FP32) para pegar bugs que dependem do tamanho ou de limites.
Estabilidade numérica: Teste com entradas "adversárias" para garantir que o kernel não se quebre em condições extremas.
Verificação de determinismo: Roda a mesma entrada três vezes e exige saídas bit a bit idênticas, pegando problemas de concorrência ou operações não-determinísticas.
Casos de borda: Testa dimensões que não são potências de dois (tipo 1023, 4097) para expor bugs de mascaramento ou de preenchimento.

Os testes de tolerância são específicos para cada tipo de dado, garantindo que os resultados sejam sempre precisos. Nos testes em H100, todos os 34 kernels passaram sem falhas!

Versatilidade Sob o Capô: Triton e CUDA C++

O AutoKernel suporta tanto backends Triton quanto CUDA C++ dentro do mesmo framework.

Triton: É uma linguagem que lembra Python, feita sob medida para esse tipo de tarefa, e que compila na hora (JIT – Just In Time) em 1 a 5 segundos. Perfeito para o agente iterar rapidamente, ajustando tamanhos de bloco, contagens de warp, estágios de pipeline, etc. O Triton consegue alcançar de 80% a 95% da performance do cuBLAS (uma biblioteca superotimizada) para multiplicação de matrizes.
CUDA C++: É incluído para casos que exigem acesso direto a primitivas de baixo nível, como instruções WMMA dos tensor cores, carregamentos vetorizados e layouts de memória compartilhada sem conflitos.

Ambos os backends expõem a mesma interface kernel_fn(), então a infraestrutura de benchmark funciona de forma idêntica. O sistema cobre nove tipos de kernels, que são as operações dominantes em arquiteturas de transformers modernas, como matmul, flash_attention, softmax, layernorm e outros.

Resultados Que Impressionam (Principalmente em H100)!

Testado em uma GPU NVIDIA H100 80GB HBM3, contra o PyTorch eager e torch.compile com max-autotune, os resultados para kernels limitados pela memória são de cair o queixo!

RMSNorm: Alcançou 5,29x mais rápido que o eager e 2,83x mais rápido que o torch.compile nos maiores tamanhos testados.
Softmax: 2,82x mais rápido que o eager e 3,44x mais rápido que o torch.compile.
Cross-entropy: 2,21x mais rápido que o eager e 2,94x mais rápido que o torch.compile.

Esses ganhos vêm da fusão de várias operações do ATen (o backend do PyTorch) em um único kernel Triton, que minimiza o tráfego de dados na HBM (High Bandwidth Memory). O AutoKernel superou o torch.compile em 12 das 16 configurações testadas, mesmo com o torch.compile rodando seu próprio autotuning Triton. Isso mostra que as otimizações genéricas do TorchInductor nem sempre conseguem encontrar as estratégias superespecíficas que um kernel feito sob medida consegue explorar.

Matmul é um desafio maior, pois o cuBLAS do PyTorch já é extremamente otimizado. No entanto, em um tamanho específico (2048³), o AutoKernel ainda conseguiu bater o torch.compile em 1,55x!

E a prova social não para por aí: um kernel otimizado pelo AutoKernel ficou em primeiro lugar no ranking do vectorsum_v2 B200, com uma latência ligeiramente melhor que o segundo colocado. Outro usuário da comunidade relatou que um único "prompt" para o AutoKernel gerou um kernel de multiplicação de matrizes FP4 em Triton que superou o CUTLASS (código C++ otimizado à mão para tensor cores da NVIDIA!) em até 2,15x em algumas configurações na H100. Isso é simplesmente INACREDITÁVEL!

Key Takeaways (As Mensagens Chave!)

Automação Extrema: O AutoKernel transforma semanas de ajuste fino de GPU por especialistas em um processo autônomo que pode ser feito da noite para o dia, rodando centenas de experimentos sem intervenção humana.
Correção Acima de Tudo: Antes de qualquer ganho de velocidade ser registrado, cada kernel candidato passa por um crivo rigoroso de cinco estágios, garantindo que a otimização não comprometa a precisão dos resultados.
Ganhos Massivos em Kernels Limitados por Memória: Na H100 da NVIDIA, o AutoKernel demonstrou ganhos impressionantes em kernels como RMSNorm (5,29x), softmax (2,82x) e cross-entropy (2,21x), principalmente pela fusão de operações que minimizam o tráfego na HBM.
Otimização Inteligente e Estratégica: Em vez de otimizar kernels isoladamente, o AutoKernel perfila o modelo PyTorch completo e direciona seus esforços para as partes que mais impactam o tempo total de execução, seguindo a Lei de Amdahl. Isso garante que as melhorias se traduzam em ganhos reais no nível do modelo.

Minha Visão

Cara, o AutoKernel é um daqueles projetos que me deixam genuinamente animado com o futuro da IA. A otimização de código GPU sempre foi um gargalo gigantesco, uma barreira de entrada para muitos desenvolvedores e, francamente, um dreno de tempo para equipes experientes. Ver uma IA conseguir automatizar um processo tão complexo, com resultados que superam até códigos otimizados à mão, é um game-changer!

Isso não só vai democratizar o acesso à computação de alto desempenho, permitindo que mais pessoas criem modelos mais rápidos e eficientes, mas também vai liberar os engenheiros mais talentosos para se concentrarem em problemas mais criativos e de alto nível, em vez de ficarem ajustando parâmetros de baixo nível. É como ter um "super-engenheiro de otimização" trabalhando 24/7 para você. O impacto na velocidade de pesquisa e desenvolvimento de IA pode ser monumental.

E você, o que achou dessa novidade? Você acredita que o AutoKernel realmente vai democratizar a otimização de GPUs e mudar a forma como desenvolvemos IA? Deixa seu comentário aqui embaixo, vou adorar saber sua opinião!

Para quem quiser se aprofundar, confira o Artigo e o Repositório no GitHub.