Toward Compiler World Models: Learning Latent Dynamics for Efficient Tensor Program Search
作者: Haolin Pan, Lianghong Huang, Xvlin Zhou, Mingjie Xing, Yanjun Wu
分类: cs.LG, cs.PL
发布日期: 2026-06-08
💡 一句话要点
提出世界模型启发的评估器以优化张量程序搜索
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 张量程序优化 自动调度 潜在动态 机器学习 性能提升
📋 核心要点
- 现有的自动调度器在评估候选程序时，往往忽视了调度轨迹，导致对动作依赖性不敏感。
- 本文提出了一种基于世界模型的评估器，通过建模潜在动态来优化调度评估，避免了昂贵的代码变换。
- 实验结果显示，该方法在GPU和CPU上分别提高了1.37倍和1.54倍的性能，并显著减少了测量次数。
📝 摘要（中文）
张量程序优化对现代机器学习系统至关重要，但其搜索空间巨大。现有的自动调度器通过学习成本模型来降低测量成本，但通常将每个候选者视为静态代码快照，忽略了生成它的调度轨迹。这使得它们对动作依赖性不敏感，并容易受到表面代码变动的影响。本文提出了一种世界模型启发的评估器，将调度评估建模为程序状态上的动作条件潜在动态。该方法从初始程序开始，在连续潜在空间中展开调度动作，避免了昂贵的AST变换和重复的代码编码。最终的动态表示与动作和硬件特征结合，以对候选者进行排名。我们的实现提高了代表子图在GPU上1.37倍、在CPU上1.54倍的延迟表现。
🔬 方法详解
问题定义：本文旨在解决张量程序优化中的搜索空间巨大和评估不准确的问题。现有方法通常将候选程序视为静态快照，忽略了调度过程中的动态变化，导致评估结果不够准确。
核心思路：我们提出了一种世界模型启发的评估器，利用动作条件潜在动态来建模调度评估。这种方法通过在连续潜在空间中展开调度动作，能够更好地捕捉程序状态的变化，从而提高评估的准确性。
技术框架：整体架构包括初始程序状态、轻量级的转移模型和最终的动态表示。首先，从初始程序出发，利用转移模型在潜在空间中展开调度动作，生成动态表示。然后，将该表示与动作和硬件特征结合，进行候选程序的排名。
关键创新：本研究的主要创新在于将调度评估视为潜在动态过程，而非静态快照。这一方法能够更好地捕捉程序状态的变化和动作之间的依赖关系，从而提高了评估的准确性和效率。
关键设计：在设计中，我们采用了轻量级的转移模型，避免了复杂的AST变换和重复编码。同时，结合了动作和硬件特征，以增强评估的全面性和准确性。
🖼️ 关键图片
📊 实验亮点
实验结果显示，所提出的方法在GPU上提高了1.37倍的代表子图延迟，在CPU上提高了1.54倍。同时，该方法在使用10倍更少的测量次数的情况下，能够与Ansor-10K的表现相匹配，展示了其在性能和效率上的显著提升。
🎯 应用场景
该研究的潜在应用领域包括机器学习模型的自动调度和优化，尤其是在需要高效张量计算的场景中，如深度学习和大规模数据处理。通过提高调度效率，可以显著降低计算资源的消耗，提升系统的整体性能，对未来的智能系统发展具有重要影响。
📄 摘要（原文）
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.