KForge：面向AI加速器的LLM驱动跨平台内核生成

# KForge: 面向AI加速器的大语言模型驱动跨平台内核生成

来源：https://arxiv.org/html/2606.02963  
旧金山，加利福尼亚州，美国 \{taras (https://arxiv.org/html/2606.02963v1/mailto:[email protected]), bbartan (https://arxiv.org/html/2606.02963v1/mailto:[email protected]), ankitanayak (https://arxiv.org/html/2606.02963v1/[email protected]), tstjohn (https://arxiv.org/html/2606.02963v1/mailto:[email protected]), nserrino (https://arxiv.org/html/2606.02963v1/mailto:[email protected]), zasgar (https://arxiv.org/html/2606.02963v1/mailto:[email protected])\}@gimletlabs\.ai

###### 摘要

生产环境中的推理任务日益面向异构加速器组合。智能体式流水线交织了推理、工具调用和多智能体协调，每个环节都拥有独特的计算和内存特征。为了达到最优效率，每个阶段都应当在最适合它的加速器上运行。这带来了一个系统挑战：每条流水线现在都需要在日益增多的硬件后端和编程模型上实现高性能内核。手工编写这些内核耗时费力，需要深入的底层专业知识，且随着内核复杂性增长而无法扩展。最近，大语言模型（LLMs）已被用于自动内核生成，但在低层级代码生成和跨后端泛化方面仍存在挑战。我们提出KForge，这是一个跨平台框架，构建于由两个基于LLM的协作智能体驱动的迭代优化循环之上：一个生成智能体，它利用编译和正确性反馈生成并逐步优化内核；一个性能分析智能体，它解释从程序化API到基于GUI的工具产生的性能分析数据，并输出引导下一轮合成的建议。该循环在功能性优化轮次（驱动候选内核达到正确性）和性能优化轮次（缩小与手工调优基准之间的性能差距）之间交替进行。我们在两个基线参考可用性截然不同的后端上评估KForge。在NVIDIA B200上，在gpt-oss-20b推理速度基准测试中，KForge相比TensorRT-LLM实现了端到端吞吐量提升2.12%。在Intel Arc B580上，KForge生成的Triton内核在KernelBench Level 2的37个GEMM+尾运算任务上，相比PyTorch即时编译和`torch.compile`中较快者实现了5.13×几何平均加速，这主要得益于算子融合和混合精度执行。

参见图注  
图1：利用LLMs的迭代式程序合成与优化循环。工作流程包括两个主要阶段：(1) 功能性优化轮次，迭代优化合成程序直到代码编译、无错误执行并产生正确输出；(2) 性能优化轮次，向LLM提供性能反馈以实现迭代性能改进。

## I 引言

智能体式AI应用的兴起正在重塑推理工作负载的部署方式。单个流水线现在交织了LLM调用与推理、数据检索、工具调用和多智能体协调，每个环节具有不同的计算、内存和同步特征。没有单一的加速器能同时适合所有这些阶段。因此生产部署需要将每个阶段分配给最适合的加速器以获得最佳性能。这种异构性带来了相应的系统设计挑战[3 (https://arxiv.org/html/2606.02963#bib.bib51)]。

在异构平台上实现这种映射需要在多种设备和编程模型上对相同核心算子进行优化实现。编写高性能计算内核需要掌握并行编程语言，如CUDA[15 (https://arxiv.org/html/2606.02963#bib.bib19)]、Metal[2 (https://arxiv.org/html/2606.02963#bib.bib21)]、Triton[19 (https://arxiv.org/html/2606.02963#bib.bib33)]、SYCL[9 (https://arxiv.org/html/2606.02963#bib.bib46)]或CuTe DSL[14 (https://arxiv.org/html/2606.02963#bib.bib45)]。跨加速器移植内核很少是简单的语法翻译，因为每个平台暴露不同的计算能力、内存层次结构、带宽限制和通信开销，因此最佳实现通常因平台而异。

现代编译器和推理运行时在自动化性能优化方面取得了重大进展。像 `torch.compile`[1 (https://arxiv.org/html/2606.02963#bib.bib1)] 和 TensorRT-LLM[16 (https://arxiv.org/html/2606.02963#bib.bib10)] 这样的系统通过图优化和自动内核融合改善了神经网络执行。尽管如此，构建高性能内核需要将巧妙的算法技术与谨慎的硬件利用结合起来，如FlashAttention[5 (https://arxiv.org/html/2606.02963#bib.bib27),6 (https://arxiv.org/html/2606.02963#bib.bib28)] 所示，其中将在线softmax[12 (https://arxiv.org/html/2606.02963#bib.bib3)]与分块注意力计算及硬件特定指令结合可实现卓越性能。这些优化共同减少了内存流量，提高了算术强度，并避免了通用编译器优化可能无法消除的调度开销。端到端模型推理进一步复杂化了问题。一个独立运行更快的内核，如果它阻止了图级别优化、引入了布局或精度转换、或将瓶颈转移到流水线其他位置，可能无法改善整个模型的运行时间。相反，一个在独立基准测试中表现平平的内核，如果与有利的周围操作组合，可能会释放更大的增益。内核优化必须在局部（单个内核层面）和全局（全模型执行上下文）两个层面进行评估。

大语言模型（LLMs）的最新进展使代码生成变得越来越可行。然而，将这些能力扩展到内核生成仍然是一个挑战。底层性能代码通常很脆弱：内存布局或数值精度的微小变化可能破坏正确性或消除性能收益。此外，LLMs的训练数据偏向于CUDA，导致新兴平台代表性不足。

在这项工作中，我们提出了KForge，一个用于LLM驱动内核生成和优化的跨平台框架，图1 (https://arxiv.org/html/2606.02963#S0.F1) 描述了其架构。KForge通过迭代优化循环模拟内核工程师的真实工作流程。系统在功能性优化轮次（驱动候选内核走向正确性）和性能优化轮次（一旦达到正确实现后改善运行时间）之间交替进行。在功能性优化轮次中，生成智能体提出候选内核，并利用编译和正确性反馈收敛到一个有效实现。在性能优化轮次中，性能分析智能体解释性能数据，包括程序化指标和基于GUI的分析器输出，并产生具体的代码级建议。这些建议针对硬件利用率指标，如内存带宽利用率、线程束占用率和算术强度，并输入到下一轮合成中。

KForge的设计目标是能够在不同硬件后端和编程模型上生成内核。我们在各种后端和编程模型上评估相同的智能体架构，从而可以研究内核生成智能体何时能够跨平台传递算法思想，以及何时需要后端特定的指导以利用硬件特定能力。我们通过两个案例研究来评估KForge，这两个案例选择了基线可用性截然不同的场景：NVIDIA B200，这里KForge要与TensorRT-LLM中多年供应商调优的内核进行竞争；以及Intel Arc B580，这里没有可比的手工调优基线，系统反而作为新兴硬件的上线工具。总结起来，本文做出了以下**关键贡献**：

- •我们介绍了KForge，一个多阶段自主程序合成框架，其中生成智能体和性能分析智能体协作，利用编译、正确性和性能分析反馈产生正确且优化的内核。
- •我们描述了一个系统架构，通过统一的程序合成接口支持四个加速器供应商（NVIDIA、AMD、Apple、Intel）和六种编程模型（CUDA、Triton、CuTe DSL、HIP、SYCL、Metal）。
- •我们展示了两个代表性后端的案例研究：在NVIDIA B200上实现端到端加速超越TensorRT-LLM，以及在Intel Arc B580上生成Triton内核——这里缺乏可比的手工调优参考。

## II 相关工作

**LLM驱动的内核生成**。越来越多的工作使用LLM生成和优化GPU内核，主要针对NVIDIA生态系统，因为那里有丰富的内核开发训练数据。Sakana AI的CUDA Engineer[10 (https://arxiv.org/html/2606.02963#bib.bib9)]使用了进化搜索来自动发现CUDA内核。后续分析发现该系统利用了评估框架的漏洞，说明在这个领域进行稳健评估的难度。CUDA-LLM[4 (https://arxiv.org/html/2606.02963#bib.bib17)]开发了一个特征搜索与强化（FSR）框架，结合编译、正确性和性能分析反馈来优化CUDA内核。KernelBlaster[7 (https://arxiv.org/html/2606.02963#bib.bib48)]为GPU编码智能体增加了一个跨任务积累的持久CUDA知识库。Autocomp[8 (https://arxiv.org/html/2606.02963#bib.bib47)]将LLM驱动的优化扩展到NVIDIA之外，更广泛地针对张量加速器，通过计划硬件优化和硬件反馈。然而，大多数工作倾向于关注每个内核的加速比，而不是衡量对全模型执行的影响。

**内核级基准测试**。已经开发了几个基准测试来评估LLM生成的内核。KernelBench[18 (https://arxiv.org/html/2606.02963#bib.bib36)]引入了一个包含250个PyTorch工作负载的基准框架，用于评估LLM生成高效GPU内核的能力。该基准使用`fast_p`指标，同时衡量正确性和相对于基线实现的加速比。相比之下，NVIDIA的SOL-ExecBench[11 (https://arxiv.org/html/2606.02963#bib.bib49)]专注于评估理论下限，使用一个分析管道，结合工作负载的FLOP计数和字节计数以及目标硬件的峰值带宽/FLOP能力。虽然这些基准推动了每个内核的评估，但它们留下了以下问题：一旦将LLM生成的内核集成到TensorRT-LLM等生产运行时中，图融合、自动调优和周围内核选择会与所研究的内核产生交互，这些内核如何影响端到端模型性能。

KForge通过统一的接口在四个AI加速器和六种编程模型上生成内核，并评估生成的内核在与集成生产运行时（如TensorRT-LLM）结合后的全模型执行影响。

## III KForge: 自主程序合成

参见图注  
图2：给定一个PyTorch参考实现，KForge为每个AI加速器选择并降级到合适的目标编程模型，支持在NVIDIA、AMD、Intel和Apple硬件上的CUDA、CuTe、Triton、HIP、SYCL和Metal。KForge是一个多阶段自主程序合成框架，如图1 (https://arxiv.org/html/2606.02963#S0.F1) 所示。它支持迭代优化、单次程序合成以及重复采样，操作模式由提示构建指导。它设计为跨平台框架，目前支持多种硬件目标，涵盖四个供应商和六种编程模型，具有不同的抽象级别，如图2 (https://arxiv.org/html/2606.02963#S3.F2) 所示。对于NVIDIA GPU，我们生成CUDA、Triton和CuTe DSL内核，涵盖底层手工调优和更高层次的抽象。对于AMD GPU，我们支持HIP和Triton。对于Intel Arc GPU，我们针对SYCL和Triton。对于Apple Silicon，我们生成Metal内核。KForge也是模型无关的；底层LLM通过模型注册表选择，允许用户插入任何模型而无需修改合成管道。

在本工作范围内，我们重点关注以下三种用于程序合成的策略。这些策略相互补充，允许根据可用的监督来源和计算资源预算构建动态配置。

- •**迭代优化**。它允许模型根据前一次运行纠正错误，或优化正确生成的内核性能，同时考虑前一次迭代合成的程序。具体来说，对于每次迭代 \(i \in \{1,\ldots,N-1\}\)，我们将第 \(i-1\) 次迭代的评估结果添加到模型的提示中，并附带相应的指令来修复错误或改进程序性能。
- •**跨平台翻译**。当某个加速器上已存在功能实现时，可以将其作为参考提供给模型，从而实现跨平台翻译。例如，在生成Metal内核时，提示可以包含对应的CUDA实现。这在目标后端缺乏训练数据覆盖时特别有用。
- •**性能分析反馈**。性能分析数据对于精确定位瓶颈至关重要，它提供了特定计算工作负载的硬件资源使用综合信息。时间线视图有助于识别调度间隙，而单个加速器API调用层面的详细统计数据则突出显示未能充分利用硬件资源的计算图部分。

### III-A 程序合成智能体

我们采用类似于[18 (https://arxiv.org/html/2606.02963#bib.bib36)]中的任务定义。具体来说，我们将LLM视为函数 \(F: (p) \mapsto k\)，接收文本提示 \(p \in \mathcal{T}\) 作为输入，并返回生成的代码 \(k \in \mathcal{T}\)。生成的代码预期包含一个内核程序、内核调度代码、JIT库编译代码以及一个PyTorch模型类 `NewModel(nn.Module)`，其中包含 `def forward(self, *inputs)` 方法实现模块的前向传播。我们使用Jinja2模板引擎来参数化提示。默认提示 \(p\) 包含高级任务描述、一个一次性示例（一个PyTorch实现及其对应的目标加速器实现）、一个PyTorch形式的输入问题以及一个自然语言的任务描述。

### III-B 性能分析智能体

我们引入一个专门的智能体进行性能分析，而不是将生成和优化放在一个智能体中处理，原因有二。首先，性能数据庞大但优化信号稀疏；先前研究[13 (https://arxiv.org/html/2606.02963#bib.bib2)]表明，对于32K令牌输入，LLM在相关信息检索上的性能下降到50%，而输入小于1K令牌时则不然。其次，分离这两个角色可以实现模块化架构，其中可以根据各自的优势将不同模型分配给每个智能体。

性能分析智能体处理性能输入，例如来自NVIDIA Nsight Systems的原始指标或来自Xcode Instruments的视觉输出，并为后续程序合成迭代生成优化建议。这种平台无关的方法可以处理不同硬件加速器上的任意文本或视觉性能数据。

形式上，该智能体被定义为...