Auto-research com IA: como um kernel ficou 232 vezes mais rápido

O que é auto-research com IA?

Auto-research é um ciclo de pesquisa experimental em que um agente propõe mudanças, executa testes, observa métricas e registra o resultado para escolher a próxima tentativa. A ideia não é pedir uma solução única para a IA, mas criar um ambiente em que cada hipótese possa ser medida.

O tema ganhou atenção entre desenvolvedores de alto desempenho por combinar um objetivo quantitativo com feedback rápido. Em vez de avaliar apenas se o código parece correto, o processo acompanha tempo de execução, uso de recursos e testes de validade a cada rodada.

O relato que inspirou este artigo veio de uma competição do GPU MODE em parceria com a Core Automation. O participante descreveu uma implementação de QR compacta em GPU, terminou em 12º lugar entre 183 participantes e relatou uma melhora de 232 vezes sobre o baseline da tarefa.

Como funciona

O ponto de partida é um baseline confiável. No caso do relato, a tarefa recebia matrizes quadradas FP32 em lote e precisava devolver a representação compacta da fatoração QR. O resultado tinha de continuar correto para o verificador, não apenas rápido.

A fatoração QR separa uma matriz A em Q e R. A matriz R é triangular superior, enquanto Q precisa manter as propriedades de ortogonalidade. A representação compacta guarda R na parte superior de H, vetores de Householder abaixo da diagonal e os coeficientes tau em um vetor separado.

O ganho veio de reduzir o trabalho serial. O algoritmo bloqueado concentra as etapas dependentes em painéis pequenos e transforma o restante em atualizações parecidas com multiplicações de matrizes. Essa forma aproveita melhor as unidades de matriz da GPU do que aplicar um vetor por vez ao bloco inteiro.

Principais recursos do método

O primeiro recurso é o harness, ou seja, o conjunto de scripts que prepara entradas, executa candidatos, verifica a saída e salva as métricas. Ele transforma uma conversa com a IA em um experimento reproduzível.

O segundo é o perfil de execução. Medições de tempo dizem que algo ficou mais rápido, mas um profiler ajuda a descobrir se o gargalo está em lançamentos de kernel, cópias, sincronização, operações de memória ou baixa ocupação da GPU.

O terceiro é a diversidade de hipóteses. O autor do relato manteve várias famílias de candidatos em vez de comparar toda tentativa apenas com um único campeão. Isso reduz o risco de abandonar uma mudança estrutural antes que ela tenha recebido os ajustes necessários.

• Verificador: confirma correção numérica e formato de saída.
• Benchmark: mede o tempo em diferentes tamanhos e distribuições.
• Profiling: mostra onde o hardware está parado ou trabalhando.
• Registro: evita repetir ideias que já falharam.

Como começar: acesso e preparação

Comece com um problema pequeno e mensurável. Pode ser uma operação de álgebra linear, uma etapa de pré-processamento ou um kernel usado em inferência. O importante é ter entradas fixas, uma saída verificável e uma métrica que possa ser repetida.

Depois, implemente uma versão de referência usando uma biblioteca conhecida. No caso de QR, a função torch.geqrf pode servir como referência de formato e correção. A versão de referência não precisa ser a mais rápida, mas deve ser simples o bastante para validar as candidatas.

Em seguida, documente o ambiente: modelo da GPU, versão do driver, versão do framework, tamanho das entradas e número de repetições. Para escrever kernels próprios, consulte a documentação oficial do Triton e mantenha o caminho CUDA disponível para comparação.

1. Defina a entrada, a saída e o critério de correção.
2. Meça o baseline depois de aquecer a GPU.
3. Crie um script que salve tempo, erro e nome do candidato.
4. Altere uma hipótese por vez antes de combinar mudanças.
5. Promova somente candidatos que sejam rápidos e corretos.

Exemplo prático

O exemplo abaixo mede uma referência de QR em uma GPU. Ele é um esqueleto de benchmark, não uma reprodução completa da competição. Em um experimento sério, repita a medição, valide Q e R e cubra vários tamanhos de matriz.

import torch

A = torch.randn(32, 512, 512, device='cuda', dtype=torch.float32)
for _ in range(10):
    H, tau = torch.geqrf(A)

torch.cuda.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
H, tau = torch.geqrf(A)
end.record()
torch.cuda.synchronize()
print(f'tempo_ms={start.elapsed_time(end):.3f}')

