摘要
机器学习模型在日常软件中越来越普遍，其运行时间大多花费在少量计算内核上，如矩阵乘法、卷积和归一化。优化这些内核是减少延迟和成本的最直接方式，但传统上需要专家工程师手动编写低级GPU代码。基于大语言模型（LLM）的代理系统能够以更少的人力生成和优化内核。然而，现有工具主要在随机生成的张量和孤立内核上进行评估，生成的CUDA代码需要开发者手动重新集成，且大多仅针对LLM PyTorch模型，支持的结果检查和调试功能有限。
我们提出Kernel Forge，一个开源的端到端智能代理工具，它可以无缝接入任何未修改的PyTorch模型。Kernel Forge支持视觉、扩散和LLM工作负载，利用蒙特卡罗树搜索（MCTS）探索多个优化路径，而不是单一线性优化链，并配备图形用户界面以监控进度、检查候选内核和调试失败。
我们在NVIDIA DGX Spark和GB10 GPU上评估了Kernel Forge，使用四个覆盖视觉、扩散和LLM工作负载的PyTorch模型。在每个内核仅进行50次优化迭代的情况下，Kernel Forge优化了14个内核，超越了PyTorch的急切模式，分别在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$。
博主点评： Kernel Forge的创新之处在于将LLM与MCTS结合，极大地简化了GPU内核的优化过程。它不仅提升了性能，还降低了对专家知识的依赖，为机器学习开发者提供了极大的便利，展现了未来软件优化的新方向。