Kernel Forge：用于基于 LLM 的 CUDA 内核生成和优化的代理工具

arXiv:2607.24762v1 公告类型：新 摘要：机器学习模型越来越多地嵌入日常软件中，其大部分运行时间都花在一小组计算内核上，例如矩阵乘法、卷积和归一化。优化这些内核是减少延迟和成本的最直接方法之一，但传统上需要专家工程师手动编写低级 GPU 代码。基于大型语言模型 (LLM) 构建的代理系统现在可以用更少的人力生成和优化内核，但现有工具主要是在随机生成的张量和隔离内核上进行评估，发出开发人员必须手动重新集成的独立 CUDA 代码，大多仅针对 LLM PyTorch 模型，并且对检查和调试结果提供有限的支持。我们推出 Kernel Forge，这是一种开源的端到端代理工具，可以接受任何未经修改的 PyTorch 模型。 Kernel Forge 支持视觉、扩散和 LLM 工作负载，使用蒙特卡罗树搜索 (MCTS) 来探索多个优化路径而不是单个线性细化链，并附带图形用户界面，用于监控进度、检查候选内核和调试故障。我们在四个 PyTorch 模型上评估 Kernel Forge，这些模型涵盖视觉、扩散和配备 GB10 GPU 的 NVIDIA DGX Spark 上的 LLM 工作负载。每个内核仅进行 50 次优化迭代，它优化了 14 个内核，其性能优于 PyTorch eager 模式，在 ResNet-50 中的 Adaptive\_avgpool2d 上达到 $1.52\times$，在 Stable Diffusion 3.5 Medium 中的 group\_norm 上达到 $1.70\times$，在 Gemma 4 E2B 中的 softmax 上达到 $2.83\times$，在 Gemma 4 E2B 中的 softmax 上达到 $1.54\times$。 Qwen 3.5 35B-A3B 中的 softmax。