# Toward Compiler World Models: Learning Latent Dynamics for Efficient Tensor Program Search

arXiv: https://arxiv.org/abs/2606.09312 · HTML: https://mlsanity.com/paper/2606.09312/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in World models for language agents, one of 11 directions in World models.
- Builds on 9 of its 48 references that are in the map.
- First posted on arXiv 8 Jun 2026.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)

## Builds on (most cited first, 9 of 9)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [DeepSeek-V3 Technical Report](https://mlsanity.com/paper/2412.19437/) (2024)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [MLIR: A Compiler Infrastructure for the End of Moore's Law](https://mlsanity.com/paper/2002.11054/) (2020)
- [AutoKernel: Autonomous GPU Kernel Optimization via Iterative Agent-Driven Search](https://mlsanity.com/paper/2603.21331/) (2026)
- [A Two-Stage GPU Kernel Tuner Combining Semantic Refactoring and Search-Based Optimization](https://mlsanity.com/paper/2601.12698/) (2026)

## Abstract (arXiv)

Tensor program optimization is essential for modern machine learning systems, but its search space is enormous. Existing auto-schedulers reduce measurement cost with learned cost models, yet they usually evaluate each candidate as a static code snapshot, ignoring the schedule trajectory that produced it. This makes them insensitive to action dependencies and vulnerable to superficial code variations. We propose a world-model-inspired evaluator that models schedule evaluation as action-conditioned latent dynamics over program states. Starting from the initial program, it rolls out scheduling actions in a continuous latent space with a lightweight transition model, avoiding expensive AST mutation and repeated code encoding. The final dynamic representation is combined with action and hardware features to rank candidates. Implemented in TVM AutoScheduler, our method improves representative-subgraph latency over Ansor by 1.37× on GPU and 1.54× on CPU under the same 64-trial budget. It also matches Ansor-10K within 2.2% geometric mean using 10× fewer measurements, and accelerates full-model inference over PyTorch/PyTorch-opt(cuDNN) by 4.61×/3.67× geometric mean.