Resumo Técnico: Xe-Forge

Declaração do Problema

A portabilidade de algoritmos de aprendizado profundo para novos aceleradores de hardware, especificamente GPUs Intel, exige que os desenvolvedores apliquem manualmente otimizações repetitivas de baixo nível (quantização, coalescência de memória, ajuste de tiles, soluções específicas de arquitetura) a cada kernel Triton. Esse esforço manual cria um gargalo significativo: enquanto os padrões de otimização são consistentes, as restrições de hardware (contagens de warps, modos GRF, dimensionamento de SLM) variam entre dispositivos e frequentemente estão ausentes dos dados de treinamento de Modelos de Linguagem Grandes (LLMs). Sistemas existentes baseados em LLM focam principalmente na geração de kernels do zero para hardware NVIDIA ou AMD, deixando uma lacuna crítica na otimização de kernels Triton existentes e funcionalmente corretos para a arquitetura específica da Intel (séries Arc/Arc Pro).

Metodologia

Os autores apresentam o Xe-Forge, um pipeline multiestágio alimentado por LLM projetado para automatizar a otimização de kernels Triton existentes para GPUs Intel. O sistema opera através dos seguintes componentes principais:

1. Arquitetura de Pipeline Multiestágio

O Xe-Forge processa um kernel Triton funcionalmente correto através de até nove estágios sequenciais de otimização. Um planejador baseado em LLM determina a ordem de execução com base em problemas detectados, sujeito a restrições de dependência rígidas (por exemplo, a reestruturação algorítmica deve preceder a conversão de dtype). Os estágios incluem:

1. Análise: Detecta oportunidades de otimização.
2. Otimização Algorítmica: Reestrutura o cálculo para reduzir FLOPs e acesso à memória (por exemplo, eliminando GEMMs redundantes).
3. Descoberta: Aplica otimizações abertas e inovadoras não cobertas por estágios nomeados.
4. Correção de Dtype: Converte precisões desnecessárias (por exemplo, float64 para float32/16).
5. Fusão: Mescla kernels separados ou cadeias elementwise.
6. Acesso à Memória: Otimiza coalescência, pré-busca e layout de dados.
7. Ponteiros de Bloco: Substitui aritmética de ponteiros manual por tl.make_block_ptr.
8. Kernel Persistente: Reestrutura loops para reutilização de dados entre iterações.
9. Ajuste Específico de GPU: Aplica restrições específicas da Intel (contagens de warps, modos GRF, dimensionamento de tiles, embaralhamento GROUP_SIZE_M).
10. Autotuning: Gera uma grade de configuração @triton.autotune.

Estágios sem problemas associados são ignorados para reduzir custo e alterações no código.

2. Cadeia de Verificação e Refinamento (CoVeR)

Cada estágio é executado por um agente CoVeR (implementado em DSPy) que gera candidatos iterativamente, valida-os e refina com base no feedback. O agente opera em um loop:

• Gerar: Propõe alterações de código.
• Verificar: Executa uma cascata de quatro níveis:
  1. Sintaxe: Análise de AST Python.
  2. Estrutura: Verifica restrições de hardware (por exemplo, warps potência de dois, modos GRF válidos).
  3. Corretude: Compara saídas contra uma referência PyTorch usando torch.allclose.
  4. Desempenho: Realiza benchmarks em hardware real; rejeita se for mais lento que a linha de base.
• Refinar: Se a verificação falhar, diagnósticos de erro são alimentados de volta no contexto para a próxima iteração.

3. Base de Conhecimento Curada

Para abordar a falta de dados específicos da Intel em corpora de treinamento de LLM, o Xe-Forge utiliza uma base de conhecimento YAML legível por máquina. Esta codifica restrições rígidas (por exemplo, contagens de warps válidas, limites de capacidade de SLM) e padrões de otimização (transformações de código antes/depois) específicas para as arquiteturas Intel Xe. Isso garante que o modelo permaneça dentro de limites arquitetonicamente válidos.

4. Infraestrutura de Validação (AI Bench)

