meta-pytorch/KernelAgent

解决的问题

KernelAgent 自动化 GPU 内核的创建与性能调优。它解决了手动编写高性能 Triton 内核的困难，通过将 PyTorch 程序转换为经过验证的优化 Triton 代码，减少了实现硬件特定性能提升所需的专家级 GPU 编程知识。

工作原理

该系统通过两个主要流程运行：

1. 内核生成：使用「AutoRouter」分析 PyTorch AST 并决定最佳路径。"Fuser Orchestrator" 将代码重构为可融合的子图，然后分派给并行的 LLM 驱动的代理。这些代理迭代生成 Triton 内核并根据单元测试进行验证。最后，"Composer" 将这些内核组合成一个完整的、经过验证的 Triton 程序。

2. 内核优化：该流程采用硬件引导的循环。使用 NVIDIA NCU 对内核进行性能分析，收集硬件指标，执行屋顶线分析以判断内核是内存受限还是计算受限，并利用 LLM 诊断瓶颈并提出修复建议。该循环持续进行，直到内核达到性能极限（光速）或收敛为止。

适用人群

• 希望在不手动编写原始 Triton 内核的情况下优化 PyTorch 操作的 GPU 开发者和 ML 工程师。
• 从事 NVIDIA 和 Intel XPU GPU 的自定义算子融合与硬件特定内核调优的研究人员。

核心亮点

• 多平台支持：支持 NVIDIA CUDA 和 Intel XPU GPU。
• 多代理编排：使用并行工作器基于严格的运行时验证生成和优化内核。
• 硬件引导调优：将真实的 GPU 性能分析（NCU）和屋顶线分析集成到 LLM 优化循环中。
• 端到端自动化：从 PyTorch 代码分析和子图提取，到最终组合与基准测试，全程自动化。