Otimização de núcleos de GPU geralmente exige um profundo entendimento da arquitetura de hardware e de instruções de baixo nível. Modelos de linguagem ainda lidam com isso de forma instável: eles podem gerar código plausível, mas nem sempre ele é correto e rápido. PTXBench é um novo benchmark que tenta trazer mensurabilidade e clareza para esse processo.
O que é o PTXBench e por que ele é necessário
PTXBench é um ambiente de avaliação para grandes modelos de linguagem que tentam usar o PTX (Parallel Thread Execution), específico da arquitetura, para otimizar núcleos de GPU. PTX é uma representação intermediária de instruções para GPUs da NVIDIA, e é nesse nível que se pode controlar com precisão registradores, memória e o escalonador. Benchmarks comuns verificam apenas o código final, mas o PTXBench olha mais fundo: se o modelo realmente aplica as instruções de baixo nível alvo, ou se apenas gera código "semelhante a C".
Os autores focaram em duas cargas de trabalho clássicas: multiplicação de matrizes (GEMM) e o mecanismo de atenção (attention), que é crítico para transformadores. Os testes foram executados em aceleradores atuais — H100 e B200. Essa escolha permite verificar não apenas as habilidades básicas do modelo, mas também sua capacidade de se adaptar a arquiteturas modernas com suas características específicas.
O que o benchmark mede
O PTXBench avalia os modelos com base em três parâmetros-chave:
- Correção funcional — o resultado da execução do núcleo gerado deve coincidir com o de referência.
- Execução de instruções alvo — verifica se as instruções PTX que o modelo deveria usar aparecem em tempo de execução.
- Aceleração — o quanto o núcleo obtido é mais rápido que as bibliotecas frontier, ou seja, as melhores implementações prontas.
Essa abordagem separa três problemas distintos. O modelo pode escrever código correto, mas não usar as instruções exigidas. Ou usá-las, mas obter desempenho pior que o de uma biblioteca padrão. O PTXBench permite ver em qual estágio exato ocorre a falha.
Resultados: a otimização ainda é irregular para os modelos
A avaliação mostrou que a capacidade dos LLMs de trabalhar com PTX depende muito da tarefa. Em casos simples, os modelos apresentam bons resultados, mas a proporção de soluções bem-sucedidas cai drasticamente em tarefas backward complexas para attention. Isso é esperado: o passe reverso exige um trabalho mais cuidadoso com gradientes e memória.
Outra conclusão importante: a execução de instruções alvo por si só não garante velocidade. O modelo pode "acertar" as instruções PTX necessárias, mas escolher uma estratégia ruim de paralelização ou de posicionamento de sincronizações. No final, nenhum dos modelos testados conseguiu competir de forma estável com as bibliotecas frontier em todo o conjunto de tarefas.
Como o modelo foi ajustado e o que funcionou
Para entender se essas habilidades podem ser melhoradas, os autores realizaram o ajuste fino do modelo Qwen3.6-27B usando supervised fine-tuning. Além disso, usaram uma técnica que pode ser chamada de aprendizado com condição de correção: o modelo aprende não apenas a gerar código, mas a corrigi-lo, recebendo feedback sobre etapas incorretas.
Essa abordagem realmente melhorou várias tarefas, mas não garantiu um progresso uniforme. A generalização permaneceu fraca: o modelo podia se sair bem com alguns tipos de núcleos e falhar em outros. A análise mostrou que não é apenas o tamanho do conjunto de dados de treinamento que importa. Mais importante são sua cobertura, o equilíbrio entre diferentes cenários e a qualidade do "professor" que gera exemplos para cadeias de raciocínio. Se os dados têm um viés para casos simples, o modelo não aprenderá a complexidade.
Conclusão principal
O PTXBench oferece à comunidade um ambiente de teste verificável no qual se pode acompanhar o progresso dos LLMs na otimização de núcleos de GPU. Ele não apenas fornece uma avaliação de "rápido/lento", mas mostra em qual nível os planos do modelo se quebram: na sintaxe, na aplicação de instruções específicas ou na escolha da estratégia. Este é um passo importante para tornar a otimização automática de núcleos previsível e prática.
