KernelAgent turns a PyTorch forward pass into verified Triton kernels

The goal is a composed kernel, not a suggestion

KernelAgent's own headline is that it turns PyTorch programs into verified Triton kernels and optimizes their performance, and it was designed around KernelBench workloads.

What separates it from a code generating assistant is the last part of the sentence. The project describes five components, and the order matters: static problem analysis decides whether to run a lightweight path or a full pipeline, LLM assisted refactoring isolates fusable subgraphs, Triton kernels are generated in parallel with strict runtime verification, the original forward pass is composed end to end using only the synthesised kernels, and a hardware guided optimisation pipeline then iteratively improves performance.

That third and fourth stage are the difference. A system that emits a kernel and calls it a day leaves you with a snippet and a benchmark you have to trust. KernelAgent rebuilds the original forward pass out of the kernels it produced, and it verifies at runtime, which means a wrong kernel fails loudly rather than quietly producing wrong numbers in a training run.

The optimisation side is a separate pipeline with its own run directory, and its input is an existing Triton kernel. So the tool has two entry points: generate a kernel for a problem, or take a kernel you have and make it faster.

Both write their artefacts to disk under a run directory, .fuse/<run_id>/ for generation and .optimize/<run_id>/ for optimisation, and the final outputs are named: compose_out/composed_kernel.py in one case, output/best_kernel.py in the other.

A router reads your AST before anything else happens

The first stage is static analysis, and it is described in enough detail to be worth reading closely.

The AutoRouter, in Fuser/auto_agent.py, parses the problem's AST and looks for attention blocks, transposed convolutions, control flow and long op chains. It then decides between the direct KernelAgent path and the full Fuser pipeline. The reasoning is not documented in the readme, but the signals are: a problem with a clear fusable subgraph can go straight to kernel generation, while a complicated module with control flow or attention is better handled by the refactoring path.

Decisions are cached under .fuse/router_cache.json, and there are two flags to interfere with that. --no-router-cache avoids using or writing cached results, and --ignore-router-config ignores the tuning the router suggests and relies on the command line arguments instead.

The router also has a fallback: if the first path fails, it can fall back to the other one. That matters because the readme describes the auto-routing workflow as a problem that is picked up and run with automatic fallback if the first attempt fails.

There is a wrinkle for anyone using KernelBench as the source of examples. By default the KernelAgent UI looks for KernelBench at the same level as KernelAgent, that is in ../KernelBench, which the readme explains when it describes the optional clone step.

Extract, dispatch, compose, and the knobs for each

When you want explicit control over models and concurrency, the pipeline is run by hand, and it has three named stages.

  python -m Fuser.pipeline \
    --problem /abs/path/to/problem.py \
    --extract-model gpt-5 \
    --dispatch-model o4-mini \
    --dispatch-jobs auto \
    --compose-model o4-mini \
    --workers 4 \
    --max-iters 5 \
    --verify

Each stage can use a different model, and --dispatch-jobs auto matches the number of discovered subgraphs. The same command has an XPU variant with --target-platform xpu added.

What each stage does is documented at the component level. The Fuser orchestrator, in Fuser/orchestrator.py, rewrites the PyTorch module into fusable modules, executes them for validation, and packages a tarball of the fused code, with run IDs and directories managed through Fuser/paths.py.

The subgraph extractor, in Fuser/subgraph_extractor.py, prompts the LLM to emit a JSON array describing each unique subgraph, including ops, shapes, dtypes and parameter tensors. Entries are deduplicated by shape signature so the dispatcher can reuse a kernel rather than generating it again for a different name.

The dispatcher, in Fuser/dispatch_kernel_agent.py, converts each JSON item into a precise Triton generation spec and then spins up the workers.

The --verify flag is the one to leave on. It ensures the final composition test runs, and it is what turns a pile of generated kernels into something you can compare against the original.

The direct API and four UIs over the same machinery

Bypassing the Fuser entirely is a matter of importing one class and calling it, which is the quickest way to see whether the approach works for your problem at all.

  from triton_kernel_agent import TritonKernelAgent

  agent = TritonKernelAgent(num_workers=4, max_rounds=8, model_name="gpt-5")
  result = agent.generate_kernel(
      problem_description="Implement ReLU over a contiguous 1D tensor of length 1024"
  )

  if result["success"]:
      print("Kernel path:", result["kernel_path"])
      print("Session directory:", result["session_dir"])
  else:
      print("Failure:", result["message"])

The problem is described in plain language here, or it can be a KernelBench snippet, and the result carries the kernel path, the session directory and a failure message, so a run is inspectable whether it worked or not.

The same machinery is exposed through four console scripts in the project metadata: kernel-agent for the Triton KernelAgent UI, fuser-ui for the Fuser orchestration UI, pipeline-ui for the full pipeline, and optimization-ui for the optimisation side. Each has a script under scripts/ as an alternative entry point, and the interfaces are Gradio frontends, with gradio 5.5 or newer in the dependencies.

