A equipe de pesquisa da RightNow AI acaba de lançar o AutoKernel, uma ferramenta de código aberto que promete revolucionar a forma como otimizamos o desempenho de GPUs. Em um mundo onde a velocidade dos modelos de Inteligência Artificial é crucial, a escrita de código de GPU rápido é uma das especializações mais desafiadoras na engenharia de Machine Learning. O AutoKernel surge como uma solução para automatizar essa tarefa complexa, permitindo que desenvolvedores alcancem kernels Triton mais rápidos em modelos PyTorch, sem a necessidade de expertise aprofundada em otimização de hardware.
O que é AutoKernel e por que ele importa?
O AutoKernel é uma estrutura open-source que aplica um loop de agente autônomo baseado em Large Language Models (LLMs) para a otimização de GPU kernels em qualquer modelo PyTorch. A ideia é simples: você entrega seu modelo e o AutoKernel trabalha na otimização enquanto você se dedica a outras tarefas, entregando kernels mais eficientes no dia seguinte. Esta abordagem simplifica um processo que antes exigia anos de experiência, democratizando o acesso a otimizações de alto nível.
A complexidade da otimização de kernels de GPU
Um GPU kernel é uma função que executa paralelamente em milhares de núcleos de GPU. Operações como multiplicação de matrizes (matmul), softmax e normalização de camadas, essenciais em modelos de Transformer (como LLaMA ou GPT-2), consomem a maior parte do tempo de computação dentro desses kernels. O desafio é que extrair o máximo desempenho exige uma compreensão simultânea de intensidade aritmética, coalescência de memória, pressão de registradores, tamanhos de blocos, sincronização em nível de warp e seleção de instruções do tensor core – um conjunto de habilidades que leva anos para ser desenvolvido. Um único kernel de matmul de alto desempenho pode envolver mais de 200 linhas de código CUDA ou Triton, com dezenas de parâmetros interdependentes. A escassez dessa expertise e a dificuldade de escalonar o processo de ajuste manual, à medida que as arquiteturas de modelos evoluem, são problemas que o AutoKernel busca resolver. O KernelBench, um conjunto de benchmarks, revelou que mesmo os melhores modelos gerados com uma única tentativa não superavam o desempenho base do PyTorch em 80% dos casos, lacuna que motivou a criação do AutoKernel.
Como AutoKernel revoluciona a otimização de GPU
O ciclo autônomo: Editar, testar, aplicar ou reverter
A principal inovação do AutoKernel reside em seu insight central: o fluxo de trabalho de um engenheiro especialista em kernels pode ser replicado por um loop simples de agente. Esse loop consiste em: escrever uma nova versão de kernel, testá-la (benchmark), manter as melhorias e descartar as regressões, repetindo o processo. O framework mecaniza este ciclo: um agente LLM modifica um único arquivo, `kernel.py`; um sistema de benchmark verifica a correção e mede a taxa de transferência; e o resultado determina se a mudança é mantida ou revertida. Cada experimento é mapeado para um commit Git, permitindo que a história completa seja navegada com ferramentas Git padrão e que os resultados sejam logados em um arquivo `results.tsv` simples, legível por humanos e pelo agente.
Cada iteração leva aproximadamente 90 segundos: 30 segundos para verificação de correção, 30 segundos para benchmarking de desempenho via `do_bench` do Triton, e 30 segundos para o raciocínio do agente e modificação de código. Com cerca de 40 experimentos por hora, uma execução de 10 horas durante a noite pode gerar 300 a 400 experimentos em múltiplos kernels. Este design foi diretamente inspirado no projeto autoresearch de Andrej Karpathy, que demonstrou que um agente IA executando um loop de manter/reverter em código de treinamento de LLM poderia descobrir 20 otimizações em 700 experimentos.
O agente segue um documento de instruções de 909 linhas, o `program.md`, que codifica o conhecimento especializado em um manual de otimização de seis níveis. Estes níveis progridem desde o ajuste do tamanho do bloco (varrendo dimensões de tile por potências de 2, ajustando `num_warps` e `num_stages`) até padrões de acesso à memória, otimizações de computação, técnicas avançadas e estratégias específicas de arquitetura de GPU, como TMA em Hopper e `cp.async` em Ampere. O documento é abrangente para garantir que o agente possa operar por mais de 10 horas sem interrupções. Para mais detalhes técnicos, o artigo original pode ser acessado em arxiv.org/pdf/2603.21331.
Profiling inteligente: Otimizando onde há mais ganho
Diferente de trabalhos anteriores que tratam problemas de kernel isoladamente, o AutoKernel parte de um modelo PyTorch completo. Ele utiliza o `torch.profiler` com gravação de forma para capturar o tempo de GPU por kernel, e então classifica os alvos de otimização usando a Lei de Amdahl. Este princípio matemático estabelece que o ganho de velocidade total é limitado pela proporção que um componente representa no tempo total de execução. Um aumento de 1,5x na velocidade de um kernel que consome 60% do tempo total de execução resulta em um ganho de 1,25x de ponta a ponta. O mesmo ganho em um kernel que consome apenas 5% do tempo resultaria em um ganho de apenas 1,03x. Isso garante que os esforços de otimização sejam direcionados para onde terão o maior impacto. O profiler detecta hardware de GPU de um banco de dados de especificações conhecidas, abrangendo GPUs NVIDIA (H100, A100, L40S, L4, A10, RTX 4090/4080/3090/3080) e AMD (MI300X, MI325X, MI350X, MI355X).
Impacto e perspectivas para o futuro da IA
O AutoKernel tem o potencial de impactar significativamente o campo de Machine Learning. Para desenvolvedores e engenheiros, ele elimina uma barreira técnica considerável, permitindo que se concentrem na inovação e na arquitetura dos modelos, em vez de despenderem tempo em otimizações de baixo nível. Para empresas, isso significa ciclos de desenvolvimento mais rápidos, custos operacionais potencialmente menores devido à eficiência de hardware e a capacidade de implantar modelos mais complexos e rápidos em produção. A democratização da otimização de GPU pode acelerar a pesquisa e o desenvolvimento em áreas como visão computacional, processamento de linguagem natural e robótica, onde o desempenho é um fator crítico. A adoção de ferramentas como o AutoKernel pode se tornar um padrão na indústria, impulsionando ainda mais a inovação em IA.
Em suma, o AutoKernel da RightNow AI representa um avanço importante na automação da engenharia de Machine Learning. Ao aplicar agentes autônomos à otimização de GPU kernels, a plataforma não apenas economiza tempo e recursos, mas também abre portas para que mais inovações em IA sejam implementadas com alta performance, tornando a construção de modelos eficientes mais acessível do que nunca.
Gostou da notícia? Inscreva-se na nossa newsletter para receber as principais novidades sobre inteligência artificial diretamente no seu e-mail.
Fonte: https://www.marktechpost.com
