Resumo Técnico: GPU Forecasters: Modelos de Linguagem como Substitutos Seletivos para Otimização de Tempo de Execução de Kernels

Declaração do Problema

Os kernels de GPU são as ferramentas fundamentais dos deep learning modernos, onde ganhos de eficiência marginais traduzem-se em reduções significativas nos custos de treinamento e inferência. A otimização desses kernels envolve tipicamente um ciclo de geração-avaliação, no qual Modelos de Linguagem de Grande Escala (LLMs) propõem kernels candidatos, que são então compilados e executados em hardware físico para medir o desempenho. No entanto, este processo de avaliação é um gargalo severo. Cada avaliação requer compilação (levando de 60 a 120 segundos para CUDA) e execução no dispositivo, tornando o throughput de avaliação ordens de magnitude mais lento do que a geração. À medida que os orçamentos de busca escalam, o custo de medições físicas repetidas limita a capacidade de explorar o vasto espaço de potenciais implementações de kernels.

Embora trabalhos anteriores tenham focado no uso de LLMs para gerar melhores kernels, este artigo aborda o problema complementar de reduzir o custo de avaliação. Os autores propõem o uso de LLMs não apenas como geradores, mas como substitutos seletivos (selective surrogates) para a GPU física. O objetivo é prever o desempenho relativo de um kernel candidato em relação a uma implementação de referência e identificar quando o modelo está incerto, de modo a adiar apenas os casos incertos para a GPU física cara.

Metodologia

Formulação da Tarefa: Previsão Seletiva

Os autores formulam o problema como a previsão do ganho de velocidade relativo ($S_i = T_{ref} / T_i$) de um kernel candidato em comparação com um kernel de referência, dada a especificação do hardware. Prever o tempo de execução absoluto em segundos é considerado intratável devido à complexidade de modelar o comportamento do compilador e detalhes microarquitetônicos. Em vez disso, o alvo é discretizado em oito bins ordinais de ganho de velocidade (variando de "Lentidão Severa" a "Ganho de Velocidade Extremo").

O substituto (um LLM) recebe o kernel de referência, o kernel candidato e as especificações do hardware como entrada. Ele gera uma cadeia de raciocínio e produz uma distribuição de probabilidade sobre os oito bins. A confiança declarada do modelo é definida como a probabilidade máxima atribuída a qualquer bin individual.

Estratégia de Predição Seletiva

O sistema opera sob um framework de predição seletiva. O substituto é usado para classificar candidatos. Se a confiança do substituto exceder um determinado limiar, a previsão é aceita e o candidato não é medido na GPU. Se a confiança for baixa (indicando incerteza), o candidato é transferido para a GPU física para medição de verdade fundamental (ground-truth). Isso permite que o processo de busca considere muito mais candidatos dentro de um orçamento fixo de medição de GPU.

Aprendizado por Reforço (RL) para Calibração

Para melhorar a precisão do substituto e, crucialmente, sua calibração (o alinhamento entre a confiança declarada e a precisão real), os autores aplicam Aprendizado por Reforço a um modelo de pesos abertos (GPT-OSS-20B). Eles definem recompensas baseadas em:

1. Corretude: Se o bin previsto corresponde ao bin medido.
2. Calibração Categórica (Brier Score): Penalizando a distribuição com base na distância do alvo one-hot real.
3. Calibração Sensível à Distância (CRPS): Penalizando a distribuição cumulativa, levando em conta a natureza ordinal dos bins de ganho de velocidade (ex: prever um ganho de 2× quando a verdade é 1.5× é menos severo do que prever 4×).

Os dados de treinamento são construídos como um subproduto da própria busca de kernels, utilizando candidatos bem-sucedidos de execuções de busca anteriores como exemplos de treinamento.

Principais Contribuições

