Step-TP

Step-TP是由香港大学、蚂蚁集团及中国科学技术大学联合构建的专用于大语言模型引导的张量程序优化的后训练数据集。该数据集以LEIR中间表示为核心，提供了原子化、可验证的步骤级监督数据，并融合结构化思维链推理，覆盖了多样化的优化策略。其构建过程系统地将复杂优化轨迹分解为可组合的原子策略，通过策略过滤机制平衡覆盖范围，确保数据质量。该数据集旨在解决张量程序优化中LLM缺乏可靠单步决策能力的问题，推动其在GPU高效执行、深度学习编译器优化等领域的应用，实现可解释、可扩展的自动化性能优化。Step-TP is a post-training dataset co-developed by The University of Hong Kong, Ant Group, and the University of Science and Technology of China, specifically designed for large language model-guided tensor program optimization. Centered on the LEIR intermediate representation, this dataset provides atomic, verifiable step-level supervision data, integrates structured chain-of-thought reasoning, and covers a diverse range of optimization strategies. Its construction process systematically decomposes complex optimization trajectories into combinable atomic strategies, and balances coverage via a strategy filtering mechanism to ensure data quality. This dataset aims to address the issue that large language models (LLMs) lack reliable single-step decision-making capabilities in tensor program optimization, promote their applications in fields such as efficient GPU execution and deep learning compiler optimization, and achieve interpretable, scalable automated performance optimization.

原始信息汇总

StepTP 数据集详情

数据集概述

StepTP 数据集旨在学习张量程序变换中的单步和多步变换，包含用于训练和推理的代码及数据集。

• 数据集文件：dataset.zip 包含一个聊天模板格式的 JSONL 文件，用于微调：train_dataset_with_CoT_in_chattemplate.jsonl。
• 数据集状态：当前公开版本仅包含微调用聊天模板数据集，完整的档案字典数据集及中间构件将在论文被接收后发布。

环境要求

开发和测试环境如下：

• 操作系统：Ubuntu 22.04.4
• CUDA：12.4
• LLVM：17.0.6
• Python：3.10.20
• PyTorch：2.6.0+cu124
• TVM：0.21.dev0
• 其他 Python 包：tqdm, transformers, sympy, ms-swift

代码仓库结构

仓库包含以下核心模块和脚本：

文件/目录	功能描述
pytorch_model_to_IR.py	使用 torch.fx 追踪 PyTorch 程序并转换为 LEIR
ops/	解析 PyTorch 操作并生成 LEIR 表达式的工具
strategies/	实现变换策略（strategy.py, utils.py）
apply_strategy_to_IR.py	将选定策略应用于 LEIR，生成单步变换 LEIR
construct_CoT_dataset.py	为变换后的 LEIR 构建推理链
transform_IR_to_TIR.py	将 LEIR 转换为 TVM TIR 并执行策略级检查
TIR/	LEIR 到 TIR 的降级、TVM 构建辅助和验证工具
eval_verification.py 等	验证和评估入口点
filtering.py 等	数据集后处理和过滤工具
swift_Qwen.py	使用 MS-Swift 微调 Qwen 模型
swift_inference.py	单步变换推理
swift_*_optimization.py	多步优化搜索变体（chain、DFS、BFS、BS、GS、MCTS）
model_codes/	PyTorch 模型/算子程序，用作源工作负载
dataset.zip	公开的聊天模板训练数据集

流水线映射

论文描述的四个阶段及其对应代码：

1. PyTorch 到 LEIR 翻译

   • 主入口：pytorch_model_to_IR.py
   • 辅助代码：ops/, model_codes/, level1_model_name.json, level2_model_name.json

2. 单步策略驱动的变换

   • 主入口：apply_strategy_to_IR.py
   • 策略实现：strategies/strategy.py, strategies/utils.py
   • 推理链构建：construct_CoT_dataset.py
   • 策略列表：strategy_names.json

3. LEIR 到 TIR 翻译

   • 主入口：transform_IR_to_TIR.py
   • 辅助代码：TIR/IR_to_TIR.py, TIR/utils.py

4. 验证与过滤

   • 验证：TIR/IR_verfication.py, eval_verification.py, eval_verification_main.py, eval_verification_update.py
   • 过滤与数据集处理：filtering.py, handle_dataset_v2.py, handle_dataset_multi_IRs.py, further_handle_dataset.py

训练

使用 MS-Swift 微调 Qwen 模型示例：

