Resumo Técnico: AutoMegaKernel (AMK)

Definição do Problema

A decodificação autorregressiva de fluxo único (single-stream) com tamanho de lote (batch size) 1 é fundamentalmente limitada por largura de banda. O piso de latência é determinado pela razão entre o total de bytes de peso e a largura de banda da HBM ($t_{min} = \text{bytes\_de\_peso} / \text{largura\_de\_banda\_HBM}$). Frameworks de execução convencionais (ex: PyTorch, cuBLAS) falham em se aproximar deste piso porque lançam um kernel separado para cada operador, incorrendo em latência significativa de lançamento de CPU e forçando as ativações a fazerem o trajeto de ida e volta pela HBM entre as operações. Embora os CUDA Graphs amortizem o overhead de lançamento, eles mantem as fronteiras de inter-operação de kernels e os trajetos de ida e volta pela HMM.

Abordagens de "megakernel" existentes (ex: Mirage Persistent Kernel) ou kernels otimizados manualmente abordam o piso de latência fundindo todo o passo de propagação (forward pass) em um único kernel persistente. No entanto, esses sistemas carecem de confiança e portabilidade: frequentemente omitem portões de segurança estáticos (correndo o risco de deadlocks ou condições de corrida), visam modelos/arquiteturas específicas manualmente e não fornecem uma interface estruturada para que agentes autônomos otimizem o sistema com segurança.

Metodologia

O AutoMegaKernel (AMK) é um sistema autônomo, impulsionado por agentes, que compila modelos da família Llama do HuggingFace em um único kernel cooperativo persistente sem a necessidade de escrita manual de CUDA para cada modelo. O sistema é estruturado em quatro camadas e dois loops de pesquisa autônoma:

1. Camada 0 (A VM): Um megakernel persistente escrito à mão, congelado e por arquitetura (vm/scheduler.cu). Ele lança um threadblock cooperativo por Streaming Multiprocessor (SM). Cada bloco possui uma fila local e percorre uma ordem topológica global de tarefas, executando instruções wait → dispatch → signal. Ele depende de contadores monotônicos para sincronização, evitando locks.
2. Camada 1 (Instruções): Micro-kernels (GEMV, atenção, RMSNorm, etc.) conformes a uma ABI que são verificados quanto à correção isoladamente.
3. Camada 2 (Scheduler): A superfície de edição para o agente. Ele importa um modelo, reduz o modelo para um IR de Schedule tipado (um DAG de tarefas sincronizado por contadores monotônicos) e propõe configurações estruturais (tiling, fusão, atribuição de SM, profundidade de pipelining) via um objeto ScheduleConfig.
4. Camada 3 (Roadmap): Espaço reservado para dinamismo futuro (ex: continuous batching).

O Portão de Segurança (Safety Gate):
A inovação central é um validador estaticamente verificado (validate() em schedule/ir.py). Antes de qualquer lançamento na GPU, o validador realiza um conjunto de verificações de grafo estático no Schedule IR para certificar:

