Alcançando eficiência extrema por meio da geração de kernels de GPU especializados

Tradicionalmente, os sistemas de inferência em produção dependem de kernels genéricos para lidar com diversos modelos e cargas de trabalho. Isso é abaixo do ideal porque os formatos das operações de GPU são determinados por uma combinação de parâmetros estáticos do modelo e fatores dinâmicos no momento da requisição; por exemplo, enquanto um modelo define uma das dimensões para uma multiplicação de matrizes, a outra dimensão flutua com base na contagem específica de tokens de cada requisição. Há um interesse crescente na geração de kernels de GPU por agentes, e esforços recentes têm se mostrado promissores. Exploramos uma questão central: se a geração de kernels pode ser automatizada, por que modelos de tamanhos amplamente diferentes (de 1 bilhão a 1 trilhão de parâmetros) deveriam depender do mesmo kernel? Ao especializar os kernels para os formatos específicos encontrados em tempo de execução, podemos alcançar uma eficiência extrema.

Neste blog, compartilhamos nossos sucessos e aprendizados ao usar agentes para gerar kernels de GPU. Criamos o Proteus, um sistema projetado para alcançar uma especialização extrema, o que exige um harness sob medida para otimização rigorosa, validação e gerenciamento de contexto.

Os harnesses de codificação convencionais costumam falhar aqui porque os agentes tendem a burlar a recompensa (reward-hack): seguindo a regra ao pé da letra, em vez do seu real propósito. Se você der um benchmark a um agente, ele pode otimizar o benchmark e não a operação pretendida.

Para resolver isso, o Proteus propõe kernels, os verifica em relação a uma implementação de referência controlada, mede o tempo dos que têm sucesso e melhora iterativamente os melhores resultados. Embora o processo seja simples, seu sucesso depende inteiramente da resolução de dois desafios fundamentais. A Figura 1 mostra a arquitetura simplificada do nosso design. Usando nosso harness Proteus, geramos kernels Qwen 3.5 122B que foram de 1,8 a 5,2 vezes mais rápidos do que os melhores disponíveis no vLLM.

Validação

Originalmente, tratávamos a busca de kernels como a parte difícil: como explorar um grande espaço de programas sem ficar preso em um platô sem melhorias? Na prática, a primeira pergunta era mais básica. Estamos medindo o que achamos que estamos medindo?

Um modelo otimiza a pontuação que você dá a ele. Ele não precisa de uma exploração exótica: pode ser simplesmente que a avaliação esteja fazendo uma suposição. Um exemplo foram os kernels para rotary position embeddings (RoPE), uma etapa comum em camadas de atenção. Um candidato poderia reutilizar código compilado restante de uma tentativa anterior e parecer mais barato do que uma reconstrução do zero. Outro poderia registrar um lote de inicializações de GPU em um gráfico (por exemplo, CUDA graph) e reproduzi-los como uma única unidade, enquanto a baseline com a qual comparamos ainda iniciava cada parte separadamente, de modo que os dois lados não estavam fazendo o mesmo trabalho. Outro era forte nos tamanhos de entrada que colocamos no conjunto de testes visível e fraco nos tamanhos que não haviam sido mostrados.

Por isso, dedicamos o trabalho inicial de design ao verificador, não ao prompt. Medimos o tempo de ambos os lados da mesma forma, inclusive com mais de um temporizador (por exemplo, temporizador de eventos CUDA, tempo de relógio e temporizador CUPTI) quando precisamos de uma verificação cruzada. Limpamos o estado compilado restante que não deve persistir e mantemos a ordem de configuração (setup) e desmontagem (teardown) consistente para que um lado não possa pular o trabalho pelo qual o outro ainda paga. Medimos o tempo dos vencedores novamente antes de usá-los como ponto de partida para a próxima rodada. Mantemos alguns testes que o candidato não pode ver, para que ele não se ajuste apenas ao exame. Para evitar "trapaças" de avaliação com desempenho artificialmente inflado, implementamos verificações de consistência automatizadas para sinalizar acelerações teoricamente impossíveis (por exemplo, >100x) que excedem a largura de banda física da GPU e os limites de computação. Isso protege contra as mesmas armadilhas de reward-hacking vistas em casos anteriores do setor, onde os agentes otimizavam para as métricas do harness em vez de ganhos reais de desempenho. Sem essas restrições, gerar mais kernels produzia principalmente mais ruído.