bash python swift_Qwen.py --train_dataset dataset/train_dataset_with_CoT_in_chattemplate.jsonl --model_id_or_path /path/to/Qwen3-8B --output_dir /path/to/output/sft_8B_CoT_filtered --cuda_visible_devices 0,1 --batch_size 1 --epochs 1 --gradient_accumulation_steps 16 --max_length 10240

推理与评估

• 单步变换推理：python swift_inference.py --cuda 0 --batch_size 16
• 多步优化：使用基于搜索的脚本（如 swift_chain_based_optimization.py、swift_DFS_based_optimization.py 等），需替换默认路径为本地评估提示和生成的 store-dictionary 文件。

验证

首先生成 store-dictionary 文件（由 pytorch_model_to_IR.py 在 PyTorch 到 LEIR 翻译阶段产生）：

bash python pytorch_model_to_IR.py --level level1 --model_index 0 --dtype float64 --store_dict_dir store_dict

然后使用 TVM 验证变换后的 LEIR：

bash python eval_verification.py --idx 0 --transformed_IR "<TRANSFORMED_LEIR>" --info store_dict/<MODEL_NAME>original<MODEL_INDEX>.pkl.gz --num_IRs 1 --cuda 0 --error_record_dir /path/to/error_record.jsonl

可重现性说明

• 当前 dataset.zip 足以使用 swift_Qwen.py 进行训练。
• 完整的档案字典数据集将在论文被接收后发布。
• 单步和多步推理/评估可使用用户创建的评估提示和本地生成的 store-dictionary 文件运行，论文评估文件将在被接收后发布。
• TVM 验证脚本需要启用 CUDA 和 LLVM 的 TVM，并访问每个数据条目引用的 store-dictionary。
• 部分推理和优化脚本使用硬编码的作者路径，在新环境中运行前需更新这些路径。

搜集汇总

数据集介绍

构建方式

Step-TP的构建源于对现有张量程序优化数据集局限性的深刻洞察。现有数据集多依赖端到端的优化程序对，以token效率低下的表示形式呈现，缺乏可验证的步骤级监督与可解释性。为突破这一瓶颈，Step-TP设计了一条多阶段流水线。首先，通过PyTorch-to-LEIR翻译器，将来自KernelBench的189种不同程序类型以及随机组合的复合程序，转换为一种名为LEIR的高密度且可验证的中间表示。随后，进入核心的单步策略驱动变换阶段：对每个原始LEIR进行可行性检查以确定可应用的策略集合，通过分解-修改-重组的流程逐一执行原子化变换，并为每个变换步骤生成结构化的思维链推理轨迹，详细阐述优化逻辑。最后，经过LEIR-to-TIR翻译器生成可执行代码，再通过严格的数值等价验证和策略难度感知过滤机制，抑制模型对简单策略的偏好，确保最终数据集的高质量与均衡分布。最终，从35,878个初始样本中精炼出24,953个高质量数据条目，其中高难度策略占比达87.32%。

特点

Step-TP数据集的核心特点在于其提供了可验证的、原子化的步骤级监督，并融合了结构化的思维链推理。其独特性根植于四个设计原则。首先，它引入了一种token高效且可验证的中间表示LEIR，该表示通过将程序精简为循环结构与方程结构两大不可约语义组件，相较于TVM TIR平均减少约60%的token消耗，相较于CUDA减少约83%，从而在有限上下文长度内为复杂优化腾出推理空间。其次，该数据集实现了策略的原子化与可组合性，将错综复杂的优化轨迹分解为可解释的单步决策，涵盖图级、算子级、内存级和数学级四大类共43种不同策略。第三，结构化思维链监督将策略层面的高级语义与显式的IR到IR状态变迁映射相结合，使模型不仅学会“做什么”，更能理解“为何这么做”。最后，通过基于前置条件、参数修改和合成深度的难度感知过滤机制，精心平衡策略分布，有效避免了模型倾向于简单变换的捷径利用行为，系统性地提升了模型处理复杂优化逻辑的能力。

使用方法

