Resumo Técnico: KernelBrain

Declaração do Problema

Automatizar a otimização de kernels de GPU continua sendo um desafio prático significativo. Embora os grandes modelos de linguagem (LLMs) possam gerar variantes de kernels, a implantação em produção enfrenta três obstáculos críticos:

1. Regressões de Corretude: O código gerado frequentemente viola restrições ou falha na compilação.
2. Medições Ruidosas: As classificações de desempenho em tempo de execução são instáveis devido ao ruído de medição, levando a decisões de otimização incorretas.
3. Estagnação de Busca: Os landscapes de otimização são "não-monotônicos", onde pequenas mudanças no código podem causar oscilações massivas de desempenho dependentes de hardware ou detriments temporários, fazendo com que os processos de busca estagnem ou derivem para programas de baixo desempenho.

Frameworks existentes muitas vezes carecem da capacidade de refatoração estrutural profunda, priorizando ganhos incrementais em vez de mudanças algorítmicas, enquanto abordagens multi-agente intensivas em recursos consomem orçamento excessivo em profiling e exploração sem garantir estabilidade.

Metodologia: KernelBrain

O KernelBrain é um agente de otimização prático e consciente de orçamento, projetado para manter um loop de otimização estável e eficiente sob ruído e interrupções. O sistema opera sob um princípio de alocação assimétrica de recursos, combinando mutação guiada por LLM com feedback em malha fechada de hardware (hardware-in-the-loop).

Arquitetura Central

O sistema é organizado em torno de dois planos paralelos conectados por um escalonador adaptativo:

1. Plano de Avaliação:

   • Verificação de Política ($V(p)$): Um portão de pré-benchmark que verifica a validade da compilação, assinaturas de função e corretude da saída contra uma implementação de referência. Candidatos que falham neste estágio são rejeitados imediatamente.
   • Benchmarking de Multi-Fidelidade: Os candidatos são avaliados através de uma escada de "degraus" (níveis de fidelidade). Os degraus inferiores utilizam recursos mínimos (menos repetições, formas de entrada estreitas) para uma triagem rápida. Os degraus superiores alocam mais orçamento (mais repetições, cobertura de formas mais ampla) e, opcionalmente, acionam o Profiling de Kernel de GPU (via NVIDIA Nsight Compute) para coletar métricas como ocupação de SM, throughput de memória e stalls de pipeline.
   • Pontuação: Uma pontuação $S(p, r)$ é calculada com base na latência (throughput) e bônus específicos do degrau, priorizando candidatos que demonstram estabilidade e velocidade.

2. Plano de Geração:

   • Busca Coarse-to-Fine Guiada por Especialistas: Em vez de mutação evolutiva cega, o sistema utiliza dicas de especialistas para cobrir grandes espaços de otimização (algoritmos, tiling, layouts de memória) nas rodadas iniciais.
   • Mutação Informada por Profiling: Para candidatos promovidos aos degraus superiores, o sistema utiliza dados do profiler para diagnosticar gargalos. O Construtor de Prompts constrói prompts estruturados para o LLM, direcionando as mutações para problemas específicos identificados (ex: coalescência de memória, pressão de registradores) em vez de reescritas não guiadas.
   • Parser/Reparo: Um módulo determinístico extrai o código do kernel das respostas do LLM e aplica correções para erros comuns (importações ausentes, assinaturas quebradas) para garantir que o loop permaneça robusto.

Loop de Otimização

O processo segue um loop de multi-fidelidade online (Algoritmo 1):

1. Retira um candidato e seu degrau atribuído.
2. Verifica a corretude; se falhar, registra e descarta.
3. Realiza o benchmark no nível de fidelidade atual.
4. Computa a pontuação e, se aplicável, realiza o profiling do kernel.
5. O escalonador decide se deve promover o candidato para um degrau superior (aumentando o orçamento e acionando a mutação) ou podar (prune) o candidato.
6. Candidatos promovidos passam por mutação de LLM baseada no diagnóstico de profiling e reentram na fila no próximo degrau.

Principais Contribuições

1. Alocação Assimétrica de Recursos: O sistema aloca recursos mínimos para escanear e podar rapidamente candidatos de baixo desempenho, concentrando a maior parte do orçamento de avaliação nos sobreviventes promissores através de um regime de multi-fidelidade. Isso inclui portão de política explícito e adjudicação de estabilidade.
2. Busca Coarse-to-Fine Guiada por Especialistas: A abordagem vai além da mutação cega de LLM ao empregar dicas de especialistas para cobrir amplamente os espaços de otimização no início, e então estreitar progressivamente a busca em direção a regiões de alto desempenho usando orientação de mutação rastreável e consciente de gargalos derivada de profiling de hardware.
3. Desempenho de Estado da Arte em Triton: O sistema demonstra resultados competitivos em cargas de trabalho baseadas em Triton, alcançando acelerações significativas sobre o PyTorch e agentes existentes, enquanto reduz o tempo de otimização.

Resultados Experimentais

O sistema foi avaliado em seis benchmarks importantes de DNN (GroupConv, MatVec, RMSNorm, Softmax, LinearReLUDiv, SDPAttn) em uma GPU NVIDIA A100.

• Ganhos de Desempenho: O KernelBrain alcançou acelerações variando de 0.88x a 6.72x sobre a referência PyTorch.
• Comparação com SOTA: Alcançou até 1.4x de aceleração sobre o KernelAgent (um sistema de otimização de kernel agêntico de estado da arte).
• Eficiência: A abordagem reduziu o tempo de otimização em até 48% (média geométrica) em comparação com os baselines.
• Benchmarks Específicos:
  • GroupConv: Alcançou 5.7x–6.7x de aceleração sobre o PyTorch, superando significativamente tanto o gerador base de LLM quanto o KernelAgent, com ganhos escalando conforme o tamanho do problema aumentava.
  • RMSNorm & Softmax: Demonstraram melhorias consistentes (2.5x–2.8x e 1.25x–1.33x respectivamente) sobre baselines de não-busca.
  • SDPAttn: Liderou em comprimentos de sequência maiores (1.17x–1.54x sobre o KernelAgent), embora os custos de lançamento (launch overheads) tenham mantido alguns resultados abaixo do PyTorch em sequências menores.

Observações sobre Padrões de Otimização:
Os experimentos revelaram que os ganhos mais significativos vieram de mudanças na decomposição paralela de trabalho (ex: mudar de redução serial para redução dividida cooperativa) e de redução do tráfego de memória global (ex: substituir implementações de múltiplas passagens por reduções fundidas ou online). O sistema identificou com sucesso essas mudanças estruturais em vez de depender apenas de edições locais menores.

Significância

O artigo afirma que o KernelBrain operacionaliza o fluxo de trabalho de um programador de GPU experiente: começando com uma exploração ampla e de baixo custo de diversas estratégias algorítmicas e de implementação, e depois concentrando progressivamente o orçamento nos candidatos mais promissores à medida que as evidências se acumulam. Ao integrar o profiling de hardware diretamente no loop de mutação e impor rigorosos portões de corretude, o sistema aborda a natureza "não-monotônica" da otimização de kernels, oferecendo um método estável e eficiente em termos de orçamento para evoluir kernels de alto desempenho sem o desperdício de recursos da busca exaustiva tradicional ou a instabilidade da geração de LLM não guiada.