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

2026-05-25

AI Summary

This work addresses the challenge that large language models struggle to make reliable single-step decisions in tensor program optimization, primarily due to the absence of verifiable step-level supervision and interpretability in existing datasets. To bridge this gap, the authors introduce Step-TP, the first step-level optimization dataset enabling closed-loop reasoning. Step-TP leverages atomic and composable optimization strategies, coupled with verifiable intermediate representations, explicit state transitions, and a strategy filtering mechanism, to construct structured chain-of-thought reasoning paths. This approach not only ensures token efficiency and deterministic decompilation to TVM TIR but also substantially enhances the model’s reliability in making single-step decisions within complex optimization spaces.

Abstract

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.

Problem

• tensor program optimization
• large language models
• step-level supervision
• chain-of-thought reasoning
• combinatorial optimization

Innovation

• step-level supervision
• chain-of-thought reasoning
• tensor program optimization
• intermediate representation
• composable transformations