Optimizing PyTorch Inference with LLM-Based Multi-Agent Systems

TL;DR

This paper introduces a logical framework to evaluate LLM-based multi-agent systems for optimizing PyTorch inference, demonstrating significant speedups over standard methods on GPU hardware.

Contribution

It presents a novel framework for comparing multi-agent optimization strategies and shows that exploit-heavy strategies with error-fixing agents yield the best performance.

Findings

Exploit-heavy strategies perform best with error-fixing agents.

Performance improves with finer granularity of optimization steps.

Achieved 2.88x speedup over PyTorch Eager on H100 GPU.

Abstract

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