Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels
06:00 · July 29, 2026 · arXiv cs.AI RSS
Machine learning models are increasingly embedded in everyday software, and most of their runtime is spent in a small set of compute kernels such as matrix multiplication, convolution, and normalization. Optimizing these kernels is one of the most direct ways to reduce latency and cost, but it has traditionally required expert engineers to hand-write low-level GPU code. Agentic systems built on large language models (LLMs) can now generate and optimize kernels with far less human effort, yet existing tools are largely evaluated on randomly generated tensors and isolated kernels, emit standalone CUDA code that developers must manually reintegrate, mostly target only LLM PyTorch models, and offer limited support for inspecting and debugging results. We present Kernel Forge, an open-source, end-to-end agentic harness that accepts any unmodified PyTorch model in place. Kernel Forge supports vision, diffusion, and LLM workloads, uses Monte Carlo Tree Search (MCTS) to explore multiple optimization paths rather than a single linear refinement chain, and ships with a graphical user interface for monitoring progress, inspecting candidate kernels, and debugging failures. We evaluate Kernel Forge on four PyTorch models spanning vision, diffusion, and LLM workloads on an NVIDIA DGX Spark with GB10 GPU. With only 50 optimization iterations per kernel, it optimizes 14 kernels to outperform PyTorch eager mode, reaching $1.52\times$ on adaptive\_avgpool2d in ResNet-50, $1.70\times$ on group\_norm in Stable Diffusion 3.5 Medium, $2.83\times$ on softmax in Gemma 4 E2B, and $1.54\times$ on softmax in Qwen 3.5 35B-A3B.
Summary
Kernel Forge is an open-source agentic system that applies large language models to generate and refine CUDA kernels for production PyTorch models. It accepts any unmodified model together with example inputs, runs inference to identify the operators that dominate runtime, and then launches an optimization loop that proposes, compiles, validates, and profiles specialized kernels. Unlike prior tools that operate on synthetic tensors or emit standalone code, Kernel Forge automatically inserts validated kernels back into the model’s execution path through guarded dispatch, preserving numerical correctness relative to PyTorch eager mode.
The system supports vision, diffusion, and language workloads rather than restricting itself to transformer-only cases. It replaces the linear refinement chains common in earlier LLM-based optimizers with Monte Carlo Tree Search, allowing the agent to explore multiple optimization trajectories and retain intermediate candidates that may later prove beneficial. A graphical interface exposes the captured operator graph, per-kernel timing breakdowns, revision trees, and failure diagnostics, while a command-line mode supports scripted use.
In experiments on an NVIDIA DGX Spark with GB10 GPU, Kernel Forge processed four representative models—ResNet-50, Stable Diffusion 3.5 Medium, Gemma 4 E2B, and Qwen 3.5 35B-A3B—applying fifty iterations to each target operator. Fourteen kernels exceeded eager-mode performance, delivering measured speedups of 1.52× for adaptive average pooling in ResNet-50, 1.70× for group normalization in Stable Diffusion, 2.83× for softmax in Gemma, and 1.54× for softmax in Qwen. The implementation is publicly available and designed for direct integration into existing PyTorch pipelines without manual code changes.
Why it matters
Actionable open-source tool for efficient AI deployment on GPUs; directly supports Dutch researchers optimizing real PyTorch workloads with novel MCTS search and reproducibility via code release.
