Auto-pesquisa com Codex acelera kernel de GPU em 232x

Resultado e problema

Em um concurso promovido pela GPU Mode em colaboração com a Core Automation, o participante implementou uma fatoração QR compacta de Householder para matrizes CUDA quadradas em FP32. A tarefa exigia retornar o formato (H, tau) compatível com torch.geqrf: a parte superior de H armazenava R, a parte inferior guardava os vetores de Householder e tau continha os coeficientes dos refletores. O verificador reconstruía Q e conferia as relações A ≈ QR, QᵀQ ≈ I e QᵀA ≈ R.

O resultado foi o 12º lugar entre 183 participantes. A solução final chegou a 1.805 microssegundos de média geométrica, contra cerca de 419 mil microssegundos da linha de base baseada em torch.geqrf/cuSolver, o que corresponde a um ganho de aproximadamente 232 vezes. O concurso avaliava vários tamanhos, de n = 32 a n = 4096, além de diferentes tamanhos de lote e condições numéricas.

Estratégia algorítmica e otimização

A fatoração QR de Householder tem dependências seriais: cada refletor precisa ser calculado a partir da matriz modificada pelo refletor anterior. Para reduzir esse trabalho sequencial, a solução adotou o algoritmo bloqueado de Householder. Nele, os refletores de um painel estreito são combinados em uma representação WY, transformando a atualização da grande matriz restante em operações semelhantes a multiplicações de matrizes, mais adequadas aos tensor cores da GPU.

O autor relata desafios com precisão reduzida em entradas mal condicionadas, com a grande variedade de formatos e lotes e com a ocupação dos tensor cores. A evolução da implementação incluiu painéis em Triton, atualizações WY agrupadas, especializações para formatos fixos, fusão da montagem de V e T, reexecução por CUDA Graphs e rotas específicas, como uma abordagem baseada em Cholesky para n = 4096. Segundo o relato, o gargalo raramente era memória ou capacidade computacional bruta: custos de lançamento e de processamento dos painéis dominavam o tempo.

Auto-pesquisa com agentes

Durante 14 dias, foram feitas mais de 1.500 submissões. O autor usou o Codex com ChatGPT Pro, Claude Pro para discutir conceitos e ideias e Modal para perfilamento. O ambiente foi organizado com arquivos de instruções, registros de tentativas, perfis e variantes de código. O Codex recebeu metas quantitativas por meio de /goal, enquanto comandos como /btw e /side permitiam acompanhar o trabalho sem interromper o ciclo.

O método dependia de um ciclo estreito entre implementação, submissão, verificação, perfilamento e registro dos resultados. No início, o agente conseguiu avançar rapidamente ao implementar a rota bloqueada para n = 512, uma das dimensões mais importantes na média geométrica. Depois dos 3.000 microssegundos, o autor passou a intervir mais, fornecendo dados de perfilamento e orientando a busca por fusões, reduções combinadas, sinais estáticos para o compilador e alternativas algorítmicas.

Limites e conclusões

O autor afirma que conhecimento de domínio melhora tanto o desenho do ambiente de experimentação quanto a orientação humana do agente, embora participantes sem experiência pudessem obter ganhos respeitáveis. Para escapar de máximos locais, ele recomenda manter uma “viga” de três a cinco candidatos, preservar ideias promissoras mesmo quando inicialmente mais lentas, usar modelos conselheiros e incentivar experimentos de maior risco.

Entre as oportunidades não exploradas estavam detectar distribuições especiais dos dados, substituir mais rotinas de biblioteca por implementações próprias, manter a matriz residual em FP16, criar um perfilador mais robusto e usar instruções tcgen05 dos tensor cores Blackwell. A conclusão favorece ambientes de auto-pesquisa específicos para cada problema, em vez de um único sistema geral de autoaperfeiçoamento.