Step-TP: 체인 오브 소트 추론을 활용한 LLM 기반 텐서 프로그램 최적화를 위한 단계별 데이터셋

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

대규모 언어 모델(LLM)은 강력한 추론 능력을 가지고 있지만, 정확하고 조합 가능한 변환 결정을 필요로 하기 때문에 텐서 프로그램의 실행 효율성을 최적화하는 것은 여전히 어려운 과제입니다. 최근 LLM 기반 접근 방식에서는 텐서 프로그램 최적화를 반복적인 의사 결정 과정으로 정의하지만, 기존 데이터셋은 토큰 사용량을 줄인 표현 방식을 사용하여 전체적으로 최적화된 프로그램 쌍만을 제공하며, 검증 가능한 단계별 지도 학습 및 해석 가능성이 부족합니다. 그 결과, LLM은 큰 조합 최적화 공간에서 신뢰할 수 있는 단일 단계 결정을 내리는 데 어려움을 겪습니다. 우리는 체인 오브 소트(CoT) 추론을 통해 구조화된 단계별 지침을 제공하는 텐서 프로그램 최적화를 위한 후처리 데이터셋인 Step-TP를 소개합니다. Step-TP는 중간 프로그램 상태에 대한 폐쇄적인 추론 루프를 형성하여, 결과 모방이 아닌 신뢰할 수 있는 다단계 최적화를 가능하게 합니다. 이 데이터셋은 다음과 같은 네 가지 원칙에 따라 설계되었습니다: (i) TVM TIR로 결정적으로 변환되는 토큰 사용량이 효율적이고 검증 가능한 중간 표현(IR); (ii) 복잡한 경로를 해석 가능한 단일 단계 결정으로 분해하는 원자적이고 조합 가능한 최적화 전략; (iii) 명시적인 IR-to-IR 상태 전환과 함께 구조화된 CoT 지침; 그리고 (iv) 폭넓은 적용 범위를 유지하면서 해킹을 방지하기 위한 전략 필터링. 데이터셋 및 구현 코드는 다음 GitHub 링크에서 확인할 수 있습니다: https://github.com/LIUMENGFAN-gif/StepTP.

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.