Optimizing PyTorch Inference with LLM-Based Multi-Agent Systems

TL;DR

The paper tackles the challenge of maximizing PyTorch inference speed on GPUs by introducing a logical framework for LLM-driven multi-agent optimization and two concrete strategies, PIKE-B and PIKE-O. Through a refined KernelBench evaluation, the authors show that exploit-heavy configurations coupled with robust error fixing achieve the best average speedups, up to $2.88\times$ on H100. They analyze how per-step optimization granularity and backend characteristics (CUDA vs Triton) influence gains, offering practical guidance on configuring MAS autotuners. The work advances the state of multi-agent code optimization for ML workloads and demonstrates meaningful speedups with a principled exploration–exploitation balance and error-correction mechanisms.

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 on an H100 GPU across diverse tasks in KernelBench, a benchmark suite covering a range of machine learning architectures in PyTorch.

Figures (7)

• Figure 1: Logical framework which multi-agent systems operate in, for the task of PyTorch inference optimization. Logical steps of the optimization process are shown above, and core options/parameters are shown below near the steps they relate to.
• Figure 2: PIKE-B implementation, illustrating parallel evaluation and error fixing, followed by top-k selection and mutation
• Figure 3: Geomean speedup across Level 3-pike tasks of each PIKE implementation by (a) LLM queries per task and (b) cost in $ per task, all on an H100.
• Figure 4: Geomean speedup across Level 5 tasks of each PIKE implementation by (a) LLM queries per task and (b) cost in $ per task, all on an H100.
• Figure 5: Speedup relative to PyTorch Eager of our approaches (including extra ablations) and other approaches, using an H100. All of our approaches get a budget of exactly 300 LLM queries.