Agente com perguntas supera agente com ferramentas: 113% contra 50% da velocidade do FlashAttention-3

Kernels de GPU rápidos para atenção (attention) são a base da velocidade de treinamento e execução de grandes modelos, mas escrevê-los à mão leva tempo: cada geração de placas de vídeo exige uma nova reescrita, e equipes especializadas gastam meses nisso. Pesquisadores da Carnegie Mellon University testaram se esse trabalho pode ser delegado a um agente de código e chegaram a uma conclusão pouco óbvia. O que ajuda um agente travado não são novas ferramentas nem mais liberdade, mas uma pergunta curta no momento certo.

Na mesma tarefa, o agente com o conjunto completo de ferramentas parou em 50% da velocidade do FlashAttention-3 (o kernel de atenção de referência para as placas Hopper) após 133 passos. O agente que recebeu perguntas chegou a 113% nas mesmas 26 configurações. Para o desenvolvedor, o recado é este: se o agente empaca, provavelmente falta a ele uma hipótese sobre onde procurar o problema, não mais uma porção de dados.

Por que as ferramentas não resolvem

Os agentes atuais de otimização de kernels ficam nos dois extremos de uma escala. Os agentes roteirizados definem cada passo de forma rígida: uma política treinada, um menu de otimizações, papéis fixos. Enquanto a receita serve, eles são econômicos; quando não serve, travam. Os agentes livres mutam por conta própria um conjunto de kernels, mas pagam por essa liberdade com milhares de execuções, edições repetidas e candidatos incorretos. Em um trabalho anterior, um agente desse tipo testou mais de 500 direções e lançou 40 versões em sete dias, e precisou de um supervisor para sair de ciclos de edições malsucedidas. Cerca de 30% das linhas que agentes assim adicionam coincidem byte a byte com linhas removidas antes.

Diante de uma meta difícil, as duas escolas respondem do mesmo jeito: dar mais ao agente. Saída de profiler já interpretada, análise de disassembly, busca na documentação, um repositório de «lições aprendidas». Os autores consideram que informação o agente já tem de sobra. Os modelos se agarram com avidez à primeira solução, param cedo demais, andam em círculos entre variantes quase idênticas e declaram o trabalho concluído quando a verificação mostra o contrário.

Como a pergunta muda a tarefa

O comando «faça mais rápido» é aberto, e o modelo responde a ele de forma gulosa, mexendo no que está por perto. Já a pergunta «você está ganhando em ciclos de clock do kernel ou só em tempo de relógio (wall-clock)?» é delimitada. Para respondê-la, o agente precisa analisar e reescrever uma parte do kernel em que ele mesmo não mexeria. Uma boa pergunta não traz a resposta pronta: ela aponta uma hipótese que o agente não tinha considerado, e essa hipótese pode ser verificada por meios comuns, com uma medição.

O único código específico do método é a estrutura de verificação (harness). Um comando compila o candidato, compara suas saídas com as do kernel original e com a referência em fp32, mede a velocidade em cada cenário em relação ao kernel original e calcula o speedup médio geométrico. A meta é definida como uma condição de parada verificável, e não como um pedido ao modelo.

As perguntas surgem quando o progresso para, quando o agente declara «este é o limite» ou «está pronto» e depois de um ganho suspeitamente grande. No início, quem faz as perguntas é uma pessoa, e as que se repetem vão para um banco que o próprio agente passa a consultar. O método também foi incorporado a uma busca evolutiva, em que as instruções de mutação são formuladas como perguntas.

O que foi obtido

As campanhas completas geraram kernels de passagem direta (forward) de atenção para Hopper e Blackwell: um único arquivo CUDA, sem CUTLASS, CuTe, cuBLAS e cuDNN. Eles estão corretos em todos os 182 cenários de cada arquitetura (no caso da Blackwell, em todos os que podem ser executados). Na média geométrica, são 1,16 vez mais rápidos que o FlashAttention-3, ou seja, 16% mais rápidos, e 1,46 vez mais rápidos que o FlashAttention-4, ou 46%. Em formatos de requisição de produção, o ganho sobre o FlashAttention-4 foi de 1,10 vez. Os autores chamam esses kernels de os mais rápidos de código aberto da categoria nos cenários testados.

A estatística das intervenções chama atenção: 162 mensagens humanas e 11,5 mil chamadas de ferramentas em 16 sessões. As perguntas não fizeram o agente usar menos ferramentas; ele apenas passou a usá-las de outro modo. Todo ganho grande veio depois de uma pergunta que levou o agente a outra camada do sistema ou que derrubou o limite que ele havia declarado, e nenhum veio depois de uma mensagem que apenas fornecia informação.

As ressalvas são evidentes. Os autores montaram por conta própria o agente autônomo de comparação, seguindo o modelo de trabalhos anteriores. O experimento é um só: ele conta com uma referência forte para medir o resultado. Parte do trabalho de direcionamento é feita por uma pessoa. Os kernels, a estrutura de verificação, o banco de perguntas e todos os traces foram publicados pelos pesquisadores.

Conclusão prática: antes de conectar mais uma ferramenta ou uma camada de busca, defina uma meta verificável que cubra todos os cenários e prepare perguntas para os momentos de estagnação, de limite declarado e de vitória grande demais.