1. LLMs como Substitutos de GPU: O artigo demonstra que LLMs podem alcançar uma precisão não trivial ao prever como uma GPU física responde a kernels candidatos, recuperando uma fração substancial do ganho de velocidade disponível ao classificar candidatos com base em previsões em vez de medições.
2. Framework de Predição Seletiva: Os autores tratam a previsão de GPU como um problema de predição seletiva, mostrando que substitutos podem reportar incerteza calibrada. Isso permite uma estratégia de busca que confia em previsões de alta confiança e reserva medições físicas para casos incertos.
3. Impacto do RL no Comportamento do Substituto: Através de experimentos, o artigo caracteriza como o treinamento de RL altera as propriedades do substituto. Especificamente, adicionar recompensas de calibração (Brier ou CRPS) melhora a confiabilidade das estimativas de confiança do modelo, enquanto o treinamento focado apenas em corretude pode degradar a relação entre confiança e precisão.
4. Melhoria da Busca End-to-End: O estudo mostra que integrar o substituto em um loop de busca de kernel (especificamente uma busca baseada em PUCT) permite que a busca explore significativamente mais candidatos por medição de GPU. Na maioria das tarefas, isso leva à descoberta de kernels mais rápidos do que um baseline que mede cada proposta sob o mesmo orçamento de GPU.
5. Liberação de Dataset: Os autores liberam uma coleção reutilizável de 12.388 kernels de GPU gerados por LLM com tempos de execução medidos, abrangendo 118 problemas, múltiplos tipos de GPU (A100, L40S, H100) e vários métodos de geração. Este dataset suporta pesquisas mais amplas em previsão de desempenho de GPU e busca de kernels.

Resultados

• Precisão da Previsão: LLMs prontos para uso (ex: GPT-OSS-20B) recuperam aproximadamente 83% do melhor ganho de velocidade disponível quando apenas o top 1% dos candidatos é medido na GPU. Este desempenho melhora para 94% quando se mede o top 50%.
• Calibração e RL: O treinamento de RL altera significativamente o comportamento do substituto.
  • Treinamento focado apenas em Corretude melhora a precisão exata do bin, mas degrada a correlação entre confiança e erro (o modelo torna-se excessivamente confiante em previsões erradas).
  • Treinamento de Calibração (Brier/CRPS) melhora a monotonicidade do erro de previsão versus confiança. Notavelmente, o treinamento calibrado pelo Brier melhora a métrica de "ganho de velocidade recuperado" em cerca de 2,1 pontos percentuais sobre o modelo base, sugerindo que distribuições melhor calibradas levam a sinais de classificação melhores para o algoritmo de busca.
• Eficiência de Busca: Na busca de kernel end-to-end, a busca habilitada pelo substituto encontra kernels mais rápidos que o baseline em quatro de seis tarefas, muitas vezes atingindo o melhor ganho de velocidade do baseline com substancialmente menos medições de GPU. Nas tarefas onde o baseline vence, a busca parece saturar cedo, indicando que o substituto é mais benéfico quando o espaço de busca é grande e inexplorado.
• Momentos de Descoberta: O substituto pode identificar "momentos de descoberta" (mutações onde um kernel filho é significativamente mais rápido que seu pai) melhor do que o acaso, embora a precisão continue sendo um desafio, exigendo confirmação da GPU para decisões de alto risco.

Significância e Alegações

O artigo afirma que os LLMs podem desempenhar um papel mais amplo na otimização de kernels além da mera geração de código. Ao atuarem como modelos virtuais de uma GPU, eles podem servir como filtros de alto throughput e baixo custo que gerenciam a escassez de recursos de hardware físico.

Os autores enfatizam que o valor primário desta abordagem não é apenas velocidade, mas a capacidade de realizar otimização baseada em modelo sem consultar o sistema original em cada etapa. O substituto oferece maior throughput e menor custo por consulta do que uma GPU, trocando o ruído da previsão, que pode ser gerenciado através de estimativas de incerteza calibradas. O trabalho conecta a otimização de kernels ao campo mais amplo de construção de simuladores neurais para sistemas físicos complexos, sugerindo um caminho para otimizadores que tratam a medição de hardware como um recurso escasso a ser usado seletivamente.