mlsanityData to 14 Oct 2026

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

arXiv 2605.25954 ↗First posted 25 May 2026cs.LG cs.AIcited by 0

On the map: LLM reasoning → Test-time compute and efficient thinking

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.

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.

Builds on

20 of 56 references.

References · 20

  1. 2023GPT-4 Technical Reportcited by 18,396
  2. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
  3. 2024DeepSeek-V3 Technical Reportcited by 811
  4. 2024Long-context LLMs Struggle with Long In-context Learningcited by 286
  5. 2018Online normalizer calculation for softmaxcited by 162
  6. 2018Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractionscited by 85
  7. 2025Kevin: Multi-Turn RL for Generating CUDA Kernelscited by 78
  8. 2025CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learningcited by 61
  9. 2023Large Language Models for Compiler Optimizationcited by 52
  10. 2025TileLang: A Composable Tiled Programming Model for AI Systemscited by 43
  11. 2024Meta Large Language Model Compiler: Foundation Models of Compiler Optimizationcited by 39
  12. 2022TopoOpt: Co-optimizing Network Topology and Parallelization Strategy for Distributed Training Jobscited by 30
8 more
  1. 2025STARK: Strategic Team of Agents for Refining Kernelscited by 25
  2. 2025TritonRL: Training LLMs to Think and Code Triton Without Cheatingcited by 24
  3. 2025ConCuR: Conciseness Makes State-of-the-Art Kernel Generationcited by 13
  4. 2025CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learningcited by 13
  5. 2025From Large to Small: Transferring CUDA Optimization Expertise via Reasoning Graphcited by 10
  6. 2023Operator Fusion in XLA: Analysis and Evaluationcited by 9
  7. 2026GAC: Stabilizing Asynchronous RL Training for LLMs via Gradient Alignment Controlcited by 5
  8. 2025LOOPerSet: A Large-Scale Dataset for Data-Driven Polyhedral Compiler Optimizationcited by 1

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025REASONING COMPILER: LLM-Guided Optimizations for Efficient Model Serving
  2. 2026Toward Compiler World Models: Learning Latent Dynamics for Efficient Tensor Program Search
  3. 2026TCL: Enabling Fast and Efficient Cross-Hardware Tensor Program Optimization via Continual Learning
  4. 2025Step-Opt: Boosting Optimization Modeling in LLMs through Iterative Data Synthesis and Structured Validation
  5. 2026Prism: Symbolic Superoptimization of Tensor Programs
  6. 2024Towards LLM-guided Efficient and Interpretable Multi-linear Tensor Network Rank Selection
  7. 2025ReProbe: Efficient Test-Time Scaling of Multi-Step Reasoning by Probing Internal States of Large Language Models
  8. 2023Let's reward step by step: Step-Level reward model as the Navigators for Reasoning