Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels

cuda kernel optimizationdeep learning workload optimizationgpu programming automationllm agentmonte carlo tree searchpytorch model acceleration

Overview


Machine learning models are increasingly embedded in everyday software, and the majority of their runtime is concentrated 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 traditionally this has 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 significantly less human effort. However, existing tools are largely evaluated on randomly generated tensors and isolated kernels. They emit standalone CUDA code that developers must manually reintegrate, primarily target only LLM PyTorch models, and offer limited support for inspecting and debugging results.


Introducing Kernel Forge


We present Kernel Forge, an open-source, end-to-end agentic harness that accepts any unmodified PyTorch model in place. Kernel Forge is designed to:

  • Support vision, diffusion, and LLM workloads
  • Use Monte Carlo Tree Search (MCTS) to explore multiple optimization paths rather than a single linear refinement chain
  • Ship with a graphical user interface for monitoring progress, inspecting candidate kernels, and debugging failures

Evaluation Results


We evaluated Kernel Forge on four PyTorch models spanning vision, diffusion, and LLM workloads, running on an NVIDIA DGX Spark with GB10 GPU. With only 50 optimization iterations per kernel, Kernel Forge achieved the following speedups over PyTorch eager mode:


  • ResNet-50: 1.52× on adaptive_avgpool2d
  • Stable Diffusion 3.5 Medium: 1.70× on group_norm
  • Gemma 4 E2B: 2.83× on softmax
  • Qwen 3.5 35B-A3B: 1.54× on softmax

Availability


The source code is publicly available at: https://github.com/TheJoshBrod/KernelForge


Subject Areas


Primary: Artificial Intelligence (cs.AI)

Secondary: Performance (cs.PF)


Citation


Joshua Brodsky, Dhravid Kumar, Savini Kashmira, Jayanaka Danatanarayana, Jason Mars, Krisztian Flautner, and Lingjia Tang. "Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels." arXiv:2607.24762 [cs.AI], June 2026.




Submitted on 2 June 2026

via ArXiv AI

Related