Shared-MCTS: Multi-Model Shared Search Tree Optimizes Deep Learning Compiler Performance

Shared-MCTS is an innovative multi-model collaborative search strategy. By enabling models of varying capabilities to share the same MCTS search tree, it achieves latency improvements for some workloads in TVM tensor program optimization. The project not only demonstrates the potential of multi-model collaboration but also openly discloses current limitations, providing a new research direction for deep learning compiler optimization.

The inference efficiency of deep learning models depends on the optimization of underlying tensor operations. Compilers like TVM need to find optimal configurations in a huge scheduling space (e.g., blocking, loop unrolling). Traditional manual tuning or heuristic search can hardly reach the upper limit. The LLM-based Reasoning Compiler models optimization as a sequential decision-making problem and combines it with MCTS to achieve results, but a single model faces a dilemma: strong models perform well but have high invocation costs, while small models have low costs but limited exploration capabilities.

Shared-MCTS allows multiple models of varying capabilities to share the same MCTS search tree. At each step of the search, it jointly decides the next operation and the model to generate candidates. In the exploration phase, lightweight models are used to quickly generate a large number of candidates; when converging, it switches to strong models for deep reasoning. All models share the tree, and experience flow avoids repeated exploration.

The experimental environment matches the official implementation of Reasoning Compiler: the same LLVM CPU backend, random seeds (0-4), 32 trials, and 4 LLM invocations per tuning. The workloads are operators with reduced scale but preserved core patterns: Llama-style Attention (batch=1, etc.), FLUX-style Convolution (batch=1, etc.), and Llama-style MLP (token=128, etc.).

Strict evaluation results:

Workload	Reasoning Compiler Latency	Shared-MCTS Latency	Improvement
FLUX-style Convolution	2.508 ms	2.339 ms	+6.76%
Llama Attention	0.158 ms	0.151 ms	+4.41%
Llama MLP	2.933 ms	3.039 ms	-3.63%
Overall: 2 out of 3 workloads show improvements, with a geometric mean speedup of 1.027x and a median improvement of 4.41%. It should be noted that the number of strong model invocations has not decreased, so the cost has not been reduced.

Insights:

1. Workload Sensitivity: Attention and convolution benefit significantly, while MLP shows a slight regression. Strategies need to be selected based on operator characteristics;
2. Model Collaboration Complexity: Simple pairing of strong and weak models is not a panacea. Fine-tuning of model switching, exploration-exploitation balance, and experience transfer is required;
3. Systematic Compiler Optimization: Improvements in individual operators do not mean end-to-end gains. System-level factors (operator interaction, memory bandwidth, etc.) need to be considered.

Limitations: Course project-level implementation, not a paper-level reproduction; reduced workloads, so results cannot be directly extrapolated to production environments; number of strong model invocations not reduced, so cost advantages not demonstrated; only verified on LLVM CPU backend, not covering GPU/CUDA scenarios. Future directions: Optimize multi-model collaboration strategies, reduce LLM invocation costs, and expand to more compilation scenarios.