O sistema depende do AI Bench, um framework autônomo que desacopla a medição da lógica de otimização. Ele garante temporização reproduzível (usando temporização de eventos de hardware com limpeza de cache L2) e verificação de corretude. Crucialmente, o pipeline impõe uma separação de harness de kernel: o LLM pode apenas modificar o corpo do kernel, enquanto o wrapper Model, a geração de entrada e o caminho de execução permanecem fixos. Isso previne otimizações adversariais onde o modelo contorna a lógica de benchmark.

Contribuições Principais

1. Formulação de Otimização Triton para Triton: Diferentemente das abordagens existentes de geração PyTorch para Triton, o Xe-Forge otimiza kernels existentes preservando a semântica.
2. Primeiro Sistema Focado em GPU Intel: O primeiro sistema baseado em LLM direcionado a GPUs Intel, apresentando uma base de conhecimento consciente de hardware para restrições ausentes nos dados de treinamento.
3. Pipeline CoVeR Multiestágio: Uma abordagem decomposta onde cada estágio possui padrões dedicados e critérios de verificação, equilibrando flexibilidade com reprodutibilidade.
4. Framework AI Bench: Uma ferramenta de benchmark que desacopla execução/temporização da lógica de otimização, garantindo comparação justa entre backends.
5. Mitigação Adversarial: Um design que fixa o harness de execução para impedir que o LLM "trapaceie" o benchmark contornando o modelo.

Resultados Experimentais

O sistema foi avaliado em 97 kernels do KernelBench Nível 2 e configurações de Flash Attention na GPU Intel Arc Pro B70.

• Desempenho Geral: Alcançou um acréscimo geométrico médio de 1,17× sobre o PyTorch eager, com 67% dos kernels mostrando melhoria.
• Acelerações Extremas: Nove kernels superaram 5× de aceleração, com o mais alto atingindo 82× (Matmul_Min_Subtract). Esses ganhos foram impulsionados pela reestruturação algorítmica que eliminou cálculos redundantes, efetivamente reduzindo a contagem de FLOP abaixo da contagem original.
• Flash Attention: Alcançou acelerações de 2× a 13,3× em todas as 16 configurações testadas (variando tamanhos de modelo, comprimentos de sequência e tamanhos de lote) sem regressão. Notavelmente, o sistema permitiu inferência com sequência de 32k em uma GPU de 32GB, utilizando kernels persistentes e tiling SLM, o que era inviável com o kernel ingênuo.
• Corretude: 100% dos kernels otimizados passaram nos testes de corretude contra a referência PyTorch.
• Análise Roofline: Para kernels limitados por computação (GEMM/MatMul), o pipeline elevou o desempenho para eficiência próxima ao pico. Para kernels limitados por largura de banda, o desempenho foi geralmente competitivo, embora alguns mostraram leves regressões (0,5–0,8×) devido a limites de vazão de memória.

Significado e Alegações

O artigo afirma que conhecimento de domínio estruturado combinado com verificação hardware-no-loop é mais eficaz do que a escala do modelo isoladamente para otimização de kernels específicos de hardware.

• Preenchendo a Lacuna de Conhecimento: Ao codificar restrições específicas da Intel (contagens de warps, modos GRF) em uma base de conhecimento curada, o sistema compensa a falta de dados de treinamento do LLM sobre essas arquiteturas específicas.
• Eliminação Sistemática do Esforço de Portabilidade: O Xe-Forge demonstra que o trabalho repetitivo de portabilidade pode ser automatizado, removendo sistematicamente o gargalo que atualmente impede a implantação de algoritmos em novos aceleradores.
• Robustez: A separação do código de kernel mutável do código de harness imutável previne efetivamente a classe de falhas de "evasão" observadas em outros sistemas de geração, onde modelos contornam restrições para passar em testes de corretude.

Os autores concluem que, embora o pipeline não seja perfeito (uma minoria de kernels mostra leves regressões e padrões algorítmicos inovadores permanecem desafiadores), ele automatiza com sucesso a otimização de operadores padrão e composições fundidas que constituem a maior parte das cargas de trabalho de produção.