Vibecoding GPU Kernels

Escrever kernels de GPU na mão exige paciência, esforço e também sua sanidade. Tenho experimentado LLMs modernos (Claude Opus 5, GPT 5.6 Sol) para literalmente gerar kernels para mim. Vou falar sobre o que aprendi fazendo isso.

Honestamente, escrever kernels de GPU é uma tarefa vergonhosamente verificável.

Primeiro, você decide quais operações deseja escrever o kernel. Segundo, você escreve a primeira versão e garante que ela compila sem erros óbvios. Terceiro, você verifica a corretude com uma implementação de referência lenta. Se não corresponder, você tenta corrigir o problema de corretude. Depois disso, você mede o tempo de execução do kernel. As versões seguintes são construídas em cima disso, e você continua otimizando até atingir as métricas de roofline ou ficar satisfeito.

Loop verificável: desenvolvimento de kernel de GPU

A imagem acima mostra isso como um loop verificável:

• A verificação de compilação é um loop local rápido (B <-> C) antes mesmo de você pensar em corretude.
• A verificação de corretude (D <-> E <-> F) é o núcleo do loop de recompensa verificável. É essa parte que torna a escrita de kernels um bom problema para verificação automatizada, já que você tem uma referência de ground truth para comparar.
• A otimização (G -> H -> de volta a D) reutiliza o mesmo loop de corretude para cada nova versão, já que um kernel rápido, mas incorreto, não tem valor; a corretude precisa ser verificada em cada versão.
• A verificação de roofline/satisfação (I) é o loop externo que decide se você continua otimizando ou para.

Primeira versão

A imagem acima ainda é uma visão de alto nível, e o diabo está nos detalhes. Precisamos garantir que nosso agente de LLM tenha todo o contexto necessário para sequer começar a escrever uma boa primeira versão.

Os DSLs de CUDA entram em ação aqui. Triton, CuTeDSL e Tilelang são os mais fáceis de começar a usar em Python. A curva de aprendizado é menos acentuada do que a do CUDA C++; no entanto, as abstrações desses DSLs podem confundir ainda mais nosso agente. Precisamos de uma forma de passar o contexto dessas abstrações para o agente.

LLMs modernos já sabem escrever Triton "bom". Eles funcionam bem com as abstrações do Triton mesmo sem nenhum contexto. No entanto, para outros DSLs como CuTeDSL (que oferece muito mais controle que o Triton), descobri que ter um diretório de contexto onde o agente possa procurar entender as abstrações do DSL ajuda muito.

Por exemplo, clonar o repositório

NVIDIA cutlass

no diretório de contexto é uma boa forma de deixar o agente procurar abstrações relacionadas a

Álgebra de Layout, átomos de Copy/GEMM, hierarquia de memória, kernels de exemplo

e assim por diante ao escrever kernels em CuTeDSL.

Na minha experiência, uma boa primeira versão do kernel compila sem erros óbvios e passa no teste de corretude que vou mencionar abaixo.

Teste, Benchmark e Profiling

Dado contexto suficiente ao agente, o gargalo real agora muda para a validação. A própria implementação de referência, e a validação contra ela, torna-se cada vez mais importante. Chamo essa fase de teste de corretude ou simplesmente teste. A velocidade da implementação de referência não importa tanto quanto sua intenção. O que você pretende medir e verificar é o que seu agente vai otimizar.

Normalmente, quando o cálculo não deve ocorrer em precisão reduzida, meço o Erro Máximo Absoluto/Relativo (MAE), o Erro Quadrático Médio (MSE/RMSE) e o PSNR (Peak Signal to Noise Ratio). Quando precisões reduzidas estão envolvidas, costumo medir PSNR e similaridade de cosseno (cossim).

A forma como você faz as versões do kernel realmente executarem numa GPU depende de a GPU estar disponível localmente ou na nuvem. Independentemente disso, nosso agente deve ter a capacidade de acessar suas saídas de alguma forma.

Considero a metodologia de rung abaixo uma boa forma de ter N funções de teste:

1def rung(name):
2    def deco(fn):
3        try:
4            out = fn()
5            results[name] = {"ok": True, **(out or {})}
6            print(f"[{name}] ok  " + "  ".join(f"{k}={v}" for k, v in (out or {}).items()))
7        except Exception as e:
8            results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}
9            print(f"[{name}] FAILED  {type(e).__name__}: {e}")
10            traceback.print_exc()
11        return fn
12    return deco

que você pode chamar assim:

1out = {}
2
3@rung("pre-checks")
4def _():
5    run_pure_checks()
6    run_dsl_checks()
7
8@rung("run")
9def _():
10    out["o"] = custom_kernel(*inputs)
11    torch.cuda.synchronize()
12    return {"shape": tuple(out["o"].shape),
13            "finite": bool(torch.isfinite(out["o"]).all())}

Para o rung de benchmark, você pode fazer várias coisas:

• Medir o tempo de execução do kernel de ponta a ponta, ou seja, o tempo total gasto
• Usar rastreamento intra-kernel para medir seções dentro de um kernel e despejar os resultados na saída (usando um tracer personalizado ou CUPTI)
• Despejar o IR, PTX, SASS e CUBIN gerados em um diretório de dumps e deixar o agente examiná-lo

O último ponto pode ser expandido um pouco mais. Às vezes, pode acontecer de o DSL fazer o lowering e gerar um PTX subótimo (e, eventualmente, SASS), e você encontra instruções ou shapes melhores que podem ser usados no lugar. Nosso agente pode ler os arquivos de texto PTX/SASS e embutir código de baixo nível diretamente, em vez de deixar o DSL lidar com a parte subótima. Novamente, passar a documentação "pesquisável" do PTX como contexto é muito útil aqui.

A última peça para unir tudo é o Profiling. Se seu agente puder acessar a CLI do NCU (Nsight Compute Systems), você pode pedir que ele faça o profiling e gere um relatório para o seu kernel como parte do loop de feedback verificável acima.

Considerações finais

"Então o desenvolvimento de kernels de GPU está morto?"

"Bem, sim, mas na verdade não"

Sim, porque a parte difícil de layouts, indexação, abstrações e da estrutura geral pode ser amplamente resolvida por agentes com contexto suficiente. Você pode facilmente reduzir o trabalho de 2 a 3 semanas para 1 a 2 dias. Não, porque o gargalo real mudou dos kernels para as validações. Não existe mais um único jeito de fazer as coisas: quanto melhor seu contexto e seu harness, mais rápido o processo. Casos especializados se beneficiam ainda mais, e um bom harness é tudo o que você vai precisar.

Por fim, em vez de tratar os agentes como autônomos, ainda é necessário tratá-los como assistentes realmente inteligentes que você pode orientar. É aqui que seu entendimento fundamental de GPUs e kernels é útil. A parte humana (você) ainda é necessária aqui.

É um sentimento agridoce, eu sei :)