A primeira melhoria não precisa trocar tudo por Triton. Você pode separar os tamanhos de entrada, evitar caminhos genéricos quando a forma for conhecida e reduzir sincronizações desnecessárias. Cada mudança deve deixar um registro do antes e do depois.

Quando a referência estiver estável, introduza uma implementação candidata. A função do agente é propor e editar experimentos, mas o harness continua responsável por medir e rejeitar resultados incorretos. Esse detalhe é o que transforma tentativa e erro em engenharia.

Comparação com alternativas

A otimização manual tradicional continua sendo a melhor escolha quando o problema é bem conhecido e a equipe já domina o hardware. Ela costuma exigir menos infraestrutura inicial, mas depende bastante da experiência de quem escreve o kernel.

Um autotuner convencional explora parâmetros de uma família de implementações. Ele é eficiente quando o espaço de busca está bem definido. O auto-research amplia a busca para mudanças de algoritmo, layout, fusão e estratégia de execução, mas exige um verificador mais cuidadoso.

Uma solicitação única para um modelo de linguagem é útil para criar o primeiro protótipo. Ela é insuficiente para otimizações difíceis, porque não mede dezenas de variantes nem aprende com o histórico. O loop contínuo é a diferença central.

• Use otimização manual: quando o gargalo já está comprovado e a solução é conhecida.
• Use autotuning: quando os parâmetros são o principal espaço de busca.
• Use auto-research: quando há muitas hipóteses e um teste barato e confiável.
• Use uma IA em uma rodada: para prototipar ou explicar o código antes da medição.

Pontos positivos e limitações

A principal vantagem é a velocidade do ciclo experimental. A IA pode preparar variantes, analisar logs e sugerir o próximo teste enquanto o desenvolvedor se concentra no contrato e nas decisões de arquitetura.

Outra vantagem é a memória operacional do projeto. Um bom log mostra o que foi tentado, em qual formato, com qual resultado e por que uma alternativa foi descartada. Isso reduz trabalho duplicado entre sessões e pessoas.

A limitação é que a qualidade do resultado depende do ambiente de avaliação. Um benchmark ruim incentiva otimizações artificiais, uma validação fraca deixa passar erros numéricos e uma medição instável cria falsos ganhos.

Casos de uso reais

O método faz sentido para equipes que mantêm kernels de inferência, treinamento ou processamento numérico. Nessas áreas, poucos microssegundos repetidos milhões de vezes podem representar uma diferença operacional importante.

Também pode ajudar quem mantém extensões de frameworks. Um desenvolvedor pode testar layouts, fusões e tamanhos de bloco em Triton ou CUDA, sempre preservando uma implementação de referência para validar a extensão.

Para estudantes e pesquisadores, o loop funciona como um laboratório. A pessoa aprende uma ideia de álgebra linear, transforma a ideia em código, mede o efeito e volta à teoria quando o resultado não explica o comportamento observado.

• Engenheiros de desempenho otimizando inferência em lote.
• Pesquisadores testando kernels para álgebra linear.
• Times de plataforma reduzindo custo de GPU em serviços recorrentes.
• Desenvolvedores estudando Triton com um objetivo observável.

Dicas e boas práticas

A primeira prática é reduzir o problema antes de pedir mais esforço ao agente. Um objetivo como diminuir o tempo do caso de 512 por 10% é mais útil do que um pedido genérico para deixar o kernel melhor.

Outra boa prática é manter a intervenção humana. O desenvolvedor deve revisar hipóteses, conferir se a medição é justa e decidir quando um resultado merece ser promovido. Autonomia no loop não significa abrir mão de critérios.

Vale a pena?

Vale a pena quando o problema tem um verificador automático, um benchmark barato e espaço real para testar alternativas. A técnica é especialmente interessante em kernels GPU, onde pequenas decisões de layout e lançamento podem mudar muito o resultado.

Não vale começar por ela quando o código ainda não tem testes, quando o ambiente muda a cada execução ou quando ninguém consegue explicar o contrato da saída. Nesses casos, primeiro estabilize a referência e a medição.

O próximo passo é escolher uma função que você já executa com frequência, criar um baseline reproduzível e registrar uma única hipótese. Se o primeiro ciclo funcionar, aumente a diversidade de candidatos sem perder a validação.