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

Kernel Forge：用于基于大模型生成与优化 CUDA 核函数的智能体框架

Abstract: 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.

摘要： 机器学习模型正日益嵌入到日常软件中，而其大部分运行时间都消耗在少数几个计算核函数（compute kernels）上，例如矩阵乘法、卷积和归一化。优化这些核函数是降低延迟和成本最直接的方法之一，但传统上这需要专家工程师手动编写底层的 GPU 代码。

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.

基于大语言模型（LLM）构建的智能体系统现在可以用远少于以往的人力来生成和优化核函数。然而，现有的工具大多是在随机生成的张量和孤立的核函数上进行评估，生成的独立 CUDA 代码需要开发者手动重新集成，且大多仅针对 LLM PyTorch 模型，在结果检查和调试方面的支持也十分有限。

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.

我们提出了 Kernel Forge，这是一个开源的端到端智能体框架，可以直接接收任何未经修改的 PyTorch 模型。Kernel Forge 支持视觉、扩散模型和 LLM 工作负载，利用蒙特卡洛树搜索（MCTS）来探索多种优化路径，而非单一的线性改进链；此外，它还配备了图形用户界面，用于监控进度、检查候选核函数以及调试故障。

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.

我们在配备 GB10 GPU 的 NVIDIA DGX Spark 上，针对涵盖视觉、扩散模型和 LLM 工作负载的四个 PyTorch 模型对 Kernel Forge 进行了评估。每个核函数仅经过 50 次优化迭代，它就成功优化了 14 个核函数，使其性能超越了 PyTorch 的即时执行模式（eager mode）：在 ResNet-50 的 adaptive_avgpool2d 上实现了 $1.52\times$ 的加速，在 Stable Diffusion 3.5 Medium 的 group_norm 上实现了 $1.70\times$ 的加速，在 Gemma 4 E2B 的 softmax 上实现了 $2.83\times$ 的加速，以及在 Qwen 3.5 35B-A3B 的 softmax 上实现了 $1.54\times$ 的加速。