🤖 machine learning

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

Dit artikel introduceert Step-TP, een nieuw post-training dataset dat LLM-gestuurde tensorprogramma-optimalisatie verbetert door grondige, atomaire supervisie op stapniveau te bieden met gestructureerd chain-of-thought-redeneren en een token-efficiënte tussenrepresentatie om betrouwbare meerstapsbesluitvorming mogelijk te maken.

Oorspronkelijke auteurs: Mengfan Liu, Da Zheng, Junwei Su, Chuan Wu

Gepubliceerd 2026-05-26

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Mengfan Liu, Da Zheng, Junwei Su, Chuan Wu

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Het Grote Plaatje: Een Chef Koken Leren Sneller

Stel je voor dat je een briljante chef hebt (het Groot Taalmodel of LLM) die geweldig is in het lezen van recepten en het begrijpen van smaken. Deze chef heeft echter nog nooit in een hoge-snelheids, industriële keuken gekookt. Als je hen vraagt een recept te optimaliseren om sneller te koken, raden ze vaak op basis van hoe het eindgerecht eruitziet, in plaats van de specifieke stappen te begrijpen die nodig zijn om daar te komen. Ze kunnen zeggen: "Voeg gewoon meer hitte toe!" zonder te beseffen dat dit het eten zal verbranden.

