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

arXiv: https://arxiv.org/abs/2605.25954 · HTML: https://mlsanity.com/paper/2605.25954/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 20 of its 56 references that are in the map.
- First posted on arXiv 25 May 2026.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 20 of 20)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [DeepSeek-V3 Technical Report](https://mlsanity.com/paper/2412.19437/) (2024)
- [Long-context LLMs Struggle with Long In-context Learning](https://mlsanity.com/paper/2404.02060/) (2024)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [Kevin: Multi-Turn RL for Generating CUDA Kernels](https://mlsanity.com/paper/2507.11948/) (2025)
- [CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning](https://mlsanity.com/paper/2507.14111/) (2025)
- [Large Language Models for Compiler Optimization](https://mlsanity.com/paper/2309.07062/) (2023)
- [TileLang: A Composable Tiled Programming Model for AI Systems](https://mlsanity.com/paper/2504.17577/) (2025)
- [Meta Large Language Model Compiler: Foundation Models of Compiler Optimization](https://mlsanity.com/paper/2407.02524/) (2024)
- [TopoOpt: Co-optimizing Network Topology and Parallelization Strategy for Distributed Training Jobs](https://mlsanity.com/paper/2202.00433/) (2022)
- [STARK: Strategic Team of Agents for Refining Kernels](https://mlsanity.com/paper/2510.16996/) (2025)
- [TritonRL: Training LLMs to Think and Code Triton Without Cheating](https://mlsanity.com/paper/2510.17891/) (2025)
- [ConCuR: Conciseness Makes State-of-the-Art Kernel Generation](https://mlsanity.com/paper/2510.07356/) (2025)
- [CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning](https://mlsanity.com/paper/2512.02551/) (2025)
- [From Large to Small: Transferring CUDA Optimization Expertise via Reasoning Graph](https://mlsanity.com/paper/2510.19873/) (2025)
- [Operator Fusion in XLA: Analysis and Evaluation](https://mlsanity.com/paper/2301.13062/) (2023)
- [GAC: Stabilizing Asynchronous RL Training for LLMs via Gradient Alignment Control](https://mlsanity.com/paper/2603.01501/) (2026)
- [LOOPerSet: A Large-Scale Dataset for Data-Driven Polyhedral Compiler Optimization](https://mlsanity.com/paper/2510.10209/) (2025)

## Abstract (arXiv)

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.