Resumo Técnico: KernelEvolve

Declaração do Problema

A implantação de Modelos de Recomendação de Aprendizado Profundo (DLRMs) para sistemas de publicidade em larga escala na Meta enfrenta uma "maldição da dimensionalidade" impulsionada por três fatores que se interceptam: diversidade de arquitetura de modelo (variando de MLPs tradicionais a modelos de sequência baseados em Transformer), diversidade de primitivas de kernel (estendendo-se além do GEMM padrão para centenas de operadores complexos de pré-processamento de dados) e heterogeneidade de hardware (abrangendo GPUs NVIDIA, GPUs AMD e Aceleradores de Treinamento e Inferência proprietários da Meta, como o MTIA).

Essa complexidade cria um gargalo crítico. O desenvolvimento manual de kernels é economicamente inviável, exigindo de 2 a 8 semanas de esforço especializado por combinação de operador-hardware. Além disso, a falta de kernels otimizados para operadores de pré-processamento específicos em aceleradores proprietários força arquiteturas de "serviço desagregado". Nessas configurações, o pré-processamento deve ser executado em CPUs hospedeiras subprovisionadas enquanto a inferência da rede neural ocorre nos aceleradores, incorrendo em uma sobrecarga de latência de rede pura de 10 a 20 ms. Esse imposto arquitetural consome uma parte significativa do orçamento de latência inferior a 100 ms exigido para o serviço de anúncios, impactando diretamente a experiência do usuário e a receita. Os agentes de codificação baseados em IA existentes são limitados por escopos estreitos, benchmarks sintéticos, foco em plataforma única e uma falta de gestão de conhecimento autônoma e persistente necessária para a otimização de nível de produção.

Metodologia

KernelEvolve é um framework agêntico projetado para automatizar a geração, otimização e validação de kernels de alto desempenho em hardware heterogêneo. Ele trata a otimização de kernel como um problema de busca baseado em grafos, em vez de uma simples tarefa de geração de código.

Arquitetura Principal

1. Busca em Árvore e Máquina de Estados: O sistema modela o processo de otimização como um grafo de busca $G_t = (V_t, E_t)$. Ele emprega estratégias que incluem Busca Gananciosa (Greedy Search), Busca em Árvore de Monte Carlo (MCTS) e algoritmos evolutivos. O processo itera através de seleção, transformação e pontuação:

   • Função de Aptidão (Fitness Function): Avalia os kernels com base no ganho de velocidade (speedup) sobre os baselines do PyTorch, tendo a correção (precisão numérica) como uma restrição rígida (falha = aptidão 0).
   • Operador Universal: Ao contrário de frameworks multi-operador estáticos, o KernelEvolve utiliza um único operador sensível ao contexto. Ele sintetiza prompts dinamicamente com base no contexto de execução em tempo de execução (erros, dados de profiling) em vez de templates fixos.

2. Recuperação Agêntica e Contexto Autogerenciado:

   • Sub-agente de Memória de Contexto: Analisa artefatos de execução (logs, métricas de profiling, resultados de correção) para diagnosticar gargalos e sintetizar diretrizes de otimização. Ele mantém um armazenamento de metadados persistente (banco de dados relacional) rastreando a árvore de busca, pontuações de aptidão e relações pai-filho, permitindo tolerância a falhas e reutilização de conhecimento entre sessões.
   • Sub-agente de Busca Profunda: Realiza a recuperação direcionada de uma Base de Conhecimento Persistente. Esta base de conhecimento codifica restrições específicas de hardware (NVIDIA, AMD, MTIA), padrões de otimização e metodologias de depuração em um sistema de arquivos hierárquico. Para hardware proprietário como o MTIA, que carece de representação nos corpora de treinamento de LLMs, a injeção deste conhecimento de domínio (ex: operações SFU do MTIA, primitivas de comunicação inter-PE) é crítica para gerar código correto.

