OptiML: An End-to-End Framework for Program Synthesis and CUDA Kernel Optimization

1. O Problema

A geração de kernels CUDA de alto desempenho permanece um gargalo central para a computação científica e sistemas de aprendizado de máquina. Embora os GPUs ofereçam alto throughput, atingir desempenho competitivo exige expertise profunda em hierarquias de memória, configuração de execução, sincronização e comportamento do compilador.

Os desafios principais identificados são:

• Gap entre Correção e Desempenho: Modelos de Linguagem Grande (LLMs) conseguem frequentemente sintetizar código CUDA sintaticamente válido e funcionalmente correto, mas os kernels resultantes muitas vezes falham em compilar, violam restrições numéricas ou deixam desempenho significativo inexplorado devido à falta de transformações de baixo nível (ex: tiling, vetorização).
• Espaço de Otimização Combinatório: O espaço de otimização para kernels de GPU é não suave e altamente interativo. Pequenas decisões de código (padrões de acesso, estratégia de reutilização) podem alterar drasticamente se um kernel é limitado por computação ou memória.
• Feedback Caro e Ruidoso: A otimização baseada apenas em medições de tempo final (end-to-end timing) é ineficiente e frágil, pois o feedback é caro (requer compilação e perfilamento) e ruidoso, dificultando a identificação da causa raiz da ineficiência.
• Separação entre Geração e Otimização: Trabalhos anteriores tratam a geração de código e a otimização como problemas separados, ou traduzem programas de alto nível para CUDA, ou ajustam kernels existentes isoladamente, sem um fluxo unificado.

2. Metodologia: O Framework OptiML

O OptiML é um framework end-to-end que unifica a síntese de programas e a otimização de kernels, mapeando tanto intenções em linguagem natural quanto código CUDA existente para kernels otimizados. O sistema é composto por dois estágios desacoplados:

A. OptiML-G (Geração de Proposta)

Atua quando a entrada é uma descrição em linguagem natural.

• Arquitetura: Baseado no framework Mixture-of-Thoughts (MoT).
• Mecanismo: Mantém um pool de especialistas pré-treinados (ex: Qwen2.5-Coder, HPC-Coder-V2, StarCoder2). Um roteador leve calcula scores de relevância e seleciona os top-K especialistas.
• Colaboração Latente: Diferente de métodos de ensemble tradicionais, o OptiML-G permite colaboração granular entre os especialistas selecionados em um espaço latente compartilhado. As representações intermediárias (estados ocultos) dos especialistas são projetadas e fundidas via mecanismos de cross-attention.
• Objetivo: Gerar um programa executável inicial estruturalmente forte, que já exponha oportunidades significativas para transformação, servindo como um ponto de partida robusto para a otimização.

B. OptiML-X (Otimização Baseada em Busca)

Refina kernels (seja os gerados pelo OptiML-G ou fornecidos pelo usuário) através de um processo de busca guiada.

• Motor de Busca: Utiliza Monte Carlo Tree Search (MCTS) para explorar trajetórias de otimização de múltiplos passos.
• LLM-as-a-Judge: Um LLM atua como juiz e diagnosticador. Dado o estado atual do kernel e métricas de perfilamento, o LLM:
  1. Propõe hipóteses de gargalos (ex: limitado por memória, por registro, por ocupação).
  2. Gera patches concretos (edições de código) para abordar esses gargalos.
• Feedback de Hardware (Profiling): O sistema compila, testa e perfila cada candidato usando o Nsight Compute.
• Modelo de Recompensa Composta: A busca é guiada por uma recompensa que combina:
  • Tempo de Execução: Melhoria direta no latency.
  • Métricas Proxy (Hardware-Aware): Utiliza métricas do Nsight (SOL - Stall on Load, bytes de DRAM, setores L1, instruções executadas) para identificar gargalos específicos.
  • Guardrails: Penalizações severas para candidatos que falham em compilação, testes de correção (L0/L1) ou que degradam métricas críticas (ex: aumento de register pressure ou shared memory).
  • Julgamento do LLM: Avalia se a edição proposta faz sentido lógico em relação ao gargalo diagnosticado.

