cs.LGFeb 2, 2026

MOSAIC: More LLMs, Less Cost through Collaborative Reasoning for Compiler Optimizations

Authors: Annabelle Sujun Tang, Christopher Priebe, Lianhui Qin, Hadi Esmaeilzadeh

Organizations: Alternative Computing Technologies (ACT) Lab, University of California San Diego

Abstract

LLM-guided compiler optimization has recently shown promise. However, current LLM-guided optimizers query a single LLM at every search step, so their cost grows with both that model's price and the search budget. A smaller model is cheaper but degrades the search: in our experiments, a recent single-LLM optimizer reaches 10.2% (6.0%) lower speedup on GPU (CPU) with gpt-5-mini than with GPT-5.2. Agentic multi-LLM frameworks add controllers, concurrent calls, and shared memory, whose overhead can offset the savings of cheaper models. This paper shows that the optimization search tree itself can coordinate heterogeneous LLMs: the tree already records which proposals improved the program, and we can augment the tree to record which model made them. That evidence can then be used to decide which model should propose next. To this end, we introduce MOSAIC, in which each Monte Carlo tree search (MCTS) node records a program and the LLM that expands it, and the active LLM proposes both a compiler transformation and the LLM for the next step. We also contribute an LLM-aware UCT that favors smaller models and a course-alteration rule that invokes the largest model only after repeated small-model regressions. Against that single-LLM optimizer run with GPT-5.2 alone and the same number of search iterations and tree parameters, on an NVIDIA GPU (Intel CPU), MOSAIC with eight LLMs improves the geometric-mean speedup on five kernels by 7.2% (14.4%) and the end-to-end Llama-3-8B speedup by 1.61x (1.41x), while reducing kernel compilation time by 1.95x (1.74x) and API cost by 4.47x (4.32x), with GPT-5.2 making only 23.1% (23.9%) of LLM calls. With the same eight models, random or round-robin model choice improves over the single-LLM baseline by only 2.5% on CPU, showing that the gains come from letting the LLMs choose the next model rather than from mixing the pool at random.

Figures & tables

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

May 19, 2026

Prior Knowledge or Search? A Study of LLM Agents in Hardware-Aware Code Optimization

LLM discovery and optimization systems are increasingly applied across domains, implementing a common propose-evaluate-revise loop. Such optimization or discovery progresses via context conditioning on received feedback from an environment. However, as modern LLM agents are increasingly complex in their structure, it is difficult to evaluate which components contribute the most, and when and how this exploration may fail. We answer these questions through three controlled experiments. Our findings: (1) In pure black-box optimization, LLMs act as greedy optimizers. (2) In zero-shot kernel generation, providing explicit input-size information has no measurable effect, models converge to the same kernel parameters regardless of size or temperature, as though the size instruction were invisible. Moreover, when tasked to perform kernel optimization for uncommon kernel sizes, performance sharply degrades regardless of the language used. (3) In feedback-loop kernel optimization, CUDA improves monotonically under iterative feedback, while TVM IR actively degrades, which demonstrates that kernel optimization degrades when models operate with low-density language. Our results conclude that LLMs in code optimization tasks highly depend on pretrained priors rather than provided feedback or agentic structure.
Jan 31, 2025

LLM Program Optimization via Retrieval Augmented Search

Recent work has demonstrated the potential of large language models (LLMs) for program optimization, a key challenge in programming languages. We propose a blackbox adaptation method called Retrieval Augmented Search (RAS) that performs beam search over candidate optimizations; at each step, it retrieves in-context examples from a given training dataset of slow-fast program pairs to guide the LLM. Critically, we find that performing contextual retrieval based on an LLM-generated natural language description significantly outperforms retrieval based on the source code. We also propose AEGIS, a method for improving interpretability by decomposing training examples into ''atomic edits'' that are significantly more incremental in nature. We show that RAS performs up to 2.06×\times better than prior state-of-the-art blackbox adaptation strategies on optimizing C++ programs, and that AEGIS performs up to 1.37×\times better while making significantly smaller edits. We also show that using RAS improves the mean runtime percentile of Python programs by 10.27 compared to baselines.
Aug 8, 2025

Optimizing Prompt Sequences using Monte Carlo Tree Search for LLM-Based Optimization

Mathematical optimization is a core tool in engineering, operations research, and data science, but modeling real problems remains difficult and time-consuming. In practice, an expert must work with clients to translate informal domain descriptions into precise objectives, constraints, and solver-ready code; small mistakes in this process can lead to infeasible, nonconvex, or suboptimal formulations, and the iterative process is often slow. Large language models (LLMs) could automate this translation from natural language to optimization code, but standard prompting often fails on realistic tasks: it struggles with long, structured task descriptions, mishandles constraints, and produces infeasible, non-convex, or non-executable programs. We introduce Monte Carlo Tree Search for Reasoning Decomposition Plans (MCTS-RDP), which performs pre-hoc test-time scaling by searching over decomposition plans and their resulting programs before committing to the final output. MCTS-RDP treats plan construction as a sequential decision problem: given a problem description, it decomposes the task descriptions into logical components and generates candidate plan steps for each component. Monte Carlo Tree Search (MCTS) then explores and refines Reasoning Decomposition Plans (RDPs) prior to code synthesis, balancing exploration and exploitation across alternative plan sequences. Each candidate plan is rolled out into executable code, executed, and scored with an execution-grounded evaluator measuring executability, constraint satisfaction, and solution quality; this feedback is backpropagated through the search tree to guide subsequent search. Across wireless optimization tasks and the NL4OPT and NLP4LP benchmarks, MCTS-RDP consistently outperforms strong baselines, improving execution success, reward, stability, and optimality attainment.