Optimizing PyTorch Inference with LLM-Based Multi-Agent Systems

• URL: http://arxiv.org/abs/2511.16964v1
• Date: Fri, 21 Nov 2025 05:37:38 GMT
• Title: Optimizing PyTorch Inference with LLM-Based Multi-Agent Systems
• Authors: Kirill Nagaitsev, Luka Grbcic, Samuel Williams, Costin Iancu,
• Abstract summary: We present a framework for comparing multi-agent PyTorch optimization systems.<n>We show that exploit-heavy strategies perform best when paired with error-fixing agents.<n>The best implementation achieves an average 2.88x speedup on an H100 GPU.
• Score: 1.2289544895833646
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning [26.264303471292845]
  We propose StitchCUDA, a multi-agent framework for end-to-end GPU program generation.<n> Experiments show that StitchCUDA achieves nearly 100% success rate on end-to-end programming tasks, with 1.72x better speedup than the multi-agent baseline.
  arXiv  Detail & Related papers  (2026-03-03T06:04:49Z)
• PRAGMA: A Profiling-Reasoned Multi-Agent Framework for Automatic Kernel Optimization [12.24680414520151]
  PRAGMA is a profile-guided AI kernel generation framework.<n>It integrates execution feedback and fine-grained hardware profiling into the reasoning loop.<n>We evaluate PRAGMA on KernelBench, covering GPU and CPU backends.
  arXiv  Detail & Related papers  (2025-11-09T12:01:43Z)
• STARK: Strategic Team of Agents for Refining Kernels [23.717055490630596]
  We introduce an agentic framework for GPU kernel optimization that explores the design space through multi-agent collaboration.<n>This framework mimics the workflow of expert engineers, enabling LLMs to reason about hardware trade-offs, incorporate profiling feedback, and refine kernels iteratively.<n>We evaluate our approach on KernelBench, a benchmark for LLM-based kernel optimization, and demonstrate substantial improvements over baseline agents.
  arXiv  Detail & Related papers  (2025-10-19T20:41:46Z)
• Astra: A Multi-Agent System for GPU Kernel Performance Optimization [10.715861478214961]
  We introduce Astra, the first multi-agent system for GPU kernel optimization.<n>Within Astra, specialized agents collaborate through code generation, profiling, and planning to produce kernels that are both correct and high-performance.
  arXiv  Detail & Related papers  (2025-09-09T08:39:50Z)
• NGPU-LM: GPU-Accelerated N-Gram Language Model for Context-Biasing in Greedy ASR Decoding [54.88765757043535]
  This work rethinks data structures for statistical n-gram language models to enable fast and parallel operations for GPU-optimized inference.<n>Our approach, named NGPU-LM, introduces customizable greedy decoding for all major ASR model types with less than 7% computational overhead.<n>The proposed approach can eliminate more than 50% of the accuracy gap between greedy and beam search for out-of-domain scenarios while avoiding significant slowdown caused by beam search.
  arXiv  Detail & Related papers  (2025-05-28T20:43:10Z)
• KGym: A Platform and Dataset to Benchmark Large Language Models on Linux Kernel Crash Resolution [59.20933707301566]
  Large Language Models (LLMs) are consistently improving at increasingly realistic software engineering (SE) tasks. In real-world software stacks, significant SE effort is spent developing foundational system software like the Linux kernel. To evaluate if ML models are useful while developing such large-scale systems-level software, we introduce kGym and kBench.
  arXiv  Detail & Related papers  (2024-07-02T21:44:22Z)
• Hardware-Aware Parallel Prompt Decoding for Memory-Efficient Acceleration of LLM Inference [23.633481089469836]
  Auto-regressive decoding of Large Language Models (LLMs) results in significant overheads in their hardware performance.<n>We propose a novel parallel prompt decoding that requires only $0.0002$% trainable parameters, enabling efficient training on a single A100-40GB GPU in just 16 hours.<n>Our approach demonstrates up to 2.49$times$ speedup and maintains a minimal memory overhead of just $0.0004$%.
  arXiv  Detail & Related papers  (2024-05-28T22:19:30Z)
• Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures [67.47328776279204]
  This work introduces a framework to develop efficient, portable Deep Learning and High Performance Computing kernels. We decompose the kernel development in two steps: 1) Expressing the computational core using Processing Primitives (TPPs) and 2) Expressing the logical loops around TPPs in a high-level, declarative fashion. We demonstrate the efficacy of our approach using standalone kernels and end-to-end workloads that outperform state-of-the-art implementations on diverse CPU platforms.
  arXiv  Detail & Related papers  (2023-04-25T05:04:44Z)
• Kernel methods through the roof: handling billions of points efficiently [94.31450736250918]
  Kernel methods provide an elegant and principled approach to nonparametric learning, but so far could hardly be used in large scale problems. Recent advances have shown the benefits of a number of algorithmic ideas, for example combining optimization, numerical linear algebra and random projections. Here, we push these efforts further to develop and test a solver that takes full advantage of GPU hardware.
  arXiv  Detail & Related papers  (2020-06-18T08:16:25Z)
• MPLP++: Fast, Parallel Dual Block-Coordinate Ascent for Dense Graphical Models [96.1052289276254]
  This work introduces a new MAP-solver, based on the popular Dual Block-Coordinate Ascent principle. Surprisingly, by making a small change to the low-performing solver, we derive the new solver MPLP++ that significantly outperforms all existing solvers by a large margin.
  arXiv  Detail & Related papers  (2020-04-16T16:20:53Z)
• PolyScientist: Automatic Loop Transformations Combined with Microkernels for Optimization of Deep Learning Primitives [55.79741270235602]
  We develop a hybrid solution to the development of deep learning kernels. We use the advanced polyhedral technology to automatically tune the outer loops for performance.
  arXiv  Detail & Related papers  (2020-02-06T08:02:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.