3. Principais Contribuições

1. Formulação Unificada: Define o problema de otimização de kernels CUDA condicionada a intenção, tratando síntese e otimização como um pipeline contínuo.
2. Arquitetura OptiML: Apresenta a combinação de um gerador Mixture-of-Thoughts (para síntese inicial de alta qualidade) com um otimizador baseado em MCTS guiado por perfis de hardware e julgamento de LLM.
3. Otimização Interpretável: Diferente de "caixas pretas", o sistema explica por que um kernel melhorou (ex: redução de tráfego de memória, aumento de utilização de SM), baseando-se em evidências de perfilamento.
4. Validação Empírica: Demonstra que o framework supera consistentemente baselines de LLMs fortes e variantes de otimização isolada em uma suíte diversificada de kernels.

4. Resultados

Os experimentos foram realizados em uma GPU NVIDIA A100 80GB com uma suíte de kernels CUDA variados (Multiplicação de Matrizes, Max Pooling 3D, Atenção Multi-Head, ReLU, Cross-Entropy, etc.).

• Desempenho Geral: O OptiML (G + X) alcançou speedups consistentes em relação aos melhores LLMs de base. Por exemplo:
  • Multiplicação de Matrizes: Redução de tempo de 7.20s (GPT-5.1) para 4.40s (OptiML), um speedup de 1.64x.
  • ReLU: Redução de 7.50s para 4.49s (1.67x).
  • Max Pooling 3D: Muitos LLMs falharam em compilar (CF), enquanto o OptiML-X conseguiu gerar e otimizar kernels válidos, alcançando 1.08x de speedup sobre a melhor tentativa compilável.
• Qualidade de Geração (Pass@k): O OptiML-G superou todos os modelos de base no benchmark ParEval, alcançando 12.17% Pass@1 (vs. 10.46% do GPT-5.1), indicando maior capacidade de gerar código compilável e correto na primeira tentativa.
• Análise de Gargalos:
  • Para kernels limitados por memória (ex: Max Pooling 3D), o sistema reduziu significativamente o tráfego de DRAM e setores L1.
  • Para kernels limitados por instrução (ex: ReLU), reduziu o número de instruções executadas e melhorou a vetorização.
  • O sistema evitou "trocas" de gargalos (ex: melhorar memória mas piorar ocupação), focando em melhorias líquidas.
• Eficiência da Busca: O uso do OptiML-G como inicialização reduziu o número de iterações necessárias para encontrar a melhor solução (round-to-best) em comparação com a otimização partindo de kernels gerados apenas por LLMs padrão.

5. Significado e Impacto

O trabalho OptiML representa um avanço significativo na automação de programação para aceleradores:

• Superação da "Caixa Preta": Ao integrar feedback de hardware (Nsight Compute) diretamente no loop de busca e decisão do LLM, o sistema torna a otimização interpretável e direcionada a gargalos reais, não apenas a métricas de tempo cegas.
• Robustez: A abordagem de "gerar-então-otimizar" mitiga a fragilidade de tentar otimizar código que não compila ou é estruturalmente deficiente. O OptiML-G fornece uma fundação sólida, e o OptiML-X refina com precisão.
• Escalabilidade: O framework demonstra que é possível automatizar a criação de kernels de alto desempenho que rivalizam com implementações manuais de especialistas, reduzindo a barreira de entrada para o desenvolvimento de software para GPUs.
• Futuro: O trabalho abre caminho para sistemas que não apenas escrevem código, mas "entendem" a arquitetura de hardware subjacente para tomar decisões de otimização informadas, potencialmente estendendo-se a outros backends (HIP, Triton) e bibliotecas de GPU.

Em resumo, o OptiML estabelece um novo paradigma onde a síntese de código e a otimização de desempenho são tratadas como um processo contínuo e guiado por dados, superando as limitações das abordagens atuais de geração única ou ajuste pós-hoc isolado.