In de wereld van computers is "koken" het uitvoeren van complexe wiskundige programma's (zogenaamde tensorprogramma's) op krachtige grafische kaarten (GPU's). Het sneller laten draaien van deze programma's is ontzettend moeilijk, omdat er miljoenen kleine manieren zijn om de code te herschikken, en één verkeerde zet het hele geheel kan doen crashen.

Dit paper introduceert Step-TP, een nieuw "kookboek" dat de chef precies leert hoe deze wijzigingen stap-voor-stap moeten worden aangebracht, met een duidelijke uitleg van waarom elke stap werkt.

Het Probleem: Het "Black Box" Kookboek

Voordat dit paper verscheen, hadden de datasets die werden gebruikt om deze AI-chefs te trainen drie hoofdproblemen:

Alleen het Resultaat Getoond: De meeste oude kookboeken toonden alleen het "Voor"-recept en het "Na"-recept. Ze toonden niet de stappen ertussen. Het was alsof je een rauwe kip en een geroosterde kip liet zien, maar de kruiden, de oventemperatuur en de timing verborg. De AI zou dan alleen de look van het eindgerecht memoriseren in plaats van de kooktechniek te leren.
Te Veel Rommel: De recepten waren geschreven in een zeer rommelige, technische taal (zoals ruwe computercode) vol met onnodige details. Het was alsof je probeerde een recept te lezen waarbij elke instructie in een ander lettertype was geschreven, met voetnoten over het merk van het fornuis. Dit maakte het moeilijk voor de AI om zich te focussen op de daadwerkelijke kooklogica.
Geen "Waarom": De AI leerde het redeneren niet. Als de chef een nieuw type groente zag, kon hij niet uitzoeken hoe hij het moest koken, omdat hij alleen specifieke gerechten had gememoriseerd, niet de algemene regels van koken.

De Oplossing: Step-TP

De auteurs creëerden een nieuwe dataset genaamd Step-TP die deze problemen oplost. Hier is hoe ze dat deden, met drie hoofdingrediënten:

1. Een Schoner Taal (LEIR)

De auteurs bedachten een nieuwe manier om de "recepten" te schrijven, genaamd LEIR (Loop-Equation Intermediate Representation).

De Analogie: Stel je voor dat je een rommelig, handgeschreven briefje vol krabbels en koffievlekken vertaalt naar een schone, getypte lijst met opsommingstekens.
Wat het doet: LEIR haalt alle computer-"boilerplate" (het saaie, herhalende gedoe) weg en laat alleen de essentiële logica over: de lussen (de herhalende stappen) en de vergelijkingen (de wiskunde). Dit maakt het recept veel korter en gemakkelijker voor de AI om te lezen en te begrijpen.

2. Stap-voor-stap Training (Atomaire Stappen)

In plaats van de AI de hele transformatie van begin tot eind te tonen, breekt Step-TP het op in tiny, enkele stappen.

De Analogie: In plaats van te zeggen "Zet de rauwe kip om in een geroosterde kip", zegt de dataset:
- Stap 1: Kruid de kip.
- Stap 2: Verwarm de oven voor.
- Stap 3: Plaats de kip in de oven.
Waarom het belangrijk is: Dit leert de AI om één kleine, veilige beslissing per keer te nemen. Het leert dat "Stap 1" leidt tot een specifieke staat, die vervolgens "Stap 2" mogelijk maakt. Dit voorkomt dat de AI enorme, risicovolle sprongen maakt die het programma kunnen doen crashen.

3. Redeneren in Gedachtenketens (Het "Waarom")

Elke stap in de dataset wordt vergezeld van een "Chain-of-Thought" (CoT) uitleg.

De Analogie: Het is als een kookshow waarbij de chef niet alleen de actie uitvoert, maar uitlegt: "Ik draai de kip nu om omdat de onderkant goudbruin is."
Wat het doet: De dataset vertelt de AI expliciet waarom een specifieke wijziging is aangebracht (bijvoorbeeld: "We herschikken deze lussen om beter gebruik te maken van het geheugen"). Dit helpt de AI om de onderliggende logica te leren, zodat het deze regels kan toepassen op nieuwe, onbekende problemen.

De Resultaten: Een Meesterchef

Het paper testte deze nieuwe dataset op verschillende maten van AI-modellen. Hier is wat ze vonden:

Efficiëntie: Omdat de "recepten" (LEIR) zo schoon waren, kon de AI ze veel sneller verwerken, met veel minder "tokens" (woorden) dan voorheen. Het was alsof je overschakelt van het lezen van een 100-pagina's tellende handleiding naar het lezen van een 10-pagina's tellende cheat sheet.
Nauwkeurigheid: De AI werd veel beter in het maken van wijzigingen die daadwerkelijk werkten. In tests konden de modellen programma's met succes transformeren en ze in meer dan 90% van de gevallen correct laten draaien, zelfs voor zeer complexe taken.
Lange Reizen: De AI kon lange ketens van optimalisaties aan. In plaats van slechts één kleine wijziging te maken, kon het een reeks van 10 of 15 stappen plannen om een programma honderden keren sneller te laten draaien. Het was alsof de chef leert om een heel banketmenu te plannen in plaats van alleen een bijgerecht te koken.

Samenvatting

Step-TP is een gespecialiseerde trainingsdataset die AI-modellen leert hoe ze computerprogramma's kunnen optimaliseren door:

Het gebruik van een schone, vereenvoudigde taal (LEIR) om rommel te verwijderen.
Het opsplitsen van complexe taken in kleine, hanteerbare stappen.
Het verstrekken van uitleg voor elke stap zodat de AI de logica begrijpt, niet alleen het patroon.

Het resultaat is een AI die kan optreden als een betrouwbare expert-engineer, die nauwkeurige, meerstapsverbeteringen aanbrengt in softwareprestaties zonder iets te breken.

Technische Samenvatting: Step-TP

Probleemstelling

Het optimaliseren van de uitvoeringsefficiëntie van tensorprogramma's op GPU's is een kritieke uitdaging voor deep learning, vooral naarmate modelarchitecturen complexer worden. Hoewel Large Language Models (LLM's) sterke redeneercapaciteiten hebben aangetoond, wordt hun toepassing op de optimalisatie van tensorprogramma's gehinderd door het ontbreken van geschikte trainingsdata. Bestaande datasets bieden doorgaans alleen end-to-end geoptimaliseerde programmaparen met behulp van omvangrijke, token-inefficiënte representaties (zoals ruwe CUDA of TVM TIR). Deze datasets missen verifieerbare supervisie op stap-niveau, waardoor LLM's gedwongen worden om te vertrouwen op imitatie van uitkomsten of oppervlakkig patroonherkenning in plaats van het leren van het onderliggende besluitvormingsproces. Bijgevolg worstelen LLM's om betrouwbare beslissingen op één enkele stap te nemen binnen de enorme, combinatorische optimalisatieruimtes die vereist zijn voor effectieve afstemming van tensorprogramma's.

Methodologie

Om deze beperkingen aan te pakken, introduceren de auteurs Step-TP, een post-training dataset ontworpen om onderbouwde, atomaire supervisie op stap-niveau te bieden met gestructureerde Chain-of-Thought (CoT)-redenering. De methodologie is gebaseerd op vier kernontwerpprincipes:

1. Loop-Equation Intermediate Representation (LEIR)

De auteurs stellen LEIR voor, een token-efficiënte en verifieerbare intermediaire representatie die is toegespitst op LLM-gebaseerde optimalisatie. In tegenstelling tot bestaande IR's (CUDA, TVM TIR) die vol zitten met compiler-basiscode en hardware-specifieke details, reduceert LEIR tensorprogramma's tot twee onreduceerbare semantische componenten:

Lusstructuur: Definieert het iteratieruimte met behulp van compacte symbolen voor seriële ( $L$ ), parallelle ( $P$ ), vectoriële ( $V$ ), ontvouwde ( $U$ ) en thread/block-gebonden ( $B$ ) lussen.
Vergelijkingsstructuur: Specificeert de algebraïsche berekening bij elke iteratie met behulp van element-voor-element operatoren, tensorvariabelen (met expliciete datatypen en memory hierarchy) en scheidingstekens.
LEIR maakt gebruik van impliciete initialisatie voor veelvoorkomende reductiepatronen en gebruikt een op LaTeX gebaseerde syntaxis om het voorafgaande kennis van LLM's te benutten, wat het tokenverbruik aanzienlijk verlaagt in vergelijking met CUDA en TIR.

2. Constructie van Dataset op Stap-niveau

De dataset wordt samengesteld via een meerfasige pijplijn:

Brongeneratie: PyTorch-programma's worden gegenereerd vanuit enkele operatoren en samengestelde architecturen (bijvoorbeeld attention-modules), die diverse vormen en datatypen bestrijken.
Transformatie gestuurd door één-stapsstrategie: In plaats van end-to-end mapping, decomposeert het systeem complexe optimalisatietrajecten in atomaire, composable strategieën. Voor elke oorspronkelijke LEIR worden toepasbare strategieën geïdentificeerd op basis van precondities, individueel toegepast, en de resulterende getransformeerde LEIR wordt gegenereerd.
Gestructureerde CoT-supervisie: Elke transformatie wordt begeleid door een redeneerspore die de rationale van de strategie uitlegt en de specifieke structurele wijzigingen documenteert (IR-naar-IR staatsovergangen).
Verificatie en filtering: Getransformeerde programma's worden verlaagd naar TVM TIR en uitgevoerd om numerieke equivalentie te verifiëren. Een strategie-filtermechanisme wordt toegepast om de datasetverdeling in evenwicht te brengen, waardoor LLM's worden beschermd tegen overfitting op eenvoudige strategieën door "makkelijke" transformaties minder zwaar te wegen en "moeilijke" transformaties (zoals online softmax, FlashAttention-varianten) te behouden.

3. Strategie Taxonomie

De dataset bestrijkt 43 verschillende optimalisatiestrategieën die zijn gecategoriseerd in vier niveaus:

Graf-niveau: bijv. herschikking van expressies, fusie van operatoren.
Operator-niveau: bijv. herschikking van lussen, tiling van lussen, splitsing van lussen.
Geheugenniveau: bijv. cache lezen/schrijven, instellen van opslagscope.
Wiskundig niveau: bijv. online softmax, factorisatie, annulering.

Belangrijkste Bijdragen

Ontwerp van LEIR: Een hoogdichte, verifieerbare IR die een compacte representatie van intermediaire programmastaten mogelijk maakt en de naadloze toepassing van atomaire optimalisatiestrategieën, wat de tokenvoetafdruk met ongeveer 60% verkleint in vergelijking met TIR en met 83% in vergelijking met CUDA.
Constructie van Step-TP: De eerste post-training dataset voor LLM-gebaseerde tensoroptimalisatie die onderbouwde, atomaire supervisie op stap-niveau biedt met gestructureerde CoT. Het transformeert de optimalisatiezoekruimte naar interpreteerbare beslissingen op één enkele stap.
Empirische Validatie: Uitgebreide experimenten die aantonen dat Step-TP modellen in staat stelt om uitvoerbare, onderbouwde transformaties te genereren over diverse strategieën heen en lange-horizon optimalisatietrajecten te navigeren.

Resultaten

De auteurs evalueerden modellen (Qwen3-serie) die op Step-TP waren getraind voor taken met optimalisatie op één stap en meerdere stappen:

Token-efficiëntie: LEIR vereiste consequent minder tokens dan TIR en CUDA, waardoor complexe optimalisatieprocessen binnen praktische contextvensters pasten.
Transformaties op één stap: Modellen getraind op Step-TP behaalden hoge fideliteit bij het genereren van uitvoerbare en numeriek equivalente transformaties. Bijvoorbeeld, het Qwen3-32B-model behaalde een "Equal Pass"-percentage van 92% (bit-identieke uitvoer) op vastgehouden testgevallen. De modellen dekten succesvol een divers scala aan strategiecategorieën, evenredig aan de datasetverdeling.
Optimalisatie op meerdere stappen: Met behulp van zoekalgoritmen (Greedy, Beam, BFS, DFS, MCTS, Chain Search) behaalden de modellen aanzienlijke versnellingen in runtime. Beam Search bereikte bijvoorbeeld een gemiddelde versnelling van 42,90× en een piek van 561,82×. De modellen slaagden erin sequenties van 2 tot 10 verschillende strategietypen te combineren om lange-horizon trajecten te optimaliseren.
Generalisatie: De optimalisatiekennis die door Step-TP is vastgelegd, bleef effectief over verschillende GPU-architecturen (getest op NVIDIA H20-3e en A100).
Ablatiestudies: De resultaten bevestigden dat de combinatie van LEIR, gestructureerde CoT en strategie-filtering cruciaal is. Het verwijderen van CoT verlaagde de efficiëntie van zoekopdrachten op meerdere stappen aanzienlijk (van 0,87 naar 0,10), en het verwijderen van strategie-filtering leidde tot een bias naar eenvoudige transformaties.

Betekenis

Het artikel beweert dat Step-TP een kritieke kloof in het veld aanpakt door LLM's te transformeren van grove heuristische generators naar betrouwbare optimalisatieagenten die in staat zijn om op schaal te redeneren over tensorprogramma's. Door verifieerbare supervisie op stap-niveau te bieden, maakt de dataset data-efficiënt leren en stabiel redeneren mogelijk, waardoor modellen kunnen generaliseren over werklasten en hardware-backends zonder afhankelijk te zijn van dure versterkende leer of zoekgebaseerde fine-tuning. Het werk legt de basis voor flexibelere en data-efficiëntere frameworks voor tensoroptimalisatie die systematisch kunnen redeneren over GPU-uitvoergedrag.