用动态模型优化张量程序搜索,提速超1.5倍
Toward Compiler World Models: Learning Latent Dynamics for Efficient Tensor Program Search

- 将调度过程建模为潜空间中的动态演化,避免重复编码
- 64次尝试下比Ansor快1.37倍(GPU)和1.54倍(CPU)
- 只需少量测量即可逼近万次尝试效果,适合资源受限场景
张量程序优化对现代机器学习系统至关重要,但搜索空间巨大。现有自动调度器虽用学习成本模型降低测量开销,却通常将每个候选程序视为静态快照,忽略生成它的调度轨迹,导致对操作依赖不敏感且易受表面代码变化干扰。本文提出一种受世界模型启发的评估器,将调度评估建模为基于动作条件的程序状态潜空间动态演化。从初始程序出发,通过轻量级转移模型在连续潜空间中滚动执行调度动作,避免昂贵的AST修改与重复编码。最终动态表征结合动作与硬件特征进行候选排序。在TVM AutoScheduler中实现,该方法在相同64次尝试预算下,对代表性子图的延迟分别比Ansor提升1.37倍(GPU)和1.54倍(CPU);仅用10%测量次数即达到Ansor-10K的97.8%几何平均性能;在PyTorch/PyTorch-opt(cuDNN)上加速全模型推理4.61倍/3.67倍几何平均。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。