Kernel Forge：用LLM智能体自动生成优化GPU内核

在深度学习模型日益普及的今天，大部分推理和训练时间都集中在少数几个计算内核上，例如矩阵乘法、卷积和归一化操作。优化这些内核是降低延迟和成本最直接的方式之一，但传统上需要专家工程师手动编写低层GPU代码，过程繁琐且耗时。近年来，基于大语言模型（LLM）的智能体系统展现出自动生成和优化内核的潜力，然而现有工具大多存在局限性：它们通常在随机生成的张量和孤立内核上评估，输出独立的CUDA代码需要开发者手动集成，主要针对LLM PyTorch模型，且对结果检查和调试支持有限。

针对这些问题，研究人员提出了Kernel Forge——一个开源的端到端智能体框架。Kernel Forge的核心创新在于它能够接受任意未经修改的PyTorch模型作为输入，无需开发者事先进行代码适配或模型转换。该框架支持视觉、扩散和LLM等多种工作负载，覆盖了当前主流AI应用场景。其优化策略摒弃了传统的单线程线性改进链，转而采用蒙特卡洛树搜索（MCTS）来并行探索多个优化路径。MCTS是一种在决策空间中高效搜索最优解的算法，通过模拟、扩展、回溯等步骤平衡探索与利用，使得Kernel Forge能够在更广阔的优化空间中寻找更优的内核实现。

Kernel Forge还配备了一个图形用户界面，用于实时监控优化进度、检查候选内核代码以及调试失败案例。这一设计大大降低了使用门槛，让非GPU编程专家也能参与到内核优化工作中。开发者可以直观地看到每个优化步骤的效果，对比不同候选内核的性能差异，并快速定位问题所在。

在性能评估方面，研究团队在NVIDIA DGX Spark（搭载GB10 GPU）上对四个PyTorch模型进行了测试，涵盖视觉、扩散和LLM三类工作负载。每个内核仅经过50次优化迭代，Kernel Forge便成功优化了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在多种模型架构上的通用性和高效性。

Kernel Forge的开源性质意味着社区可以自由访问其代码、复现实验结果并在此基础上进行二次开发。这为AI基础设施的自动化优化开辟了新路径，有望将GPU编程从少数专家的专属技能转变为可由AI智能体辅助完成的标准化任务。未来，随着LLM能力的进一步提升和搜索算法的改进，类似Kernel Forge的工具可能会在更广泛的硬件平台上实现自动内核优化，从而推动整个AI生态系统的效率提升。