Kernel Forge：用 LLM Agent 自动优化 CUDA 核函数

Kernel Forge 是一个开源的端到端 Agent 框架，用大语言模型为 PyTorch 模型自动生成并优化 CUDA 核函数。论文于 arXiv 发布，提出用蒙特卡洛树搜索（MCTS）取代传统的单一链式迭代来探索优化空间，并附带可视化界面用于监控、检查与调试。

研究背景

机器学习模型被广泛部署后，矩阵乘法、卷积、归一化等少量计算核消耗了绝大部分运行时间。手工优化这些 CUDA 核能直接降低延迟与成本，但传统上依赖专业工程师编写底层 GPU 代码。基于 LLM 的 Agent 系统已开始涉足这一领域，但现有工具普遍存在三类局限：在随机生成的张量与孤立核上评估、输出独立 CUDA 代码需开发者手动整合、仅支持 LLM 类的 PyTorch 模型、缺少结果检查与调试手段。

Kernel Forge 的核心设计

Kernel Forge 面向未修改的 PyTorch 模型直接输入，覆盖视觉、扩散模型和大语言模型三类负载。其主要设计包含三方面：

• 用蒙特卡洛树搜索（MCTS）并行探索多条优化路径，而非线性迭代链；
• 输出可直接嵌入现有 PyTorch 代码的核函数；
• 提供图形用户界面，用于实时监控、检查候选核函数并调试失败案例。

评测结果

研究团队在搭载 GB10 GPU 的 NVIDIA DGX Spark 上对四个 PyTorch 模型做了测试，每个核函数分配 50 次优化迭代。最终共 14 个核函数的性能超过 PyTorch 的 eager 模式，关键数据如下：

• 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 倍加速。

意义与局限

Kernel Forge 把 Agent 的能力从代码生成扩展到了核函数级别的性能调优，为自动化 CUDA 优化提供了相对完整的工作流。但论文评测规模有限：仅在单一硬件平台、单一 GPU 上进行，优化迭代次数固定为 50 次，工业级部署所需的更大规模与多硬件平台验证仍有待后续工作补充。