3. Framework de Avaliação e Ferramental:

   • Interface de Implementação Dupla: Os kernels gerados são encapsulados em uma interface padronizada junto com um baseline do PyTorch para validação automática de correção (torch.allclose) e benchmarking de desempenho.
   • Interpretadores de Hardware: Ambientes dedicados (baseados em Bento) para NVIDIA, AMD e MTIA garantem execução reproduzível com toolchains atualizadas.
   • Profiling Unificado (Triton MPP): O Multi-Pass Profiler da Meta unifica a instrumentação através das camadas de abstração (MLIR, assembly, runtime), fornecendo dados de desempenho estruturados ao nível de instrução sem a fragmentação das toolchains tradicionais.
   • Integração FaaS: A avaliação é delegada a uma plataforma Function-as-a-Service, desacoplando a geração limitada pela CPU da avaliação limitada por GPU/acelerador para maximizar a utilização de recursos.

Principais Contribuições

• Sistema de Kernel de IA de Nível de Produção: O KernelEvolve é o primeiro sistema de otimização de kernel impulsionado por IA implantado em escala industrial para treinamento e inferência de modelos de recomendação. Ele gera autonomamente kernels Triton otimizados para centenas de modelos atendendo bilhões de usuários.
• Otimização Automatizada para Hardware Proprietário: Ao integrar uma base de conhecimento persistente com síntese de prompt assistida por recuperação, o sistema consegue efetivamente gerar kernels para arquiteturas proprietárias (MTIA) ausentes nos dados de treinamento de LLMs, preenchendo a lacuna entre as capacidades do modelo e as restrições de hardware.
• Ganhos de Desempenho e Eficiência: O sistema reduz o tempo de desenvolvimento de semanas para horas. Ele alcança ganhos substanciais de velocidade (1.2x a 17x) sobre os baselines do PyTorch em diversos workloads, incluindo pré-processamento de dados, transformadores convolucionais e modelos de ranking.
• Insights Operacionais: O artigo compartilha análises de modos de falha, estratégias de depuração para geração de kernels incorretos e padrões organizacionais para integrar fluxos de trabalho automatizados em ambientes de produção.

Resultados

O KernelEvolve foi validado no conjunto KernelBench publicamente disponível e em um conjunto customizado de 160 operadores ATen do PyTorch em três plataformas heterogêneas (NVIDIA H100, AMD MI350, MTIA v3).

• Correção: Alcançou uma taxa de aprovação de 100% em todos os 250 problemas do KernelBench e 100% de correção em todas as 480 configurações de operador-plataforma.
• Desempenho: Demonstrou ganhos de velocidade variando de 1.25x a 17x sobre os baselines do PyTorch. Melhorias notáveis incluem:
  • Pré-processamento de Dados: MapId (4.1x), MergeBucketizedDense (9.3x), Batch Event Truncate (9.8x).
  • Transformadores Convolucionais: Conv1d (6.5x), Conv2d (4.7x).
  • Modelos de Ranking: WuKong Optimized FM (4.0x), InterFormer PFFN (2.5x).
  • Especificidades de MTIA: RMSNorm 2D backward (17x).
• Eficiência de Busca: As trajetórias de otimização mostram que a busca baseada em grafos com feedback de execução melhora consistentemente as pontuações de aptidão ao longo dos passos iterativos, descobrindo estratégias de fusão e configurações de tiling que o desenvolvimento manual levaria semanas para explorar.

Significância

O artigo argumenta que o KernelEvolve transforma fundamentalmente o desenvolvimento de kernels de um processo manual dependente de expertise em um serviço automatizado e escalável. Sua significância reside em:

1. Viabilização de Implantação de Aceleradores Monolíticos: Ao automatizar a geração de kernels para todos os operadores (incluindo pré-processamento de baixa intensidade aritmética), ele remove a restrição binária que força arquiteturas de CPU-acelerador desagregadas. Isso elimina o imposto de latência de rede de 10–20ms, melhorando diretamente a eficiência de serviço e reduzindo o Custo Total de Propriedade (TCO).
2. Escalabilidade para Heterogeneidade: Demonstra que a IA agêntica pode navegar efetivamente na explosão combinatória de pares operador-hardware, tornando a implantação de novos modelos em novas gerações de hardware viável sem esperar pela portabilidade manual de kernels.
3. Preenchimento da Lacuna de Conhecimento: A abordagem de injetar conhecimento de hardware proprietário em uma base de conhecimento persistente permite que os LLMs gerem código de alta qualidade para aceleradores customizados sem necessidade de retreinamento, oferecendo um modelo para o design e otimização de sistemas de IA futuros em escala.