A NVIDIA detalha um fluxo com agentes de IA para traduzir kernels de GPU escritos em cuTile Python e Triton-TileIR para cuTile Rust, uma interface em Rust para escrever kernels de forma segura e idiomática. No exemplo do TileGym, o sistema foi usado para portar os 24 operadores públicos da biblioteca, somando cerca de 40 kernels, entre operações elementares, flash-attention decode, Multi-head Latent Attention (MLA) e modelos mixture-of-experts (MoE).
O texto explica que o trabalho não é feito por um único agente: há um analisador, um escritor do kernel, uma etapa para o código de host e ligações C, e outra para medir desempenho. A principal diferença que precisa ser resolvida é que o cuTile Python especializa cada kernel no momento da chamada, enquanto o Rust exige que as especializações apareçam na assinatura. Por isso, o analisador extrai constantes, tipos, grades de execução e a implementação de referência antes de qualquer código em Rust ser gerado.
A verificação também é dupla. Primeiro, o kernel em Rust passa por um teste funcional interno. Depois, a IR do Tile gerada no Rust é comparada com a da referência, para detectar mudanças estruturais que poderiam escapar dos testes. A própria NVIDIA diz que essa comparação ajuda a separar erro de tradução de falha do compilador.
Nos resultados citados no post, os operadores portados chegaram a 99,5% do desempenho da versão de referência, em média. O material também informa que o agente e os arquivos de apoio estão no repositório NVIDIA/TileGym no GitHub.