A ênfase no verificador também muda o gargalo da geração de kernels por agentes. No trabalho apenas de busca de programas (ou seja, otimização iterativa onde o sistema busca por programas gerando variantes repetidamente), bons candidatos são raros, de modo que escrevê-los domina o custo. Podemos produzir muitos rascunhos em paralelo, mas não podemos pular a validação. Temos que elaborar a validação com cuidado, e a verificação precisa ser executada em GPUs reais, de forma isolada e mais de uma vez. O sistema se move tão rápido quanto pode confiar em um kernel, não tão rápido quanto pode escrever um.

Gerenciamento de Contexto

Outro desafio é determinar o que o modelo de geração de kernels tem permissão para ver. É um trade-off. Dê ao modelo um prompt maior e ele terá mais informações: o melhor kernel atual, falhas recentes, dicas do profiler, notas de execuções anteriores. Isso pode ajudar. Também custa mais, porque pagamos por cada token que o modelo lê. E à medida que o prompt cresce, fica mais fácil para a próxima tentativa se desviar. Sinais úteis são misturados com conselhos obsoletos, dicas conflitantes e detalhes que se aplicam a um tamanho de entrada diferente ou a uma operação diferente. O modelo nem sempre sabe em quais frases confiar, por isso segue as mais chamativas, ou todas um pouco.

Dê muito pouco e o oposto acontece. Cada tentativa começa do zero. Os mesmos becos sem saída retornam. Nada se transfere da última execução, ou de uma operação relacionada, e o loop não avança.

Queríamos uma camada de conhecimento para ajudar com isso: lembrar o que funcionou, reutilizar mais tarde e fazer isso sem intervenção humana. Essa camada tem um segundo trade-off, entre o quão detalhada é uma lição armazenada e o quão amplamente ela se aplica.

Uma nota muito específica (“neste kernel, com este tamanho de entrada, desenrole este loop”) pode ser exatamente o que a próxima tentativa precisa. Também é fácil usá-la incorretamente na próxima operação, na próxima GPU ou em um tamanho de entrada diferente. Uma nota muito geral (“faça melhor uso da memória on-chip”) se aplica a quase qualquer lugar e não diz quase nada ao modelo sobre o que fazer. Vimos ambos os modos de falha. Quando as lições eram muito gerais, elas reafirmavam uma falha sem uma ação. Quando armazenávamos mais detalhes, eles frequentemente estavam muito vinculados a uma única execução para ajudar a próxima. Em uma execução longa, a maior parte do que o modelo leu e escreveu foi gasta buscando e roteando essa memória, em vez de escrever kernels. A camada de memória estava trabalhando muito, mas não estava melhorando o próximo candidato. A Figura 2 mostra o detalhamento do custo de tokens de um sistema desse tipo. O custo é dominado pelas camadas de conhecimento.

A versão de conhecimento que vale a pena manter é menor e encontra um equilíbrio entre generalidade e especificidade. Quando o modelo está prestes a escrever um kernel, seu prompt deve incluir apenas contexto de alta confiança: aprendizados práticos que associam situações específicas a ações (destilados do mapeamento anterior de modificação para impacto) e notas de falha concisas de execuções pai intimamente relacionadas. Recuperadas por meio de filtragem hierárquica de tags combinada com busca híbrida (palavra-chave + semântica), as lições devem ser específicas o suficiente para que se possa agir sobre elas e delimitadas o suficiente para esclarecer onde não se aplicam. Operações mais profundas, como reorganizar e destilar ainda mais o repositório de lições, pertencem a tarefas em segundo plano, e não a uma travessia síncrona de múltiplos saltos (multi-hop) sobre execuções anteriores a cada tentativa. Se um aprendizado não puder nomear a situação e a ação, não vale a pena colocá-lo no prompt. A Figura 3 mostra o detalhamento do custo de tokens após a correção da camada de conhecimento, sendo que a maioria dos tokens é gasta na geração de candidatos após essa correção.

Estudo de caso: decode compactado do Gated DeltaNet

