论文
KForge:LLM 驱动的 AI 加速器跨平台内核生成
KForge: LLM-Driven Cross-Platform Kernel Generation for AI Accelerators
摘要
生产推理越来越多地针对异构加速器组合。代理管道交织推理、工具调用和多代理协调,每个管道都具有不同的计算和内存配置文件。为了获得最佳效率,每个阶段都应该在最适合它的加速器上运行。这带来了系统挑战:每个管道现在都需要跨越来越多的硬件后端和编程模型的高性能内核。手动编写这些内核非常耗时,需要深厚的底层专业知识,并且不会随着内核复杂性的增长而扩展。最近,大语言模型(LLM)已被用于自动内核生成,但底层代码生成和跨后端泛化方面的挑战仍然存在。我们提出了 KForge,一个围绕迭代细化循环构建的跨平台框架,由两个协作的基于 LLM 的代理驱动:生成代理使用编译和正确性反馈生成并逐步细化内核,性能分析代理解释从编程 API 到基于 GUI 的工具的分析数据,并发出指导下一轮综合的建议。该循环在功能通道和优化通道之间交替,功能通道使候选者达到正确性,优化通道缩小了与手动调整基线的性能差距。我们在两个具有不同基线参考可用性的后端上评估 KForge。在 NVIDIA B200 上,KForge 在 gpt-oss-20b 推理速度基准测试中与 TensorRT-LLM 相比,端到端吞吐量提高了 2.12$\%$。在 Intel Arc B580 上,KForge 生成的 Triton 内核在 KernelBench Level 2 的 37 GEMM + tail-ops 工作负载上,主要通过运算符融合和混合精度执行,比更快的 PyTorch eager 和 torch.compile 实现了 5.13$\times$ 几何平均加速。