🤖 machine learning

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

This paper introduces Step-TP, a novel post-training dataset that enhances LLM-guided tensor program optimization by providing grounded, atomic, step-level supervision with structured chain-of-thought reasoning and a token-efficient intermediate representation to enable reliable multi-step decision-making.

Original authors: Mengfan Liu, Da Zheng, Junwei Su, Chuan Wu

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mengfan Liu, Da Zheng, Junwei Su, Chuan Wu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Teaching a Chef to Cook Faster

Imagine you have a brilliant chef (the Large Language Model or LLM) who is amazing at reading recipes and understanding flavors. However, this chef has never actually cooked in a high-speed, industrial kitchen before. When you ask them to optimize a recipe to cook faster, they often guess based on what the final dish looks like, rather than understanding the specific steps needed to get there. They might say, "Just add more heat!" without realizing that will burn the food.

In the world of computers, "cooking" is running complex math programs (called tensor programs) on powerful graphics cards (GPUs). Making these programs run faster is incredibly difficult because there are millions of tiny ways to rearrange the code, and one wrong move can break the whole thing.

This paper introduces Step-TP, a new "cookbook" designed to teach the chef exactly how to make these changes, step-by-step, with a clear explanation of why each step works.

The Problem: The "Black Box" Cookbook

Before this paper, the datasets used to train these AI chefs had three main problems:

Only Showed the Result: Most old cookbooks just showed the "Before" recipe and the "After" recipe. They didn't show the steps in between. It was like showing a raw chicken and a roasted chicken, but hiding the seasoning, the oven temperature, and the timing. The AI would just memorize the look of the final dish instead of learning the cooking technique.
Too Much Clutter: The recipes were written in a very messy, technical language (like raw computer code) that was full of unnecessary details. It was like trying to read a recipe where every instruction was written in a different font, with footnotes about the brand of the stove. This made it hard for the AI to focus on the actual cooking logic.
No "Why": The AI didn't learn the reasoning. If the chef saw a new type of vegetable, they couldn't figure out how to cook it because they only memorized specific dishes, not the general rules of cooking.

The Solution: Step-TP

The authors created a new dataset called Step-TP that fixes these issues. Here is how they did it, using three main ingredients:

1. A Cleaner Language (LEIR)

The authors invented a new way to write the "recipes" called LEIR (Loop-Equation Intermediate Representation).

The Analogy: Imagine translating a messy, handwritten note full of scribbles and coffee stains into a clean, typed list with bullet points.
What it does: LEIR strips away all the computer "boilerplate" (the boring, repetitive stuff) and leaves only the essential logic: the loops (the repeating steps) and the equations (the math). This makes the recipe much shorter and easier for the AI to read and understand.

2. Step-by-Step Training (Atomic Steps)

Instead of showing the AI the whole transformation from start to finish, Step-TP breaks it down into tiny, single steps.

The Analogy: Instead of saying "Turn the raw chicken into a roasted chicken," the dataset says:
- Step 1: Season the chicken.
- Step 2: Preheat the oven.
- Step 3: Place the chicken in the oven.
Why it matters: This teaches the AI to make one small, safe decision at a time. It learns that "Step 1" leads to a specific state, which then allows "Step 2." This prevents the AI from making giant, risky jumps that might break the program.

3. Chain-of-Thought Reasoning (The "Why")

Every step in the dataset comes with a "Chain-of-Thought" (CoT) explanation.

The Analogy: It's like a cooking show where the chef doesn't just do the action, but explains, "I am flipping the chicken now because the bottom is golden brown."
What it does: The dataset explicitly tells the AI why a specific change was made (e.g., "We are reordering these loops to make better use of the memory"). This helps the AI learn the underlying logic so it can apply these rules to new, unseen problems.

The Results: A Master Chef

The paper tested this new dataset on different sizes of AI models. Here is what they found:

Efficiency: Because the "recipes" (LEIR) were so clean, the AI could process them much faster, using far fewer "tokens" (words) than before. It was like switching from reading a 100-page manual to reading a 10-page cheat sheet.
Accuracy: The AI became much better at making changes that actually worked. In tests, the models could successfully transform programs and keep them running correctly over 90% of the time, even for very complex tasks.
Long Journeys: The AI could handle long chains of optimizations. Instead of just making one small change, it could plan a sequence of 10 or 15 steps to make a program run hundreds of times faster. It was like the chef learning to plan a whole banquet menu rather than just cooking a single side dish.

Summary

Step-TP is a specialized training dataset that teaches AI models how to optimize computer programs by:

Using a clean, simplified language (LEIR) to remove clutter.
Breaking complex tasks into small, manageable steps.
Providing explanations for every step so the AI understands the logic, not just the pattern.

The result is an AI that can act like a reliable expert engineer, making precise, multi-step improvements to software performance without breaking anything.

Technical Summary: Step-TP

Problem Statement