Um exemplo concreto é o kernel de decode compactado (packed decode) no caminho do Gated DeltaNet no Qwen 3.5 122B. A operação atualiza um estado recorrente e escreve a saída do decode a partir de entradas QKV compactadas, parâmetros de porta (gate) e índices de estado. Usamos essa tarefa para exercitar o loop completo do Proteus em GPUs NVIDIA B200 com um backend Triton: validar o contrato da tarefa, medir a implementação de referência, solicitar kernels candidatos aos agentes, executar verificações estáticas e compilações, verificar a correção em relação à referência controlada, realizar o benchmark apenas dos candidatos verificados e, em seguida, medir novamente os melhores candidatos.

A Figura 4 é lida da esquerda para a direita. O nó da baseline ancora o benchmark em 0,025 ms. O Candidato 0000 é a semente segura: ele reproduziu a estrutura de decode compactado e passou na validação, mas foi mais lento que a referência, então o Proteus o manteve como um pai medido em vez de tratá-lo como uma vitória. A partir daí, o Proteus parou de otimizar um kernel genérico para cada formato e dividiu a busca em caminhos específicos para cada formato.

O caminho de reparo Batch-1 produziu um kernel específico para o formato no Candidato 012, atingindo 1,5x no formato de decode de lote único (single-batch). Os resultados mais fortes vieram no caminho de decode de serviço (serving-decode): o Candidato 030 encontrou a menor latência de kernel medida em 0,018 ms, e o Candidato 036 produziu a melhor aceleração de formato em 1,6x. Esse fragmento de serviço vencedor se especializou para o layout Batch=4, Key=128, Value=128 e processou a dimensão de valor em blocos de largura 64, por isso é um kernel seguro para esse formato específico, em vez de uma substituição universal.

O desvio final na linha do tempo mostra por que o trace importa. Tentativas posteriores de geração em C++ (em vez de Triton) encontraram falhas de compilação e geração, e a execução longa terminou após esgotar o orçamento de tentativas daquela ramificação. O artefato útil, portanto, não é apenas o candidato mais rápido. É o caminho completo mostrado na figura: falhas semânticas foram rejeitadas, kernels corretos, mas mais lentos, foram medidos, e os ganhos reais de desempenho foram mantidos vinculados ao formato que os torna seguros para compor um kernel de produção.

No que estamos trabalhando a seguir

O loop evolutivo que construímos é rigoroso como escritor. Ele costuma chamar o modelo em um padrão fixo: pegar o melhor kernel atual, tentar uma pequena edição, verificar, repetir. Isso retira a autonomia de que o agente precisa. Ele não pode facilmente alterar a estrutura, mudar de linguagem ou abandonar um design sem saída.

O loop ainda é necessário. Não para criar o kernel, mas para dar ao agente uma próxima dica confiável. Essa dica precisa vir de dois lugares.

Primeiro, a comunicação com a camada de conhecimento: alguns aprendizados que sejam específicos o suficiente para se agir sobre eles, e delimitados para sabermos onde não se aplicam. Sem isso, cada tentativa começa do zero.

Segundo, resultados de um verificador confiável: correção e tempo de execução que o próprio agente não mediu. Esses números são dicas para la próxima tentativa. Eles também são as únicas pontuações em que devemos acreditar. Se o agente medir o tempo de seu próprio trabalho, voltamos a caches residuais, comparações incompatíveis e testes que ele consegue ver.

Portanto, a divisão que queremos é mais estreita do que “agente versus loop”. Dê ao agente autonomia sobre como um kernel é escrito. Mantenha o loop como o canal para memória e avaliação. O agente propõe. O loop retorna o que ele tem permissão para ver e se a última proposta realmente venceu.

O Proteus construiu kernels especializados para partes do Qwen 3.5 122B em seu caminho Gated DeltaNet (um bloco no estilo de atenção linear) executado em GPUs NVIDIA B200. Os speedups nos kernels individuais ficaram na faixa de 1,8x a 5,2x.

A lição é que a geração é a etapa barata. A validação e o gerenciamento de contexto são a parte difícil. É aí que o tempo de design cuidadoso e as inovações são necessários.

A geração agentiva de kernels de GPU desbloqueou o incrível potencial da especialização extrema, mas a construção de harnesses confiáveis e prontos para produção continua sendo uma fronteira desafiadora. Estamos enfrentando os desafios mais difíceis na interseção de AI e sistemas, e estamos procurando engenheiros ousados para se juntarem a nós na construção do futuro da inferência eficiente. Se você tem paixão por superar os limites do que é possível, estamos contratando!

(Esta publicação no blog foi traduzida utilizando ferramentas baseadas em inteligência artificial) Publicação original