RightNow-AI/autokernel: pesquisa automática para kernels de GPU. Dê a ele qualquer modelo PyTorch, durma e acorde com kernels Triton otimizados. · Github

Pesquisa automática para kernels de GPU. Dê a ele qualquer modelo PyTorch, durma e acorde com kernels Triton otimizados.

Inspirado por @karpathy/autoresearch – que demonstrou agentes autônomos de IA para pesquisa de treinamento LLM. O AutoKernel aplica a mesma filosofia à otimização do kernel da GPU: o agente modifica um arquivo, executa uma avaliação fixa, mantém ou reverte, repete para sempre.

Dê ao AutoKernel qualquer modelo PyTorch. Ele vai:

1. Perfil o modelo para descobrir quais kernels de GPU são gargalos
2. Extrair cada gargalo como um kernel Triton independente
3. Otimizar cada kernel de forma autônoma (editar, comparar, manter/reverter – para sempre)
4. Verificar correção de ponta a ponta e relatar a aceleração total

O agente lê program.md – o “código organizacional de pesquisa” – que contém instruções abrangentes para operação autônoma. Ele edita kernel.py um kernel por vez, executa bench.py (benchmark fixo com verificações de correção em 5 estágios + análise da linha do telhado) e mantém ou reverte a alteração. O orquestrador decide quando passar para o próximo kernel usando a lei de Amdahl.

Cada experimento leva cerca de 90 segundos. São aproximadamente 40 experimentos/hora, aproximadamente 320 durante a noite, em todos os kernels.

Requisitos: GPU NVIDIA (testada em H100/A100/RTX 4090), Python 3.10+, uv.

# Install uv (if you don't have it)
curl -LsSf https://astral.sh/uv/install.sh | sh

# Clone and setup
git clone https://github.com/RightNow-AI/autokernel.git
cd autokernel
uv sync

# One-time setup: test data + baselines
uv run prepare.py

# Profile a model (ships with GPT-2, LLaMA, BERT -- no transformers needed)
uv run profile.py --model models/llama_7b.py --class-name LlamaModel \
 --input-shape 1,512 --dtype float16

# Extract top bottleneck kernels
uv run extract.py --top 5

# Verify benchmark works
uv run bench.py

Abra Claude, Codex ou qualquer agente de codificação neste diretório:

Read program.md and let's kick off a new experiment. Start with setup.

O agente irá:

1. Crie o perfil do seu modelo e apresente o plano de otimização
2. Crie uma ramificação (por exemplo, autokernel/mar10-llama7b)
3. Otimize cada kernel de gargalo em ordem de prioridade
4. Verifique a correção de ponta a ponta e relate a aceleração total

program.md é intencionalmente abrangente para que o agente possa funcionar por mais de 10 horas sem ficar preso. Inclui um manual de otimização de 6 níveis, estrutura de decisão, tratamento de falhas e raciocínio da lei de Amdahl.

                 profile.py              extract.py           bench.py (loop)         verify.py
Any PyTorch  ──>  Rank kernels  ──>  Generate baseline  ──>  Optimize each  ──>  End-to-end
   model          by GPU time       Triton kernels          kernel (agent)       verification

9 tipos de kernel que cobrem as operações principais do aprendizado profundo moderno:

Cada um tem uma referência PyTorch em reference.py e um kernel Triton inicial em kernels/.

Definições de modelo independentes são fornecidas com o AutoKernel (sem transformers biblioteca necessária):

Para modelos HuggingFace (uv sync --extra models):

uv run profile.py --module transformers --class-name AutoModelForCausalLM \
 --pretrained meta-llama/Llama-2-7b-hf --input-shape 1,2048 --dtype float16

autokernel/
  kernel.py             the file the agent modifies (one kernel at a time)
  program.md            agent instructions -- the "research org code"

  bench.py              fixed benchmark + 5-stage correctness harness
  reference.py          PyTorch reference implementations (ground truth)
  prepare.py            one-time setup: test data, baselines

  profile.py            profile any PyTorch model, rank kernels by GPU time
  extract.py            extract bottleneck kernels into workspace/
  orchestrate.py        multi-kernel scheduler (Amdahl's law)
  verify.py             end-to-end model verification + speedup report
  analysis.py           experiment visualization (generates progress.png)

  kernels/              starter Triton kernels (9 types)
  models/               self-contained model definitions (GPT-2, LLaMA, BERT)
  workspace/            runtime artifacts (gitignored)

Por que Tritão. Sintaxe legível semelhante a Python que o agente pode entender e modificar sem dominar PTX ou SASS inline. O Triton bem ajustado atinge regularmente 80-95% do cuBLAS. O agente precisa iterar rapidamente – o Triton compila em segundos, não em minutos.

Correção primeiro. O benchmark verifica a saída do kernel em relação ao PyTorch antes de medir o desempenho. Um kernel rápido, mas errado, é imediatamente revertido. Isso evita que o agente “otimize” produzindo lixo.

Orquestração da lei de Amdahl. O orquestrador prioriza por impacto. Uma aceleração de 1,5x em um kernel de 60% (1,25x de ponta a ponta) supera uma aceleração de 3x em um kernel de 5% (1,03x de ponta a ponta). Ele segue em frente quando os retornos decrescentes se instalam.

Arquivo único para modificar. O agente apenas toca kernel.py. O escopo permanece gerenciável, as diferenças podem ser revisadas e revertem limpo.

Registro de TSV. Os resultados vão para uma planície results.tsv arquivo. Legível por humanos, compatível com git, analisável trivialmente, sem infraestrutura.

Cada experimento é registrado em results.tsv (separado por tabulação):

Este projeto é pesquisa automática para kernels de GPU – diretamente inspirado na pesquisa automática de Andrej Karpathy, o experimento original em agentes autônomos de pesquisa de IA para treinamento LLM. Karpathy mostrou que um agente de IA pode realizar centenas de experimentos durante a noite, explorando metodicamente um espaço de busca e registrando cada resultado. O AutoKernel aplica o mesmo loop – o agente edita um arquivo, executa uma avaliação fixa, mantém ou reverte – ao domínio de otimização do kernel da GPU com Triton.

Construído pela equipe por trás do Forge.

COM