There is also a list-models script, which is a small thing but tells you something about how the project is used: the model names come from configuration rather than being hard coded, so a team can switch providers per stage.

Your hardware and your model provider are the two constraints

The requirements section is short and it is where the practical limits are.

The operating system is Linux or macOS, and the readme lists Python 3.8 to 3.12, while the project metadata pins requires-python to 3.10 up to but not including 3.13. The GPU has to be one of two things: an NVIDIA GPU with CUDA support, or an Intel GPU with oneAPI support, covering Arc, Data Center GPUs or integrated Xe graphics. Triton is installed separately, either with pip install triton or a nightly built from source, and PyTorch comes from its own install page. An LLM provider is required, and the readme lists OpenAI, Anthropic or a self hosted relay.

Configuration goes in an .env file that the CLIs load automatically, or in environment variables.

OPENAI_MODEL=gpt-5            # default model for extraction
NUM_KERNEL_SEEDS=4            # parallel workers per kernel
MAX_REFINEMENT_ROUNDS=10      # retry budget per worker
LOG_LEVEL=INFO                # logging level

The provider keys follow the same pattern, and the Anthropic comment records the default model, claude-sonnet-4-20250514, used when OPENAI_MODEL is unset. A relay configuration is available for self hosted gateways through LLM_RELAY_URL and LLM_RELAY_TIMEOUT_S, and the readme points at triton_kernel_agent/agent.py and Fuser/config.py for the remaining knobs.

For the Intel path there is a separate wheel index and a check you can run yourself, printing whether the XPU device is available and how many devices are present.

The optimisation pipeline is a separate door in

The performance side of the project is easy to underestimate, because it is not about generating a kernel from scratch but about improving one you already have.

The readme describes a multi-worker optimisation pipeline with its own artefacts: the input Triton kernel, the artefacts themselves, individual optimisation worker sessions, and the final output/best_kernel.py. The pipeline is hardware guided and iterative, meaning it measures on the machine and keeps what is faster.

The examples directory shows what that looks like in practice, with an optimise_01_matvec example, an optimise_02_rmsnorm example and an optimise_03_max_pooling example, plus a run_opt_manager.py driver, and separate examples for the generation side such as triton_01_element_add, triton_02_fused_reduction_gemm and triton_03_fused_dcpp.

This is the most useful entry point for most teams, and it is where the alternative approach sits. Writing a Triton kernel by hand and then letting this pipeline tune it is a different starting point from asking a model to invent one. The hand-written kernel is yours, you know it is correct, and the optimiser is improving something that already works. Generating from a problem description is the more impressive path and the more likely to need several attempts.

The trade is that the optimisation path needs a kernel to start from and a working benchmark, so it is not a way to discover whether a fused implementation is worth having at all.

Alpha software with no releases, and Apache-2.0

The status of this project is worth stating plainly before anyone plans around it. The project metadata marks it as Development Status :: 3 - Alpha, the version is 0.1.0, and the repository has no GitHub releases, so there is no tagged artefact to pin and no changelog to read.

The activity record is current enough to be encouraging rather than alarming. The repository is not archived and the last push was on 2026-09-08, and the project publishes blog posts on pytorch.org about the approach, one on autonomous GPU kernel generation through deep agents and one on hardware guided GPU kernel optimisation through multi agent orchestration.

The licence is Apache-2.0, recorded in the project metadata and in the LICENSE file at the repository root, which is the permissive choice with an explicit patent grant and reasonable for a research tool that others may build on.

The code layout shows how many moving parts this is. There is a Fuser directory for the orchestration pipeline, a triton_kernel_agent directory for the generation agent, a kernel_perf_agent directory for the optimisation side, an oink directory, a utils directory, a scripts directory for the UI entry points, a tests directory, an e2e_test.py at the root, and an examples directory with configuration files and the sample optimisations. The package metadata ships Jinja templates as package data for the triton_kernel_agent package, which is what the generation prompts are rendered from.

The honest summary is that this is a research project from the PyTorch Labs team that happens to have a working command line, and the right way to approach it is on a side problem, with verification on, rather than as the foundation of a production build.

Editorial conclusion

KernelAgent fits a PyTorch team with a specific hot op that they want rewritten as Triton and verified on real hardware, and that can point the pipeline at a KernelBench problem or their own module. It is the wrong tool for a general PyTorch optimisation, for a codebase with no clear fusable subgraph, and for anyone who needs reproducible build artefacts today, since the project is marked alpha and publishes no releases. Check first that you have a supported GPU, either CUDA or an Intel XPU, that your model provider and keys are configured, and that the composed kernel actually passes verification before you compare its speed to the original.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. meta-pytorch/KernelAgent on GitHub
4. README