Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
The paper explores the use of LLM-based multi-agent systems for optimizing PyTorch inference on GPUs, aiming to surpass traditional compilers and manual kernel development. They introduce a logical framework to compare different multi-agent optimization strategies, focusing on the dynamics between exploration, exploitation, and error-fixing. Their best system, combining exploit-heavy strategies with error-fixing agents and fine-grained optimization steps, achieves a 2.88x speedup on an H100 GPU across KernelBench.
LLM-powered agents can now optimize PyTorch inference by 2.88x on H100 GPUs, but only if you carefully balance exploit-heavy search with agents that fix errors.
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.