Resumo Técnico: KForge – Geração de Kernels Multiplataforma Impulsionada por LLM para Aceleradores de IA

Declaração do Problema

A proliferação de aplicações de IA agêntica deslocou as cargas de trabalho de inferência de execuções monolíticas para pipelines heterogêneos. Esses pipelines intercalam raciocínio, chamadas de ferramentas e coordenação multiagente, com cada estágio possuindo perfis distintos de computação e memória. Consequentemente, o deployment ideal requer o mapeamento de estágios específicos do pipeline para o acelerador mais adequado. Essa heterogeneidade cria um desafio significativo de sistemas: desenvolvedores devem produzir kernels de alto desempenho em uma crescente diversidade de backends de hardware (NVIDIA, AMD, Intel, Apple) e modelos de programação (CUDA, Triton, SYCL, Metal, etc.).

Escrever e portar esses kernels manualmente é demorado, exige profunda expertise de baixo nível e não escala conforme a complexidade do kernel aumenta. Embora os Grandes Modelos de Linguagem (LLMs) tenham mostrado promessa na geração de código, aplicar-os à otimização de kernels de baixo nível permanece difícil devido à fragilidade do código de performance (onde mudanças menores no layout de memória ou precisão podem quebrar a corretude) e à escassez de dados de treinamento para plataformas que não sejam CUDA. Além disso, os benchmarks existentes frequentemente avaliam kernels isoladamente, falhando em capturar o impacto da otimização do kernel na execução end-to-end de modelos completos dentro de runtimes de produção.

Metodologia: O Framework KForge

O KForge é um framework de síntese de programas autônomo, multiplataforma e de múltiplos estágios, projetado para espelhar o fluxo de trabalho iterativo de um engenheiro de kernels. Ele emprega um loop colaborativo entre dois agentes especializados baseados em LLM:

1. Agente de Geração: Responsável por produzir e refinar progressivamente candidatos a kernels. Ele opera dentro de um loop impulsionado por prompts que incorpora:

   • Refinamento Iterativo: Utilizando erros de compilação, erros de tempo de execução e discrepâncias numéricas de iterações anteriores para corrigir o código.
   • Tradução Cross-Platform: Aproveitando implementações funcionais existentes (ex: CUDA) como referências para gerar código para plataformas sub-representadas (ex: Metal ou SYCL).
   • Guardrails de Candidato: Aplicando restrições para evitar comportamentos indesejados, como redução de tipo de dados (downcasting) ou mudanças de precisão.

2. Agente de Análise de Performance: Um agente especializado, distinto do gerador, para lidar com a esparsidade dos sinais de otimização dentro de grandes conjuntos de dados de profiling. Ele interpreta diversos inputs de profiling (métricas programáticas de ferramentas como Nsight Systems ou saídas visuais do Xcode Instruments) para gerar recomendações concretas de nível de código. Essas recomendações visam métricas de utilização de hardware, como largura de banda de memória, ocupação de warps e intensidade aritmética.

O Loop de Síntese:
O sistema alterna entre duas fases:

• Passagens Funcionais: O agente de geração itera até que o código compile, execute sem erros e produza saídas numericamente corretas.
• Passagens de Otimização: Uma vez estabelecida a corretude, o agente de análise de performance fornece feedback para direcionar a próxima rodada de síntese visando fechar a lacuna de performance em relação às baselines ajustadas manualmente.

O framework suporta uma interface uniforme para quatro fornecedores (NVIDIA, AMD, Intel, Apple) e seis modelos de programação, utilizando um registro de modelos para permanecer agnóstico ao LLM subjacente.

Principais Contribuições

• Framework KForge: Introdução de um sistema de síntese autônomo de múltiplos estágios onde os agentes de geração e análise de performance colaboram usando feedback de compilação, corretude e profiling.
• Arquitetura Cross-Platform: Um design de sistema que suporta quatro fornecedores de aceleradores e seis modelos de programação através de uma interface unificada, permitindo a transferência de ideias algorítmicas entre plataformas.
• Avaliação de Modo Duplo: Casos de estudo demonstrando a eficácia do framework em dois cenários distintos: competindo contra runtimes maduros e otimizados pelo fornecedor e servindo como uma ferramenta de bring-up para hardware emergente carente de implementações de referência.

Resultados Experimentais

Os autores avaliaram o KForge em dois backends com diferentes disponibilidades de baseline:

1. NVIDIA B200 (Competindo com Runtimes Otimizados pelo Fornecedor)

• Contexto: Otimização de kernels de caminho de decodificação para o modelo gpt-oss-20b contra o TensorRT-LLM v1.3.0rc9.
• Micro-benchmarks: Os kernels gerados pelo KForge superaram consistentemente a baseline em vários tamanhos de batch. Notavelmente, o kernel "MoE finalize" mostrou um aumento de velocidade partindo de 1.06× em pequenos batches para 1.43× em um batch de 128.
• Performance End-to-End: A integração desses kernels otimizados na pipeline de inferência completa resultou em uma melhoria de 2.12% no throughput end-to-end (2656.61 vs. 2601.55 tokens/seg) e uma redução de 2.07% no tempo de wall-clock. Os autores observam que, embora esses ganhos pareçam modestos, eles são significativos dado os anos de esforço de engenharia já investidos na baseline do fornecedor.

2. Intel Arc B580 (Bring-up de Hardware Emergente)

• Contexto: Geração de kernels Triton para o Intel Arc B580 (Battlemage), onde não existem referências ajustadas manualmente comparáveis.
• Benchmark: Avaliado em 37 workloads de GEMM + operações de cauda do KernelBench Level 2.
• Resultados: O KForge alcançou um ganho de velocidade de média geométrica de 5.13× sobre o mais rápido entre o modo eager do PyTorch e o torch.compile.
• Técnicas de Otimização: A análise dos kernels gerados revelou a aplicação bem-sucedida de:
  • Fusão por Camada (Layer-wise Fusion): Colapso de operações pós-GEMM em kernels únicos para eliminar viagens de ida e volta de tensores completos.
  • Execução de Precisão Mista: Utilização de FP16 para matmuls em unidades XMX com acumulação em FP32.
  • Reduções de Passagem Única: Substituição de algoritmos multi-pass para GroupNorm e LogSumExp por algoritmos de streaming online.

Significância e Alegações

O artigo posiciona o KForge como uma ferramenta prática tanto para otimização de nível de produção quanto para bring-up rápido de hardware. Os autores afirmam que:

• A Síntese Agêntica é Viável: A separação entre geração e análise de performance permite que os LLMs naveguem efetivamente pelos complexos trade-offs da otimização de kernels de baixo nível.
• Escalabilidade dos Ganhos: Mesmo melhorias percentuais marginais (ex: ~2%) contra runtimes estabelecidos e otimizados pelo fornecedor são significativas quando aplicadas em escala de deployment, acumulando-se em ganhos de eficiência substanciais.
• Preenchendo a Lacuna: O framework aborda com sucesso o problema de "cold start" para novos hardwares (como o Intel Arc) ao gerar kernels performantes sem exigir referências pré-existentes ajustadas manualmente, enquanto ainda compete com runtimes estabelecidos em hardware maduro.

O trabalho conclui que a síntese de kernels agêntica está transitando de um conceito de pesquisa para um componente funcional da stack de sistemas de IA, capaz de lidar com a heterogeneidade das cargas de trabalho de inferência modernas.