OpenBMB/ForgeKernel

OpenBMB/ForgeKernel

Language: Python

License: Apache-2.0

Stars: 8

Forks: 0

Open issues: 0

Created: 2026-10-30T01:33:40Z

Pushed: 2026-10-30T02:50:23Z

Default branch: main

Fork: no

Archived: no

README:

ForgeKernel is a self-evolving GPU kernel agent that starts from a reference kernel. It reads the reference's architecture by running, timing and profiling it, turns that into an architecture-first milestone curriculum that a coding agent implements, and promotes every accepted kernel to the next reference, so optimization keeps going past the best known implementation. Findings certified by a hardware probe accumulate in a knowledge base shared across operators. In the paper, across 27 operators on Hopper and Blackwell, ForgeKernel matches or beats 18 of 22 expert-library kernels (FlashAttention-4, cuDNN, DeepGEMM and others) and all 7 published kernels of three kernel-generation systems, by up to +170.6% (+39.4% geometric mean). The kernels are released in [kernels/](kernels/).

ForgeKernel is part of the Forge family, together with ForgeTrain and ForgeStencil.

News

• 🏆 [2026-10] We reached #1 on the FlashInfer-Bench leaderboard of NVIDIA

SOL-ExecBench.

• 🏆 [2026-10] We reached #1 on Huawei CANN-Bench (Ascend A3).
• 🏆 [2026-10] We also reached #1 on Ascend 950DT on the

CANN-Bench leaderboard.

Highlights

• 🔁 Self-evolution: always one step past the reference. Every accepted kernel becomes the

next run's reference, so each generation has to beat the one before it, and the chain stops only when the hardware gives no more (a gain below the noise floor). In the paper every promoted reference was beaten again, by up to +9.9% a generation, including /goal agent kernels already far ahead of the expert library: GDN prefill ends +100% over FlashInfer, GQA attention forward +29% over FlashAttention-4 (scripts/self_evolve.py).

• 🗺️ Architecture-first curriculum: the error-prone decisions are settled first. The decisions

that are easiest to get wrong and most expensive to change later, the architecture (kernel split, tile shapes, warp roles, register and shared-memory budget) and the orchestration (what stays resident, the issue order, which work hides under which), are all made in the first milestone, M0, and frozen once its gate passes. Later milestones M1 → Mk only tune inside that structure, each closed by a gate of correctness, performance and mechanism lines. When measurement shows a limit only a new structure can lift, the arch agent goes back and revises M0 instead of letting later milestones work around it.

• 📚 A knowledge base that outlives the run. Hardware mechanisms certified by a GPU probe are

stored per GPU architecture (hw_kb/) and prescribed to later operators.

🤖 Set it up with a coding agent

Clone this repository, then paste this into Claude Code (or any coding agent that can run shell commands) started in the clone:

Read README.md in this directory and set up ForgeKernel on this machine:
install it in a virtual environment, write local.env for my GPU (this
machine, or the ssh host I name) and run `forge new tasks/rmsnorm --dry-run`.
Show me the plan it prints.

Quick Start

Requirements. The machine you run forge on: Linux or macOS, Python ≥ 3.11, git, and the Claude Code CLI (logged in, or an ANTHROPIC_API_KEY); tmux if a task has no reference. The GPU machine (the same one, or an ssh host): an NVIDIA Hopper (sm90) or Blackwell (sm100) GPU, python3 with a CUDA build of PyTorch and, for CuTe DSL tasks, nvidia-cutlass-dsl; Nsight Compute (ncu) is recommended, and without it a run proceeds on timings alone.

1. Install

cd ForgeKernel # a clone of this repository
python3 -m venv .venv && source .venv/bin/activate
pip install -e .
pip install nvidia-cutlass-dsl # recommended, see below
forge --version # forgekernel 0.1.0

The coding agent works offline; installing nvidia-cutlass-dsl next to ForgeKernel is how it reads the CuTe DSL API it writes against.

2. Say where the GPU is

cp local.env.example local.env
# FK_DELIVERY=local -> the GPU on this machine (FK_GPU picks the device)
# FK_DELIVERY=ssh -> REMOTE_HOST (an ssh alias that logs in without a
# prompt) and REMOTE_WORKDIR (a directory on it)

3. Run the example task: fuse a PyTorch RMSNorm (nine kernel launches a call) into one CuTe DSL kernel on a Blackwell GPU (set arch = "sm90" in task.toml for Hopper).

forge new tasks/rmsnorm --dry-run # check the task, print the plan
forge new tasks/rmsnorm # -> .artifacts/rmsnorm_
tmux new -s forge 'forge run .artifacts/rmsnorm_' # a run takes hours
forge status .artifacts/rmsnorm_ # progress
forge run .artifacts/rmsnorm_ # continue after an interruption

forge run exits 0 when the run reaches its target -- the loop stopped on it, or the final confirmation of the best kernel meets it -- and 2 otherwise.

Expected output. The run writes the correctness judge, calibrates it against the reference, measures the reference, distils a milestone, and iterates until the kernel is 1.2× the reference or the budget ends. It finishes with .artifacts/rmsnorm_/report.md and the best kernel in .artifacts/rmsnorm_/best/kernel.py. On a B200 our run of this example stopped at 5.2× the PyTorch reference after about 1.5 hours:

- task: rmsnorm (source mode)
- best: 77856a9d9561 -- device median 0.0307 ms, 5.2363x the reference
- kernel: best/kernel.py

> Full launch reference → [docs/design/v1-launch.md](docs/design/v1-launch.md)

How It Works

A run has two nested loops. In the inner loop the coding agent implements one decision row of the milestone at a time, and one GPU job judges correctness, times the candidate against the reference and profiles it; a correctness failure is reverted. In the outer loop the arch agent re-enters after iterations without a new best, at an exhausted decision graph or at a gate impasse: it prices levers, runs verdict probes and revises the curriculum. A shipped kernel can become the next run's reference, and what probes certify feeds the knowledge base.

Three artifacts carry the work,...