Multi agent teams improve gpu kernel optimization performance
Evolving Agent Teams
Artificial Intelligence
Summary
Optimizing GPU code is hard because many smart programs get stuck improving only a little and miss better solutions. The authors developed a system called Evolving Agent Teams, where multiple software agents work together and a supervisor reshapes their efforts based on past experience. This approach helps avoid getting stuck in small improvements and finds faster GPU code. Their system outperformed other coding programs and even made new, faster GPU kernels from scratch on real-world tasks.
What this means in practice
- •For gpu software engineers: Create faster GPU kernels by dynamically coordinating multiple coding agents to escape local performance plateaus.
- •For robotics developers: Improve robotic manipulation behaviors by evolving agent teams that adapt strategies based on accumulated experience.
Authors
Shiyi Cao, Ziming Mao, Dacheng Li, Joseph E. Gonzalez, Ion Stoica
Abstract
The rapid progress of software engineering agents drives interest in applying them to open-ended problems such as GPU kernel optimization. However, we find that existing coding agents converge to local optima on real-world, expert-written kernels (e.g., FlashInfer). We introduce Evolving Agent Teams, a multi-agent framework that escapes such local optima. It features (1) a central supervisor that diagnoses each local optimum and restructures a team of agents; and (2) a persistent file system that comprises a learning document, an evolution history, and per-team artifacts, grounding restructuring with accumulated experience. Under a matched token budget, Evolving Agent Teams consistently outperforms strong coding-agent baselines. With an extended token budget, Evolving Agent Teams generates a CUDA kernel that reaches 1.2x mean speedup over the production FlashInfer MLA paged decode kernel from scratch and wins on 36 of 47 production workloads from FlashInfer-Bench (0.86-0.99x on the remaining 11). It also demonstrates superior performance on tasks such as KDA prefill on AMD MI350X (3.01x over vLLM) and NVIDIA B200 (3.44x over FlashKDA), MoE expert-offload serving on Apple silicon (30.3 tokens/s decode for Qwen3.8-Flash-Next in MLX 4-bit on a 48 GB M4 Pro), and RoboDojo robot manipulation (62.1 on held-out seeds vs. 55.3 for the best leaderboard entry). Code is available at https://github.com/caoshiyi/EAT.