用大模型多智能体系统优化PyTorch推理，平均提速2.88倍。
Optimizing PyTorch Inference with LLM-Based Multi-Agent Systems
- 设计多智能体协作框架，自动优化PyTorch代码以适配GPU
- 在H100上实现2.88倍于原生Eager模式的加速，优于torch.compile
- 适合追求极致推理性能的工程师与研究者
在现有GPU硬件上最大化AI推理性能仍是持续挑战。传统方法包括编写自定义GPU内核和使用专用模型编译器对高层代码进行目标特定调优。近期研究表明，基于大模型的多智能体系统可有效完成此类调优，通常优于现有编译器并省去手动内核开发。然而，该任务中多智能体系统的动态机制尚未被探索。本文提出一种逻辑框架，用于比较多智能体PyTorch优化系统。评估显示，结合纠错智能体时，以利用为主策略表现最佳，且性能与优化步骤粒度正相关。最佳实现方案在KernelBench基准测试集（涵盖多种PyTorch机器学习架构）上，于H100 GPU上对多样化任务实现了平均2.88倍于PyTorch Eager的加速（较torch.compile提升1.85倍）。代码已公开：https://github.com/pike-project/pike
原文摘要 · Abstract (English)
Maximizing performance on available GPU hardware is an ongoing challenge for modern AI inference systems. Traditional approaches include writing custom GPU kernels and using specialized model compilers to tune high-level code for specific GPU targets. Recent work shows that LLM-based multi-agent systems can effectively perform such tuning, often outperforming existing compilers and eliminating the need for manual kernel development. However, the dynamics of multi-agent systems for this task remain unexplored. In this work, we present a logical framework for comparing multi-agent PyTorch optimization systems. Our evaluation shows that exploit-heavy strategies perform best when paired with error-fixing agents, and that performance correlates with the granularity of optimization steps. The best implementation achieves an average 2.88x speedup over PyTorch Eager (1.85x over torch.compile) on an H100 GPU across diverse tasks in KernelBench, a benchmark suite covering a range of machine learning architectures in PyTorch. Code is publicly available at: https://github.com/pike-project/pike
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。