Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper introduces a world-model-inspired evaluator for tensor program optimization that leverages action-conditioned latent dynamics to efficiently search the vast space of scheduling actions. By modeling the evaluation process in a continuous latent space, the approach avoids costly AST mutations and repeated code encodings, allowing for a more nuanced understanding of action dependencies. The implementation in TVM AutoScheduler demonstrates significant performance improvements, achieving up to 1.54× latency reduction on CPU and 4.61× acceleration in full-model inference compared to existing methods.
A novel evaluator using latent dynamics achieves over 1.5× faster tensor program optimization while requiring 10× fewer measurements.
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 \emph{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$\times$ on GPU and 1.54$\times$ on CPU under the same 64-trial budget. It also matches Ansor-10K within 2.2% geometric mean using 10$\times$ fewer measurements, and accelerates full-model inference over PyTorch/PyTorch-opt(cuDNN) by 4.61$\times$/3.67$\times$ geometric mean.