Optimizing the execution efficiency of tensor programs on GPUs is a critical challenge for deep learning, particularly as model architectures grow in complexity. While Large Language Models (LLMs) have demonstrated strong reasoning capabilities, their application to tensor program optimization has been hindered by the lack of suitable training data. Existing datasets typically provide only end-to-end optimized program pairs using verbose, token-inefficient representations (such as raw CUDA or TVM TIR). These datasets lack verifiable, step-level supervision, forcing LLMs to rely on outcome imitation or superficial pattern matching rather than learning the underlying decision-making process. Consequently, LLMs struggle to make reliable single-step decisions within the vast, combinatorial optimization spaces required for effective tensor program tuning.

Methodology

To address these limitations, the authors introduce Step-TP, a post-training dataset designed to provide grounded, atomic, step-level supervision with structured Chain-of-Thought (CoT) reasoning. The methodology is built upon four core design principles:

1. Loop-Equation Intermediate Representation (LEIR)

The authors propose LEIR, a token-efficient and verifiable intermediate representation tailored for LLM-based optimization. Unlike existing IRs (CUDA, TVM TIR) that are cluttered with compiler boilerplate and hardware-specific details, LEIR distills tensor programs into two irreducible semantic components:

Loop Structure: Defines the iteration space using compact symbols for serial ( $L$ ), parallel ( $P$ ), vectorized ( $V$ ), unrolled ( $U$ ), and thread/block-binding ( $B$ ) loops.
Equation Structure: Specifies the algebraic computation at each iteration using element-wise operators, tensor variables (with explicit data types and memory hierarchy), and delimiters.
LEIR employs implicit initialization for common reduction patterns and uses a LaTeX-based syntax to leverage LLMs' prior knowledge, significantly reducing token consumption compared to CUDA and TIR.

2. Step-Level Dataset Construction

The dataset is constructed via a multi-stage pipeline:

Source Generation: PyTorch programs are generated from single operators and composite architectures (e.g., attention modules), covering diverse shapes and data types.
One-Step Strategy-Driven Transformation: Instead of end-to-end mapping, the system decomposes complex optimization trajectories into atomic, composable strategies. For each original LEIR, applicable strategies are identified based on preconditions, applied individually, and the resulting transformed LEIR is generated.
Structured CoT Supervision: Each transformation is accompanied by a reasoning trace that explains the strategy rationale and documents the specific structural changes (IR-to-IR state transitions).
Verification and Filtering: Transformed programs are lowered to TVM TIR and executed to verify numerical equivalence. A strategy filtering mechanism is applied to balance the dataset distribution, preventing LLMs from overfitting to simple strategies by down-weighting "easy" transformations and retaining "difficult" ones (e.g., online softmax, FlashAttention variants).

3. Strategy Taxonomy

The dataset covers 43 distinct optimization strategies categorized into four levels:

Graph Level: e.g., expression reorder, operator fusion.
Operator Level: e.g., loop reorder, loop tiling, loop split.
Memory Level: e.g., cache read/write, set storage scope.
Mathematical Level: e.g., online softmax, factorization, cancellation.

Key Contributions

Design of LEIR: A high-density, verifiable IR that enables compact representation of intermediate program states and seamless application of atomic optimization strategies, reducing token footprints by approximately 60% compared to TIR and 83% compared to CUDA.
Construction of Step-TP: The first post-training dataset for LLM-based tensor optimization that provides grounded, atomic, step-level supervision with structured CoT. It transforms the optimization search space into interpretable single-step decisions.
Empirical Validation: Extensive experiments demonstrating that Step-TP enables models to generate executable, grounded transformations across diverse strategies and navigate long-horizon optimization trajectories.

Results

The authors evaluated models (Qwen3 series) trained on Step-TP across single-step and multi-step optimization tasks:

Token Efficiency: LEIR consistently required fewer tokens than TIR and CUDA, allowing complex optimization processes to fit within practical context windows.
Single-Step Transformations: Models trained on Step-TP achieved high fidelity in generating executable and numerically equivalent transformations. For instance, the Qwen3-32B model achieved a 92% "Equal Pass" rate (bitwise identical outputs) on held-out test cases. The models successfully covered a diverse range of strategy categories proportional to the dataset distribution.
Multi-Step Optimization: Using search algorithms (Greedy, Beam, BFS, DFS, MCTS, Chain Search), the models achieved substantial runtime speedups. Beam Search, for example, achieved an average speedup of 42.90× and a peak of 561.82×. The models successfully composed sequences of 2 to 10 distinct strategy types to optimize long-horizon trajectories.
Generalization: The optimization knowledge captured by Step-TP remained effective across different GPU architectures (tested on NVIDIA H20-3e and A100).
Ablation Studies: The results confirmed that the combination of LEIR, structured CoT, and strategy filtering is critical. Removing CoT significantly reduced multi-step search efficiency (from 0.87 to 0.10), and removing strategy filtering led to a bias toward simple transformations.

Significance

The paper claims that Step-TP addresses a critical gap in the field by transforming LLMs from coarse heuristic generators into reliable optimization agents capable of reasoning about tensor programs at scale. By providing verifiable, step-level supervision, the dataset enables data-efficient learning and stable reasoning, allowing models to generalize across workloads and hardware backends without relying on expensive reinforcement learning or search-based fine-tuning. The work establishes a foundation for more flexible and data-efficient tensor optimization frameworks that can systematically reason about GPU execution behavior.