🤖 machine learning

Toward Compiler World Models: Learning Latent Dynamics for Efficient Tensor Program Search

Dieses Paper schlägt einen von Weltmodellen inspirierten Evaluator vor, der die latenten Dynamiken von Scheduling-Aktionen lernt, um Tensorprogrammkandidaten effizient zu ranken, wodurch signifikante Latenzverbesserungen gegenüber bestehenden Auto-Schedulern wie Ansor erzielt und die Anzahl der erforderlichen Messungen drastisch reduziert werden.

Ursprüngliche Autoren: Haolin Pan, Lianghong Huang, Xvlin Zhou, Mingjie Xing, Yanjun Wu

Veröffentlicht 2026-06-09

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Haolin Pan, Lianghong Huang, Xvlin Zhou, Mingjie Xing, Yanjun Wu

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen, die absolut schnellste Route zu finden, um von Ihrem Zuhause zum Haus eines Freundes zu fahren. Sie haben eine Karte, aber die Verkehrsbedingungen ändern sich ständig, und es gibt Millionen von möglichen Routen, die Sie nehmen könnten.

In der Welt der Informatik, speziell im Bereich des Maschinellen Lernens, ist dies genau das, was ein „Compiler“ tut. Er versucht, den effizientesten Weg zu finden, damit ein Computer komplexe mathematische Aufgaben (genannt „Tensor-Programme“) ausführt. Das Problem ist, dass es so viele Möglichkeiten gibt, den Code zu schreiben, dass das Überprüfen jeder einzelnen Variante durch das tatsächliche Ausführen auf dem Computer unglaublich langsam und teuer wäre. Es ist, als würde man versuchen, jede einzelne mögliche Route zu fahren, nur um zu sehen, welche am schnellsten ist; man würde vorher ohne Benzin ausgehen.

Der alte Weg: Eine Momentaufnahme machen

Zuvor verhielten sich Computerprogramme, die versuchten, solche Probleme zu lösen (genannt „Auto-Scheduler“), wie ein Fotograf, der eine Momentaufnahme des Zielortes macht. Sie betrachteten den fertigen Code, schätzten dessen Geschwindigkeit ein und entschieden dann, ob er gut war.

Die Arbeit argumentiert, dass dies eine schlechte Idee ist, weil:

Der Weg ignoriert wird: Es wird nicht verstanden, wie der Code dorthin gelangt ist. Zwei verschiedene Routen können am selben Punkt enden, aber eine könnte eine sanfte Autobahn gewesen sein, während die andere eine holprige Schotterstraße war. Die Momentaufnahme sieht gleich aus, aber das Erlebnis (und die Geschwindigkeit) ist unterschiedlich.
Es durch kleine Details verwirrt wird: Wenn man ein paar Wörter im Code ändert, die die Funktionsweise eigentlich nicht beeinflussen, könnte das alte System denken, es handele sich um eine völlig andere und schlechtere Route.

Die neue Idee: Ein „Weltmodell“ (Der GPS-Simulator)

Die Autoren schlagen einen neuen Ansatz vor, der von Weltmodellen inspiriert ist. Denken Sie hierbei nicht an einen Fotografen, sondern an einen hochmodernen GPS-Simulator.

Anstatt nur das Endziel zu betrachten, simuliert dieses neue System die gesamte Reise Schritt für Schritt in seinem „Geist“ (einem mathematischen Raum, der „Latent Space“ genannt wird).

So funktioniert es, unter Verwendung einer Kochanalogie:

Die Zutaten (Anfangszustand): Man beginnt mit einem rohen Rezept (dem unoptimierten Code).
Die Handgriffe des Kochs (Aktionen): Der Compiler trifft Entscheidungen, wie zum Beispiel: „Zwiebeln hacken“, „5 Minuten anbraten“ oder „Salz hinzufügen“.
Der alte Weg: Der Koch betrachtet das fertige Gericht und schätzt: „Hm, das schmeckt ganz okay.“
Der neue Weg (Weltmodell): Der Koch besitzt eine mentale Simulation. Er stellt sich vor: „Wenn ich die Zwiebeln erst hacke und dann anbrate, wird die Textur X sein. Wenn ich erst anbrate und dann hacke, wird die Textur Y sein.“ Er simuliert den Prozess des Kochens in seinem Kopf, um den endgültigen Geschmack vorherzusagen, ohne das ganze Essen zuerst tatsächlich kochen zu müssen.

