Kernel Forge Uses LLM Agents to Optimize CUDA Kernels in PyTorch Models

Most LLM code-optimization loops work by asking the model to revise a single attempt over and over. Kernel Forge instead uses Monte Carlo Tree Search to branch across alternative optimization routes. This design choice reflects the actual shape of the problem. Kernel tuning is full of dead ends: an implementation can be numerically correct but slower than the baseline, or fast but subtly wrong on edge shapes. A linear refinement chain tends to get trapped polishing one such branch. Tree search lets the harness abandon a promising-but-invalid path and keep exploring, which matters when the space of valid, fast implementations for an operator is large and non-obvious.

The workloads where this pays off are predictable. A small set of operators, matrix multiplication, convolution, normalization and softmax, absorbs most of the runtime in modern models, and the same operators are hit repeatedly during inference. Shaving latency off them compounds. Historically that work has demanded engineers fluent in GPU programming, a scarce and heavy skill. Kernel Forge's pitch is that an agent can cover routine per-model bottlenecks without pulling those engineers onto every workload, and without waiting for a framework upgrade to deliver the improvement.

The evidence is concrete but narrow. The authors ran Kernel Forge on an NVIDIA DGX Spark with a GB10 GPU, allowing 50 optimization iterations per kernel. The reported wins are specific: 1.52 times on adaptive average pooling in ResNet-50, 1.70 times on group normalization 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. That the harness improved kernels across three different workload families, rather than only LLMs, is the strongest part of the result, because it directly answers the criticism the authors level at earlier tools.

Two caveats sit inside those same numbers. First, four models on one hardware configuration is enough to demonstrate that the approach generalizes past a single category, but not enough to claim broad coverage across PyTorch models, GPU generations, or precision regimes, all of which CUDA performance is acutely sensitive to. Second, and more consequential, the baseline is PyTorch eager mode. Eager execution is the least optimized way most production systems run a model. Real serving stacks lean on graph compilers, vendor libraries such as cuDNN, framework-level fusion, or bespoke inference runtimes, and the paper does not establish how Kernel Forge's kernels stack up against those. A 2.83 times gain over eager softmax is a meaningfully different figure than a gain over a compiled or vendor-tuned baseline.

This is a research paper, and open source does not mean production-ready. Automatically generated CUDA introduces the usual reliability burden: teams would need strong correctness tests, numerical tolerance checks, fallback paths and human review before any generated kernel touches a customer-facing system. The paper's emphasis on a GUI for inspection and debugging is a tacit acknowledgment of exactly this, generated GPU code still needs validation and failure analysis, not blind trust.

The useful next evidence is straightforward to name. Independent reproduction across more models and more GPUs would test generality. Comparisons against compiler-optimized PyTorch and specialized inference runtimes would test whether the wins survive a serious baseline. And measuring end-to-end application latency, rather than per-kernel speedup, would show whether the optimized operators actually move the number that users feel. If Kernel Forge can slot into continuous performance testing, it could help teams catch regressions and maintain workload-specific optimizations over time. For now it is a clear signal that agent harnesses, not just standalone code generation, are pushing into AI infrastructure.

Source: arXiv paper.