• Ausência de Deadlock: Garantir que o DAG de produtor-consumidor seja acíclico, que os limiares de espera sejam satisfatórios ($1 \le t \le \#producers$) e que as ordens de fila por SM sejam válidas.
• Ausência de Race Conditions: Garantir que os contadores compartilhados sejam unidos corretamente (todos os produtores devem terminar antes que um consumidor espere por um contador de múltiplos produtores) e que os acessos à memória respeitem as relações de happens-before.
• Corretude: Verificar a alcançabilidade da saída e a integridade referencial.

Se um schedule falha nessas verificações, ele é rejeitado imediatamente. Isso evita que schedules propostos por agentes de forma insegura causem travamentos nas GPUs.

Loop de Pesquisa Autônoma (Autoresearch Loop):
O AMK inclui um loop não supervisionado e controlável por agentes que propõe iterativamente mudanças de ScheduleConfig, valida essas mudanças contra uma referência de CPU, mede a latência e mantém ou reverte as alterações. Este loop melhora autonomamente o megakernel sobre sua própria linha de base.

Principais Contribuições

1. Schedule IR Estaticamente Verificado: Um validador congelado que certifica a ausência de deadlock e race conditions via verificações de grafo estático (não provas mecanizadas) sobre uma base confiável. Ele alcançou zero falsos-positivos em 7.160 schedules construídos adversariamente (6.091 inseguros) enquanto aceitou todos os 360 reduções de modelos reais.
2. Cobertura de Geração Automática: O AMK gera automaticamente megakernels corretos para 10 de 10 modelos suportados (incluindo checkpoints reais como SmolLM2-135M e TinyLlama-1.1B) sem zero escrita manual de CUDA. Ele rejeita com sucesso variantes não suportadas (ex: projeções enviesadas) ou documenta pontos cegos honestamente.
3. Retargeting Automático: Um único código-fonte gera megakernels corretos para sm_80 (A100), sm_90 (H100) e sm_120 (RTX 5090), correspondendo ao modo eager bit a bit ou dentro de tolerâncias estritas de fp32 ($\le 3.8 \times 10^{-5}$).
4. Megakernels Quantizados: O sistema gera automaticamente kernels de pesos quantizados int8 (lossless-greedy, 1.12× mais rápido) e int4 (lossy, 2.42× redução de tráfego de peso) a partir do mesmo caminho de IR.
5. Compilador Drivável por Agentes: O sistema trata a corretude como um invariante estrutural, permitindo que um agente busque no espaço de schedules de forma segura sem o risco de travamentos de runtime.
6. Estudo de Performance Honesto: O artigo reporta resultados em GPUs de classe de inferência e de treinamento, divulgando explicitamente onde o sistema vence e onde ele fica atrás das linhas de base estabelecidas.

Resultados Experimentais

Corretude e Portabilidade:

• Na RTX 5090, A100 e H100, o AMK reproduz o decode token-a-token do HuggingFace e iguala a perplexidade de teacher-forced para $2.5 \times 10^{-7}$ no checkpoint SmolLM2-135M.
• O validador rejeitou todos os 6.091 schedules inseguros no teste de estresse, demonstrando uma taxa de falso-aceite de 0%.

Performance em GPUs de Classe de Inferência (A Vitória):

• RTX 5090 (Consumidor): Um megakernel auto-otimizado de peso apenas int8 (W8A16) supera o cuBLAS (bf16) com CUDA-graph por 1.19–1.23× no decode de batch-1. Esta é uma comparação de assimetria de precisão (int8 vs. bf16) onde a vitória vem de transmitir menos bytes de peso (0.61× os bytes), não de um kernel por byte mais rápido.
• Frota de Inferência de Datacenter: O megakernel int8 supera o cuBLAS (bf16) com CUDA-graph em toda a frota de classe de inferência:
  • NVIDIA L4: Até 1.33× mais rápido (cresce com o tamanho do modelo).
  • L40S: 1.25–1.27× mais rápido.
  • A10G: Até 1.08× mais rápido em escala.
• Auto-melhoria: O loop de pesquisa autônoma melhorou o megakernel autonomamente sobre o seu padrão em 1.25–1.72× na RTX 5090.

Performance em GPUs de Classe de Treinamento (A Derrota):

• Em GPUs de alta largura de banda de classe de treinamento (A100, H100), o megakernel gerado automaticamente fica atrás do cuBLAS com CUDA-graph e do vLLM.
  • Na H100, o AMK foi 1.65× mais lento que o caminho padrão de CUDA-graph do vLLM.
  • Na A100, o GEMV otimizado em bf16 atingiu apenas ~63% do pico de HBM medido (contra ~90% do cuBLAS).
• Diagnóstico: O gargalo é o custo de sincronização entre SMs por tile (um custo fixo que o megakernel paga e o cuBLAS não). Em silício de alta largura de banda, esse overhead de sincronização domina, enquanto em GPUs de classe de inferência, a economia de largura de banda proveniente da fusão e quantização amortiza esse custo.

Significância e Alegações

O artigo afirma que o AutoMegaKernel representa uma mudança de "velocidade bruta" para segurança sistemática e portabilidade na síntese de megakernels. Sua principal contribuição é uma arquitetura de compilador estaticamente verificada que permite a agentes autônomos explorar com segurança o espaço de schedules de modelos completos.

• Confiança: Substitui o risco de travar GPUs por um portão de segurança estático que rejeita schedules inseguros antes do lançamento.
• Generalidade: Alcança retargeting automático entre arquiteturas (sm_80/90/120) e famílias de modelos sem ajuste manual de CUDA.
• Honestidade: Os autores enquadram explicitamente os resultados de performance como dependentes de regime. O sistema vence em GPUs de classe de inferência ao fundir operações e utilizar quantização para reduzir o tráfego de memória, mas atualmente perde para bases de alto desempenho em hardware de classe de treinamento devido aos overheads de sincronização que exigem agendamentos de grão mais grosso (trabalho futuro), em vez de apenas melhores micro-kernels.

O artigo conclui que, embora a qualidade atual do kernel ainda não iguale as bases otimizadas manualmente em hardware de alto nível para treinamento, o harness estaticamente verificado fornece a fundação necessária para automatizar com segurança as melhorias futuras (ex: sincronização de grão mais grosso, pipelining mais profundo) necessárias para fechar essa lacuna.