Wie sie es gebaut haben

Die Forscher haben ein System mit drei Teilen entwickelt:

Der Übersetzer (Encoder): Er verwandelt den unordentlichen Computercode in einen sauberen, mathematischen „Gedanken“ (einen Vektor), den der Computer leicht verstehen kann.
Der Simulator (Transition Model): Dies ist die Kerninnovation. Er nimmt den „Gedanken“ des aktuellen Codes und wendet die „Handgriffe des Kochs“ (Scheduling-Aktionen) nacheinander an. Er sagt voraus, wie der Code nach jedem Schritt aussehen wird – und das alles innerhalb des Computerspeichers, ohne den Code tatsächlich auszuführen.
Der Richter (Ranking Model): Sobald die Simulation abgeschlossen ist, betrachtet der Richter das vorhergesagte Endergebnis und sagt: „Diese Route ist wahrscheinlich die schnellste“ oder „Jene sieht langsam aus“.

Die Ergebnisse

Sie testeten dies auf zwei Arten von Computern: einer leistungsstarken CPU (Intel Xeon) und einer High-End-Grafikkarte (NVIDIA RTX 4090).

Schnellere Ergebnisse: Sie fanden deutlich bessere Code-Schedules schneller als die bisher beste Methode (genannt Ansor).
Weniger Aufwand: Sie erreichten Ergebnisse, die genauso gut waren wie die der alten Methode, mussten aber 10-mal weniger „Testfahrten“ (Messungen) durchführen.
Reale Geschwindigkeit: Wenn sie dies zur Ausführung echter KI-Modelle (wie Bilderkennung oder Sprachmodelle) verwendeten, liefen die Programme 4- bis 5-mal schneller als Standardversionen und in einigen Fällen sogar bis zu 58-mal schneller.

Das Fazit

Die Arbeit behauptet, dass wir, indem wir den Computer lehren, den Prozess der Optimierung (den Weg) anstatt nur das Ergebnis (die Momentaufnahme) zu verstehen, die schnellsten Codes viel effizienter finden können. Es ist wie ein GPS, das den Verkehr in seinem Kopf simuliert, um die beste Route zu finden, anstatt nur basierend auf einem Foto des Ziels zu raten.

Im Paper erwähnte Einschränkungen:

Dieses System ist ein „Richter“, der hilft, die beste Route auszuwählen; es erfindet die Routen nicht selbst. Wenn die Suchmaschine von vornherein keine guten Routen vorschlägt, kann der Richter das nicht korrigieren.
Wenn die „Reise“ extrem lang und komplex ist, könnten die kleinen Fehler in der Simulation im Kopf des Computers sich aufsummieren und die Vorhersage ungenauer machen.
Es ist darauf ausgelegt, Optionen miteinander zu vergleichen (welches ist schneller?), nicht um die exakte Zeit bis auf die Millisekunde vorherzusagen.

Technisches Resümee: Auf Compiler-Weltmodelle zugehen

Problemstellung

Die Optimierung von Tensorprogrammen ist für moderne maschinelle Lernsysteme essenziell, doch die Identifizierung des effizientesten Programms wird durch einen massiven kombinatorischen Suchraum erschwert. Um die prohibitiven Kosten von On-Device-Messungen zu vermeiden, verlassen sich bestehende Auto-Scheduler (z. B. Ansor, AutoTVM) auf gelernte Kostenmodelle, um Kandidatenprogramme zu ranken. Diese Modelle betrachten Kandidaten jedoch typischerweise als statische Code-Snapshots (z. B. finalen TensorIR-Text).

Diese statische Perspektive führt zu zwei kritischen Einschränkungen:

