构建可解释的分步优化数据集,让大模型学会一步步改进张量程序。
Step-TP: A Grounded, Step-Level Dataset with Chain-of-Thought Reasoning for LLM-Guided Tensor Program Optimization

- 提供原子级优化步骤与链式推理标注,支持逐步决策。
- 基于确定性中间表示,实现从代码到TVM TIR的可验证转换。
- 适合研究大模型在编译优化中的可解释性与多步推理能力。
尽管大语言模型具备强大的推理能力,但张量程序的执行效率优化仍因需要精确、可组合的变换决策而困难重重。现有基于LLM的方法将优化视为迭代决策过程,但数据集仅提供端到端优化结果对,使用低效的标记表示,缺乏可验证的分步监督与可解释性。导致模型难以在庞大的组合优化空间中做出可靠单步决策。本文提出Step-TP,一个面向张量程序优化的后训练数据集,提供基于结构化链式思维(CoT)的接地、原子、分步监督。该数据集通过中间程序状态构建闭环推理流程,支持可靠的多步优化而非仅模仿最终结果。其设计遵循四个原则:(i) 高效且可验证的中间表示(IR),能确定性地降至TVM TIR;(ii) 原子且可组合的优化策略,将复杂轨迹分解为可解释的单步决策;(iii) 结构化CoT监督与明确的IR-to-IR状态转移;(iv) 策略过滤机制,在覆盖广度与防止捷径利用间取得平衡。数据集与实现已开源至GitHub:https://github.com/LIUMENGFAN-gif/StepTP。
原文摘要 · Abstract (English)
Despite the strong reasoning capabilities of large language models (LLMs), optimizing the execution efficiency of tensor programs remains challenging due to the need for precise, composable transformation decisions. Recent LLM-guided approaches frame tensor program optimization as an iterative decision process, but existing datasets provide only end-to-end optimized program pairs using token-inefficient representations, lacking verifiable step-level supervision and interpretability. As a result, LLMs struggle to make reliable single-step decisions in large combinatorial optimization spaces. We introduce Step-TP, a post-training dataset for tensor program optimization that provides grounded, atomic, step-level supervision with structured chain-of-thought (CoT) reasoning. Step-TP forms a closed reasoning loop over intermediate program states, enabling reliable multi-step optimization rather than outcome imitation. Its design is guided by four principles: (i) a token-efficient, verifiable intermediate representation (IR) that deterministically lowers to TVM TIR; (ii) atomic and composable optimization strategies that decompose complex trajectories into interpretable single-step decisions; (iii) structured CoT supervision coupled with explicit IR-to-IR state transitions; and (iv) strategy filtering to balance coverage while preventing shortcut exploitation. The dataset and implementation are available at a GitHub link, https://github.com/LIUMENGFAN-gif/StepTP.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。