Kernels de GPU para Vibecoding

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

Sinceramente, escrever kernels de GPU é uma tarefa embaraçosamente verificável.

Primeiro, você decide para quais operações deseja escrever o kernel. Em segundo lugar, 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 bater, você tenta corrigir o problema de corretude. Feito isso, você faz benchmark do tempo de execução do kernel. As versões seguintes se baseiam nisso, e você continua otimizando até atingir as métricas de roofline ou ficar satisfeito.

Loop verificável: desenvolvimento de kernels 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 loop central de recompensa verificável. É essa a 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 ao menos começar a escrever uma boa primeira versão.

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

LLMs modernos já sabem escrever código Triton "bom". Eles conseguem trabalhar 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 buscar compreensão das abstrações do DSL ajuda muito.

Por exemplo, clonar o repositório

NVIDIA cutlass

no diretório de contexto é uma boa maneira de permitir que o agente busque 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, sobre o qual falarei a seguir.

Teste, Benchmark e Profiling

Com contexto suficiente para o agente, o verdadeiro gargalo agora muda para a validação. A própria implementação de referência — e a validação em relação a ela — torna-se cada vez mais importante. Chamo esta fase de teste de corretude ou simplesmente teste. A velocidade da implementação de referência não importa tanto quanto o seu propósito. Aquilo que você pretende medir e verificar é o que seu agente vai otimizar.

Normalmente, quando o cálculo não deve ser feito em precisão reduzida, eu 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 maneira de fazer as versões do kernel realmente rodarem em uma GPU depende de a GPU estar disponível localmente ou na nuvem. Independentemente disso, nosso agente deve ter a capacidade de acessar seus resultados de uma forma ou de outra.

Considero a metodologia rung abaixo uma boa maneira 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:

• Fazer benchmark do tempo de execução do kernel de ponta a ponta, para medir o tempo total gasto
• Usar tracing intra-kernel para fazer benchmark de seções internas de um kernel e despejá-las 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 analisar o conteúdo

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

A última peça para unir tudo é o profiling. Se o seu agente tiver acesso à CLI do NCU (Nsight Compute Systems), você pode pedir que ele faça o profiling e gere um relatório do seu kernel como parte do loop de feedback verificável descrito 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 verdadeiro gargalo mudou dos kernels para as validações. Não existe mais um único jeito de fazer as coisas: quanto melhor o seu contexto e o 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 o seu conhecimento fundamental de GPUs e kernels se mostra útil. A parte humana (você) ainda é necessária aqui.

É um sentimento agridoce, eu sei :)