Verlust des Transformationskontexts: Sie ignoriert den schrittweisen Transformationsprozess, der das Programm erzeugt hat. Infolgedessen können Evaluatoren nicht verstehen, wie eine Scheduling-Entscheidung in einem späteren Schritt von dem durch frühere Entscheidungen etablierten Kontext abhängt.
Anfälligkeit für syntaktisches Rauschen: Statische Encoder werden leicht durch oberflächliche Code-Variationen in die Irre geführt. Beispielsweise können zwei strukturell äquivalente Programme (z. B. Matrixmultiplikation mit Schleifengrößen $N=1024$ vs. $N=1023$ ) aufgrund des Boundary-Handlings (Loop Peeling) signifikant unterschiedliche TensorIR-Texte erzeugen, was dazu führen kann, dass das Modell deren Performance-Potenzial falsch klassifiziert, obwohl das Hardware-Verhalten identisch ist.

Methodik

Die Autoren schlagen ein Compiler-Weltmodell-Framework vor, das die Bewertung von Schedulings als aktionsbedingte latente Dynamik über Programmazustände neu definiert. Anstatt ein finales statisches Code-Artefakt zu bewerten, simuliert das Framework die Evolution des Programmazustands durch einen kontinuierlichen latenten Raum.

Das Framework besteht aus drei gelernten Komponenten, die innerhalb des TVM AutoSchedulers instanziiert sind:

1. Repräsentation des Tensorprogrammzustands (Encoder)

Input: TensorIR-Text (sowohl Vor- als auch Nach-Scheduling-Zustände).
Architektur: Eine trainierbare CodeBERT-basierte Architektur, die für lange Eingaben angepasst wurde. Sie verwendet ein Sliding-Window, um überlappende Chunks zu tokenisieren und zu kodieren, welche dann via Masked Mean Pooling zu einer dichten Vektorrepräsentation $z = E(s)$ aggregiert werden.
Training: Verwendet kontrastives Lernen, um ähnliche TensorIR-Zustände (aus demselben Workload-Source) im latenten Raum nah beieinander zu platzieren, während unähnliche Zustände voneinander weggestoßen werden. Dies schafft einen Zustandsraum, in dem semantische Äquivalenz trotz syntaktischer Variationen erhalten bleibt.

2. Aktionsbedingte mehrstufige Zustandsübergänge

Ziel: Die durch eine Sequenz von Scheduling-Aktionen induzierte Terminal-Zustandsrepräsentation vorherzusagen, ohne intermediäre ASTs explizit zu materialisieren.
Mechanismus: Ausgehend von einer initialen Zustands-Einbettung $z_0$ $z_{0}$ rollt das Modell die Wirkung einer Scheduling-Aktionssequenz $a_{1:T}$ $a_{1 : T}$ Schritt für Schritt aus.
- Aktionen werden serialisiert und in dichte Einbettungen kodiert.
- Ein Transitionsmodell $F(\cdot)$ , implementiert mittels einer TransH-basierten Architektur (ursprünglich aus dem Erlernen von Knowledge-Graph-Repräsentationen), sagt den nächsten latenten Zustand voraus: $z_t = F(z_{t-1}, a_t)$ .
- Dies modelliert Scheduling-Aktionen als geometrische Transformationen in einer relationsspezifischen Hyperebene, wodurch die Herkunft von Entscheidungen nachverfolgt und referenzielle Ambiguitäten (z. B. was eine Schleifenvariable nach einer Fusion repräsentiert) aufgelöst werden.
Training: Verwendet einen Multi-Step-Rollout-Loss, der einen State-Matching-Term (Minimierung des Abstands zu Ground-Truth-Zukunftszuständen) und einen Margin-basierten Ranking-Term (stellt sicher, dass der vorhergesagte Zustand näher am Ground-Truth als an negativen Samples liegt) kombiniert.

3. Ranking-basierte Kandidatenbewertung

Input: Die vorhergesagte Terminal-Zustandsrepräsentation $\hat{z}_T$ , strukturierte Aktionsmerkmale (z. B. Aktionsanzahl, Sequenzlänge) und optionale Hardware-Merkmale.
Output: Ein Skalar-Score für die relative Priorisierung von Kandidaten innerhalb desselben Workloads.
Architektur: Ein XGBoost-Ranking-Modell, das mit einem LambdaRank-Ziel trainiert wurde.
Trainingsstrategie: Das Modell sagt keine absolute Laufzeit voraus. Stattdessen lernt es, Kandidaten innerhalb einer Workload-Gruppe basierend auf gemessenen Ausführungszeiten zu ranken, wodurch das Problem als Learning-to-Rank-Aufgabe behandelt wird.

