arXiv:2607.24762cs.AIcs.PF2026-07

用AI自动生成并优化GPU计算核函数,显著提升模型运行速度。

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

论文配图:Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels
图 1 · 摘自论文原文
  • 基于大模型和蒙特卡洛树搜索,自动探索多种优化路径。
  • 50次迭代内优化14个核心函数,最高提速2.83倍。
  • 支持视觉、生成、大模型多类任务,可视化界面便于调试。

机器学习模型越来越多地嵌入日常软件,其大部分运行时间集中在少数计算核函数上,如矩阵乘法、卷积和归一化。传统优化需专家手动编写低级GPU代码,成本高。如今基于大语言模型的智能体可大幅降低人工投入,但现有工具大多在随机张量和孤立核上评估,生成的CUDA代码需手动集成,主要针对LLM PyTorch模型,且缺乏结果分析与调试支持。我们提出Kernel Forge,一个开源的端到端智能体框架,可直接接收任意未修改的PyTorch模型。它支持视觉、扩散和大模型工作负载,采用蒙特卡洛树搜索(MCTS)探索多条优化路径而非单一线性优化链,并配备图形化界面用于监控进度、检查候选核函数及调试失败。我们在NVIDIA DGX Spark搭载GB10 GPU的设备上对四个跨视觉、扩散和大模型的PyTorch模型进行评估。每核仅50次优化迭代,即成功优化14个核函数,超越PyTorch急模式:ResNet-50中adaptive_avgpool2d提速1.52倍,Stable Diffusion 3.5 Medium中group_norm提速1.70倍,Gemma 4 E2B中softmax提速2.83倍,Qwen 3.5 35B-A3B中softmax提速1.54倍。

原文摘要 · Abstract (English)

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.

GPU优化大模型智能体CUDA

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。