Step-TP数据集的使用方法灵活多样，旨在赋能大语言模型成为可靠的张量程序优化智能体。在数据格式上，数据集提供了两种版本。一是字典格式的存档仓库，每条记录包含原始与变换后的LEIR、TIR及CUDA代码、应用策略及对应思维链。二是基于聊天模板的训练数据集，细分为单答案版本（一个原始LEIR对应一个具体变换）和多答案版本（一个原始LEIR对应多个不同变换），前者用于基础的单步变换学习，后者则鼓励模型进行多样化的优化推理，并为波束搜索等多步优化场景提供候选节点。在应用层面，研究者可将此数据集用于大语言模型的后训练监督微调。模型在微调后展现了卓越性能：在单步变换中，Qwen3-32B模型达到了92%的等价通过率；在多步优化中，结合波束搜索能实现平均42.90倍的加速比。模型能够自主地从43种策略中选择合适的变换，并组合成长序列，在H20-3e和A100等多种GPU架构上均表现出一致的有效性，证明了其强大的泛化能力。

背景与挑战

背景概述

在大规模深度学习模型迅猛发展的当下，张量程序的GPU高效执行已成为决定推理延迟、能耗与部署成本的关键瓶颈。传统编译器依赖手工调优或黑盒搜索，难以应对日益复杂的模型结构与庞大的组合优化空间。为此，香港大学、蚂蚁集团与中国科学技术大学的研究团队于2025年联合提出了Step-TP数据集，旨在为基于大语言模型的张量程序优化提供细粒度的监督信号。该数据集的核心研究问题在于：如何使LLM摆脱对最终优化结果的简单模仿，转而具备可解释、可验证的逐步推理能力，以在巨大的离散搜索空间中做出可靠的组合决策。Step-TP的提出填补了该领域后训练数据的空白，显著推动了LLM从粗粒度启发式生成向可信优化代理的范式转变。

当前挑战

张量程序优化面临的首要挑战在于其固有的迭代决策本质：每一次变换都需在保证语义等价的前提下，精确地对循环结构与代数方程进行组合式修改，而传统的端到端优化数据仅提供最终程序对，缺乏可验证的步骤级监督，导致LLM难以泛化至未见程序。此外，现有数据集采用的低级中间表示（如CUDA、TIR）包含大量与变换逻辑无关的语法噪声，严重浪费LLM有限的上下文窗口。构建过程中，研究者需确保策略分解的原子性与可组合性，同时设计精细的难度感知过滤机制以抑制模型对简单变换的偏好，最终在35,878条初始样本中仅保留24,953条经过严格数值等价验证的高质量数据，平衡了复杂推理与基础知识的覆盖。

常用场景

经典使用场景

Step-TP数据集专为基于大语言模型的张量程序优化而设计，其经典使用场景在于为模型提供细粒度、可验证的逐步监督信号。在深度学习编译器优化领域，传统方法往往仅提供端到端的优化程序对，缺乏中间状态的显式建模与可解释的推理链条。Step-TP通过引入一种高密度中间表示LEIR，将复杂的优化轨迹拆解为原子化、可组合的单步决策，并辅以结构化思维链推理，使大语言模型能够在一个封闭的推理循环中，基于当前程序状态进行可靠的多步优化决策，而非简单模仿最终输出结果。

衍生相关工作

Step-TP的提出催生了一系列衍生工作。其核心思想——逐步监督与结构化思维链推理——已被借鉴至更广泛的编译器与代码优化领域。例如，后续研究在此基础上探索了面向CUDA与Triton内核生成的多轮强化学习框架，以及基于推理图的跨模型优化知识迁移方法。此外，Step-TP所定义的LEIR中间表示与原子化策略分类体系，为构建统一的张量程序优化基准提供了标准化工具，促进了不同方法之间的可比较性评价。这些工作共同推动了大语言模型在程序优化场景中从启发式生成器向可靠推理代理的转化。

数据集最近研究

最新研究方向

Step-TP数据集的提出标志着大语言模型在张量程序优化领域从粗粒度的端到端监督迈向细粒度的步骤级推理。该研究通过设计一种高效且可验证的中间表示LEIR，将复杂的优化轨迹解耦为原子化的单步策略，并辅以结构化链式思维推理，从而赋予模型可靠的组合优化能力。这一前沿方向紧密关联当前深度学习编译器自动化与硬件高效执行的热点，尤其在FlashAttention等在线优化和长序列计算场景中，Step-TP提供的步骤级监督与策略难度感知过滤机制显著提升了多步优化的搜索效率与泛化性能。其影响在于为LLM在程序优化中的可解释性与泛化性奠定了数据基础，推动了从模式模仿到逻辑推理的范式转变。

相关研究论文

• Step-TP: A Grounded, Step-Level Dataset with Chain-of-Thought Reasoning for LLM-Guided Tensor Program Optimization

  香港大学; 蚂蚁集团; 中国科学技术大学 · 2026年