COLT: Lightweight Multi-LLM Collaboration through Shared MCTS Reasoning for Model Compilation
Annabelle Sujun Tang, Christopher Priebe, Lianhui Qin, Hadi Esmaeilzadeh
TL;DR
COLT tackles the high cost of LLM-driven compiler optimization by embedding endogenous model selection into a single shared Monte Carlo tree search, allowing multiple smaller LLMs to collectively explore the transformation space. A model-aware MA-UCT policy biases search toward using smaller models while preserving exploration, and a course-alteration mechanism intermittently escalates to the largest LLM to prevent persistent regressions. Proposals from different LLMs share and backpropagate through one MCTS tree, enabling cross-model learning without external controllers. Empirically, COLT delivers substantial speedups over a single large-model baseline across five kernels on CPU and GPU, while reducing reliance on the largest model and maintaining robustness across hardware.
Abstract
Model serving costs dominate AI systems, making compiler optimization essential for scalable deployment. Recent works show that a large language model (LLM) can guide compiler search by reasoning over program structure and optimization history. However, using a single large model throughout the search is expensive, while smaller models are less reliable when used alone. Thus, this paper seeks to answer whether multi-LLM collaborative reasoning relying primarily on small LLMs can match or exceed the performance of a single large model. As such, we propose a lightweight collaborative multi-LLM framework, dubbed COLT, for compiler optimization that enables coordinated reasoning across multiple models within a single Monte Carlo tree search (MCTS) process. A key contribution is the use of a single shared MCTS tree as the collaboration substrate across LLMs, enabling the reuse of transformation prefixes and cross-model value propagation. Hence, we circumvent both heavy internal reasoning mechanisms and conventional agentic machinery that relies on external planners, multiple concurrent LLMs, databases, external memory/versioning of intermediate results, and controllers by simply endogenizing model selection within the lightweight MCTS optimization loop. Every iteration, the acting LLM proposes a joint action: (compiler transformation, model to be queried next). We also introduce a model-aware tree policy that biases search toward smaller models while preserving exploration, and a course-alteration mechanism that escalates to the largest model when the search exhibits persistent regressions attributable to smaller models.