Optimising the compute kernels that dominate neural network inference has traditionally been the preserve of specialist GPU engineers. Writing efficient CUDA for convolutions, matrix multiplications, or normalisation layers requires deep knowledge of memory hierarchies, warp scheduling, and hardware-specific intrinsics. The promise of large language models as code generators raises an obvious question: can an agentic system replace or at least augment that expertise? Kernel Forge is a serious attempt to answer that question in a deployment-relevant setting, rather than on the synthetic benchmarks that have dominated prior work.
What Problem Does This Actually Solve?
Existing LLM-based kernel optimisation tools share a cluster of related weaknesses. They typically evaluate generated kernels against randomly constructed tensors rather than operators captured from real model executions. They produce standalone CUDA files that developers must manually reintegrate. Most focus narrowly on LLM serving workloads, ignoring vision or diffusion models. And they tend to use linear or beam-style refinement loops, which can commit early to suboptimal design choices and fail to recover.
Kernel Forge addresses all four of these gaps simultaneously. It accepts any unmodified PyTorch model, runs it on user-provided inputs, captures the operators that actually execute during inference, generates and validates CUDA replacements, and automatically reintegrates passing kernels into the model's execution path. The system supports vision, diffusion, and LLM workloads, and uses Monte Carlo Tree Search to explore the optimisation space rather than a single refinement chain. A graphical interface allows users to inspect candidate kernels, review MCTS revision trees, and diagnose failures without touching the command line.
Architecture and Methodology
The pipeline has two main stages. First, a model ingestion component runs the target model and records every supported operator invocation, capturing tensor shapes, dtypes, strides, scalar arguments, device placement, and reference outputs. Critically, calls to the same operator are grouped into variants only when they share concrete runtime characteristics. Two calls to conv2d with different shapes become separate variants, each carrying its own validation examples and measured eager latency. This granularity matters because kernel performance is highly sensitive to tensor dimensions and memory layout.
Each variant is packaged as an operator card that records the call count and runtime share within the captured operator region. This weighting is one of the more methodologically careful choices in the paper: a 3x speedup on an operator consuming 0.2% of runtime is worth far less than a 1.1x improvement on one consuming 50%. The optimisation pipeline is therefore guided by actual runtime responsibility, not just peak achievable speedup.
The optimisation loop itself uses MCTS to manage the search over candidate kernels. Unlike beam search, MCTS can revisit branches that appeared unpromising early, which is relevant here because a temporarily slower intermediate kernel might be a stepping stone to a faster final implementation. The agent uses Anthropic Claude Opus 4.7 for code generation and revision. Each candidate is compiled, numerically validated against the PyTorch eager reference outputs, and profiled. A guarded fallback policy ensures that if a generated kernel does not beat the eager baseline, the original PyTorch path is retained.
Results Across Vision, Diffusion, and LLM Workloads
Experiments run on an NVIDIA DGX Spark with a GB10 GPU (compute capability 12.1) across four models: ResNet-50, Stable Diffusion 3.5 Medium, Gemma 4 E2B, and Qwen 3.5 35B-A3B. Each kernel receives a budget of 50 optimisation iterations. The headline results are:
- ResNet-50: adaptive_avgpool2d reaches 1.52x over PyTorch eager; tensor add and max pooling also improve at 1.26x and 1.21x respectively.
- Stable Diffusion 3.5 Medium: group_norm reaches 1.70x; layer_norm and SiLU show modest gains around 1.05x.
- Gemma 4 E2B: softmax reaches 2.83x, the largest single-operator improvement in the evaluation.
- Qwen 3.5 35B-A3B: softmax reaches 1.54x; conv1d and embedding also improve.
Across all four models, 13 of 24 open-source or native PyTorch operators are faster than eager at opt50, compared with 1 of 9 proprietary or vendor-backed operators. This asymmetry is not surprising. Operators like linear and scaled-dot-product attention are dispatched through cuBLAS, cuDNN, and FlashAttention backends that represent years of engineering effort. Generated CUDA competing against those paths faces a much harder baseline than competing against a native PyTorch implementation.
The guarded fallback policy is therefore doing real work. In Stable Diffusion, linear and scaled-dot-product attention together account for over 80% of the captured operator region, and both fall well below eager (0.495x and 0.021x respectively). Without fallback, deploying these generated kernels would be catastrophic. With it, the system retains the PyTorch eager path for those operators and applies generated kernels only where they demonstrably improve latency.
Limitations and Open Questions
The paper is commendably honest about the gap between operator-level speedup and end-to-end model speedup. The operators where Kernel Forge succeeds most dramatically, such as softmax in Gemma 4 at 2.83x, account for a small fraction of total runtime (5.93% in that case). The dominant operators remain resistant to the approach at the current iteration budget. Whether increasing the budget or using more capable future models would close that gap is an open question the paper does not address.
The evaluation is also single-hardware. All experiments run on the GB10 GPU, which has compute capability 12.1. CUDA kernels are notoriously hardware-specific, and a kernel optimised for one architecture can regress on another. The paper does not report portability experiments, so it is unclear how generated kernels would behave on older Ampere or Hopper hardware that represents a much larger installed base.
The cost accounting section (Figures 7 and 8 in the paper) tracks API spend per optimisation arm, which is useful for practical adoption decisions. However, the paper does not report wall-clock time for the full optimisation pipeline per model, which would be equally relevant for practitioners deciding whether to run Kernel Forge as part of a deployment workflow.
Finally, the MCTS controller is presented as a design choice to avoid the pitfalls of linear refinement and beam search, but the paper explicitly states it does not claim search-policy superiority. A controlled ablation comparing MCTS against a simple iterative refinement baseline would have strengthened the architectural argument considerably.
Despite these gaps, Kernel Forge represents a meaningful step toward deployment-aware agentic kernel optimisation. The insistence on evaluating within real model executions, weighting results by runtime contribution, and accounting for fallback behaviour sets a higher methodological bar than most prior work in this space. The open-source release and GUI lower the barrier for researchers and practitioners who want to apply these techniques without specialist CUDA knowledge. The full paper is available at arXiv:2607.24762.