PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

Dentro do computador moderno, a unidade de processamento gráfico, ou GPU, atua como um motor poderoso projetado para lidar com quantidades massivas de dados simultaneamente. Para extrair o máximo deste motor, desenvolvedores de software escrevem instruções especializadas que dizem ao hardware exatamente como mover dados e realizar cálculos. Essas instruções são escritas em uma linguagem de baixo nível chamada PTX, que dá ao programador controle direto sobre o funcionamento interno da máquina. Embora existam ferramentas de nível superior para facilitar esse processo, elas às vezes perdem as características únicas das gerações mais recentes de hardware. À medida que essas máquinas evoluem rapidamente, manter o software perfeitamente ajustado às suas capacidades específicas torna-se um desafio difícil. Surge a questão: a inteligência artificial, especificamente os modelos de linguagem de grande escala, consegue aprender a escrever essas instruções precisas de baixo nível por conta própria, ou ainda depende de métodos antigos e genéricos que deixam desempenho na mesa?

Uma equipe de pesquisadores partiu para responder a isso criando um novo campo de testes chamado PTXBench. Eles queriam ver se esses modelos de IA poderiam não apenas escrever código que funciona, mas também código que utiliza ativamente os recursos específicos e avançados das GPUs mais recentes. Os pesquisadores focaram em duas das placas gráficas mais poderosas disponíveis, a H100 e a B200, e testaram os modelos em tarefas envolvendo operações matemáticas complexas usadas em inteligência artificial, como multiplicação de matrizes e mecanismos de atenção. O objetivo era medir três coisas: se o código rodava sem erros, se ele realmente acionava as instruções de hardware específicas que os pesquisadores solicitaram e se rodava mais rápido do que as melhores bibliotecas de software existentes.

Os resultados revelaram uma lacuna significativa entre o que os modelos podiam fazer e o que era necessário para um desempenho de alto nível. Embora os modelos de IA fossem frequentemente bem-sucedidos em escrever códigos que funcionavam corretamente para tarefas mais simples, eles tiveram dificuldades consideráveis com operações mais complexas, particularmente aquelas envolvendo a passagem de retropropagação (backward pass) de mecanismos de atenção, que são cruciais para o treinamento de sistemas de IA. Mesmo quando os modelos conseguiam escrever códigos que executavam as instruções de hardware solicitadas, esse código raramente igualava a velocidade das principais bibliotecas profissionais. Em muitos casos, o código gerado pela IA era mais lento, sugerindo que simplesmente fazer as instruções rodarem não era o suficiente para alcançar uma verdadeira otimização. O estudo descobriu que os modelos frequentemente recorriam a padrões de codificação genéricos em vez de aproveitar os recursos únicos e específicos da arquitetura dos chips mais novos, mesmo quando esses recursos eram explicitamente descritos para eles.

Para entender como melhorar esses modelos, os pesquisadores conduziram um experimento controlado onde ensinaram um modelo de IA específico usando um método chamado ajuste fino supervisionado (supervised fine-tuning). Eles mostraram ao modelo exemplos de seus próprios erros, juntamente com versões corrigidas do código e explicações do por que as correções funcionavam. Essa abordagem, que chamaram de Fixit, ajudou o modelo a melhorar em várias tarefas, particularmente na geração de código que é funcionalmente correto. No entanto, a melhoria não foi uniforme. O modelo ainda tinha dificuldade em generalizar suas novas habilidades para diferentes tamanhos de problemas ou variações das tarefas que não havia visto durante o treinamento. Os pesquisadores descobriram que a qualidade das correções e o equilíbrio dos dados de treinamento importavam tanto quanto a quantidade bruta de dados. Um modelo treinado em um conjunto diversificado de problemas teve um desempenho melhor do que um treinado em um conjunto maior, porém menos variado, indicando que o tipo de experiência de aprendizado era mais importante do que o volume.

O estudo também destacou a importância de fornecer o contexto certo à IA. Quando os pesquisadores deram aos modelos informações detalhadas sobre as capacidades específicas do hardware, os modelos foram muito mais propensos a usar as instruções solicitadas. Sem esse conhecimento específico, o modelo raramente utiliza o PTX solicitado. Isso sugere que, embora a IA possa aprender a programar nesse nível baixo, ela requer orientação precisa e exemplos de alta qualidade para fazê-lo de forma eficaz. Os pesquisadores concluíram que, embora os modelos de IA atuais mostrem promessa, eles ainda não estão prontos para substituir engenheiros humanos especialistas na otimização de código para o hardware mais recente. O caminho a seguir envolve melhores dados de treinamento, experiências de aprendizado mais equilibradas e uma compreensão mais profunda de como ensinar esses modelos a raciocinar sobre as restrições específicas das arquiteturas de computação em evolução.