Wichtigste Beiträge

Weltmodell-Perspektive für Compiler: Die Arbeit passt das Weltmodell-Paradigma an die Tensor-Optimierung an und etabliert ein Framework, in dem Kandidaten-Schedules durch aktionsbedingte latente Zustandsentwicklung statt durch statische Snapshots modelliert werden.
Action-State-Trajectory-Datensatz: Konstruktion eines TVM/TenSet-basierten Datensatzes, der Tuning-Logs in ausgerichteten Action-State-Trajektorien (Vor-Scheduling-Zustand, Aktionssequenz, Zwischenzustände, Nach-Scheduling-Zustand) organisiert, für die mehrstufige Vorhersage.
Effiziente Suchintegration: Implementierung des Frameworks in den TVM AutoScheduler, wobei gezeigt wird, dass die Modellierung latenter Dynamiken die Autotuning-Kosten signifikant senkt bei gleichzeitiger Beibehaltung der hohen Suchqualität.

Experimentelle Ergebnisse

Die Methode wurde auf einer Intel Xeon Gold 6430 CPU und einer NVIDIA RTX 4090 GPU über sieben neuronale Netzwerkmodelle (einschließlich ResNet, BERT, GPT und OPT) und 103 extrahierte Subgraphen evaluiert.

Performance vs. Ansor: Unter einem Budget von 64 Versuchen verbessert die Methode die Latenz repräsentativer Subgraphen gegenüber Ansor um das 1,37-Fache auf der GPU und um das 1,54-Fache auf der CPU.
Beispiel-Effizienz: Die Methode erreicht eine vergleichbare Suchqualität wie Ansor-10K (10.000 Versuche) mit nur 1.000 Versuchen (10-mal weniger Messungen). Insbesondere übertrifft sie Ansor-10K im arithmetischen Mittel leicht und liegt im geometrischen Mittel innerhalb von 2,2 % bei repräsentativen Subgraphen.
End-to-End Speedup: Im Vergleich zu PyTorch/PyTorch-opt (cuDNN) liefert die Methode geometrische Mittel-Beschleunigungen von 4,61× (CPU) und 3,67× (GPU), wobei Spitzenbeschleunigungen von bis zu 58,61× erreicht werden.
Ablationsstudien: Das Entfernen des gelernten Transitionsmoduls (Rückkehr zum statischen Initial- oder Final-State-Encoding) führte zu signifikanten Latenzzunahmen (Median ~56–69 %), was bestätigt, dass die Modellierung der intermediären Zustandsentwicklung die kritischste Komponente ist.

Bedeutung und Limitationen

Bedeutung:
Die Arbeit argumentiert, dass die Behandlung der Compiler-Optimierung als dynamischer Prozess von Zustandsübergängen – anstatt als statisches Code-Evaluierungsproblem – es ermöglicht, die semantischen Abhängigkeiten zwischen Scheduling-Entscheidungen zu erfassen. Dieser Ansatz unterdrückt effektiv syntaktisches Rauschen und löst referenzielle Ambiguitäten auf, was zu einer effizienteren Suche mit weniger Hardware-Messungen führt.

Limitationen (wie von den Autoren angegeben):

Abhängigkeit vom Suchraum: Die Methode verbessert die Bewertung von Kandidaten, nicht deren Generierung. Ihre Effektivität ist durch die Fähigkeit des zugrunde liegenden Schedulers begrenzt, starke Schedulings vorzuschlagen; ein besserer Evaluator kann keine Schedulings wiederherstellen, die nie vorgeschlagen wurden.
Fehlerakkumulation: Vorhersagefehler können sich über lange oder komplexe Scheduling-Trajektorien im latenten Raum akkumulieren, was die Zuverlässigkeit in schwierigen Fällen verringern kann.
Ranking vs. Vorhersage: Das Modell ist für das Ranking innerhalb eines Workloads trainiert, nicht für die absolute Latenzvorhersage. Dies macht es ideal für die Online-Priorisierung, aber weniger geeignet für Szenarien, die kalibrierte Performance-Schätzungen über verschiedene Workloads hinweg erfordern.