💻 computer science

HTAM: Hierarchical Transition-Attended Memory for Operator Optimization

Dieser Beitrag stellt HTAM vor, ein hierarchisches Framework, das Optimierungserfahrung in einen Grob-zu-Fein-Übergangsgraphen organisiert, um die Generierung von GPU-Operatoren durch LLMs zu steuern, wodurch Granularitätsmismatches behoben und die Korrektheit sowie Leistung von Kernen erheblich verbessert werden.

Ursprüngliche Autoren: Yining Zhang, Mingyang Yi, Chen Wang, Xuwen Xiang, Tianhe Jia, Zedong Dan, Chengqing Zong, Yue Wang

Veröffentlicht 2026-05-29

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Yining Zhang, Mingyang Yi, Chen Wang, Xuwen Xiang, Tianhe Jia, Zedong Dan, Chengqing Zong, Yue Wang

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, einem brillanten, aber unerfahrenen Lehrling beizubringen, den Motor des schnellsten Rennwagens der Welt zu bauen. Sie verfügen über eine Bibliothek mit Bauplänen, doch der Lehrling verirrt sich häufig in den Details oder schlägt Änderungen vor, die gut klingen, den Motor jedoch zerstören.

Dieser Beitrag stellt HTAM (Hierarchical Transition-Attended Memory) vor, ein neues System für einen „intelligenten Mentor", das entwickelt wurde, um Large Language Models (LLMs) dabei zu unterstützen, hochleistungsfähigen Computercode für Grafikprozessoren (GPUs) zu schreiben.

Hier ist die Funktionsweise von HTAM, erläutert durch einfache Analogien:

Das Problem: Die Falle „Zu breit vs. Zu eng"

Derzeit steht die KI, wenn sie versucht, Computercode zu korrigieren, vor einem Dilemma:

Der Ansatz „Vage Andeutung": Die KI erhält einen allgemeinen Hinweis wie „Machen Sie den Speicherzugriff schneller". Dies ist leicht zu merken, aber zu abstrakt. Die KI weiß nicht, wie sie den Code tatsächlich ändern muss, um ihn schneller zu machen.
Der Ansatz „Übermäßig detailliert": Die KI erhält eine riesige Liste jedes einzelnen winzigen Code-Änderung, die je vorgenommen wurde. Das sind zu viele Informationen. Es ist, als würde man versuchen, eine spezifische Schraube in einem Lagerhaus voller Millionen von Schrauben zu finden; die KI wird überwältigt und findet den richtigen Zug nicht.

Die Lösung: Das „Kochbuch des Meisterköchens"

HTAM löst dieses Problem, indem es das Gedächtnis der KI wie ein Kochbuch eines Meisterköchens organisiert, das in drei Schichten strukturiert ist:

Die Speisekarte (Globale Richtungen): Zuerst fragt das System: „Was ist das Hauptziel?" Ist das Problem, dass der Wagen ausfällt (Speicherzugriff)? Ist es, dass der Motor überhitzt (Datenwiederverwendung)? Ist es, dass die Räder zu schnell drehen (Parallelität)?
- Analogie: Dies ist wie die Entscheidung: „Heute reparieren wir die Bremsen", anstatt zu versuchen, das ganze Auto auf einmal zu reparieren.
Die spezifischen Rezepte (Lokale Strategien): Sobald das Ziel festgelegt ist (z. B. „Bremsen reparieren"), sucht das System nach spezifischen, bewährten Techniken für dieses Ziel.
- Analogie: Statt nur zu sagen „Bremsen reparieren", ruft es ein spezifisches Rezept auf: „Ersetzen Sie die Bremsbeläge durch keramische" oder „Stellen Sie den hydraulischen Druck ein". Dies sind konkrete, umsetzbare Schritte, die die KI tatsächlich in Code umsetzen kann.
Die „Nächster Zug"-Karte (Übergangserfahrung): Dies ist das Geheimnis. HTAM merkt sich nicht nur, was zu tun ist, sondern auch, was als Nächstes zu tun ist.
- Analogie: Ein Meisterkoch weiß, dass nach dem „Anbraten des Fleisches" der nächste logische Schritt das „Ablöschen der Pfanne" ist. Wenn Sie versuchen, das Pfanne abzudampfen, bevor Sie das Fleisch anbraten, funktioniert es nicht. HTAM lernt diese Sequenzen. Es weiß, dass wenn Sie gerade den „Speicherzugriff" behoben haben, der nächste beste Versuch möglicherweise „Datenwiederverwendung" ist und nicht „Grenzwertbehandlung".

Funktionsweise in der Praxis

Das System läuft in einer Schleife und fungiert wie ein Trainer, der den Lehrling anleitet:

Prüfen der Anzeigetafel: Die KI betrachtet den aktuellen Code und sieht, wo er langsam ist oder fehlerhaft.
Auswählen eines Ziels: Mit Hilfe ihrer „Speisekarte" (Globales Gedächtnis) wählt sie eine übergeordnete Richtung aus (z. B. „Optimieren wir, wie Daten bewegt werden").
Auswählen eines Zuges: Mit Hilfe ihrer „Rezepte" (Lokales Gedächtnis) wählt sie eine spezifische Code-Änderung aus (z. B. „Verwenden Sie einen schnelleren Weg zum Laden von Daten").
Betrachten der Geschichte: Bevor sie den Zug ausführt, prüft sie ihre „Nächster Zug-Karte" (Übergangsgedächtnis). Sie fragt: „Wir haben gerade X getan; sagt uns die Geschichte, dass Y als Nächstes eine gute Idee ist?"
Schreiben und Testen: Die KI schreibt den neuen Code, testet ihn, und wenn er funktioniert, aktualisiert sie ihr Kochbuch mit diesem neuen Erfolg. Wenn er fehlschlägt, aktualisiert sie das Buch mit dem, was man nicht tun sollte.

Die Ergebnisse: Ein schnellerer, intelligenterer Lehrling

Die Autoren testeten dieses System auf KernelBench, einem Standard-Testsuite für die Leistung von GPU-Code.

Genauigkeit: Das System erhielt den Code in 98,4 % der Fälle richtig (im Vergleich zu deutlich niedrigeren Raten bei Standard-KI).
Geschwindigkeit: Es fand die schnellste Lösung in 84 % der Fälle.
Leistung: Der von ihm geschriebene Code war im Durchschnitt fast 2-mal schneller als der Code, der von Standard-KI-Methoden geschrieben wurde.

Warum dies wichtig ist

Der Beitrag behauptet, dass HTAM durch die Organisation des Gedächtnisses auf diese Weise – indem es das „große Ganze" von den „winzigen Details" trennt und die Reihenfolge der Operationen merkt – eine chaotische Suche nach dem perfekten Code in eine strukturierte, effiziente Reise verwandelt. Es rät nicht einfach; es folgt einem erlernten Pfad von Expertenentscheidungen, wodurch es viel besser darin wird, den komplexen, hochgeschwindigkeitsfähigen Code zu schreiben, der für moderne KI- und Grafik-Anwendungen benötigt wird.

Technischer Zusammenfassung: HTAM – Hierarchical Transition-Attended Memory zur Optimierung von Operatoren

1. Problemstellung

Das rasche Skalieren von Large Language Models (LLMs) und die Evolution spezialisierter Beschleunigerarchitekturen (z. B. GPUs) haben eine kritische Nachfrage nach der Entwicklung hochperformanter Operatoren geschaffen. Obwohl LLMs starke Fähigkeiten zur Codegenerierung demonstriert haben, bleibt die Optimierung von GPU-Kernen ein hardwarebewusstes Suchproblem, das sich nicht mit einer einmaligen Generierung lösen lässt.

Die Kernherausforderung liegt in der Granularitätsdiskrepanz bestehender LLM-basierter Optimierungsmethoden:

Grobgitterige Ansätze (z. B. textuelle Hinweise oder Engpassdiagnosen) sind wiederverwendbar und leicht abrufbar, aber oft zu abstrakt, um die Generierung ausführbarer CUDA-Codes zu steuern.
Feingitterige Ansätze (z. B. detaillierte Implementierungsspuren) liefern handlungsrelevante Details, erzeugen jedoch einen großen, verrauschten Suchraum, was den Abruf ineffizient macht und Optimierungsengpässe verschleiert.

Darüber hinaus erfordert die Kernel-Optimierung gekoppelte Mehrstufenentscheidungen (Auswahl von Kachelgrößen, Thread-Zuordnungen, Speichermustern usw.) sowie die nicht-triviale Übersetzung dieser Strategien in korrekten, hardwarebewussten Code. Bestehende Methoden haben Schwierigkeiten, Optimierungserfahrung so zu organisieren, dass sie diesen mehrstufigen Suchprozess leiten und gleichzeitig konkrete Beweise für Codeänderungen liefern.

2. Methodik: HTAM-Rahmenwerk

Um die Granularitätsdiskrepanz zu adressieren, schlägt die Arbeit HTAM (Hierarchical Transition-Attended Memory) vor, ein Grob-zu-Fein-Rahmenwerk, das Optimierungserfahrung in einer strukturierten Speicherbank organisiert. HTAM bildet den Entscheidungsprozess menschlicher Experten nach: Zuerst wird eine hochrangige globale Richtung ausgewählt, die dann durch konkrete lokale Strategien verfeinert wird.

2.1 Hierarchischer Übergangsgraph (HTG)

HTAM strukturiert seinen Speicher als Hierarchischen Übergangsgraphen $M = (V_G, V_{GL}, E_G)$ :

Globale Knoten ( $V_G$ ): Stellen grobe Optimierungsrichtungen dar (z. B. „Speicherzugriff", „Datenwiederverwendung", „Durchsatz"). Diese Knoten speichern hochrangige Absichten, statische Priors (Ziele, Auslöser, Risiken) und Laufzeitstatistiken (Erfolgsraten, Verbesserungstrends).
Lokale Knoten ( $V_{GL}$ ): Stellen spezifische, ausführbare Strategien unter einem übergeordneten globalen Knoten dar (z. B. „ausgerichtete vektorisierte Ladevorgänge" unter „Speicherzugriff"). Diese Knoten speichern Implementierungswissen, einschließlich Änderungsmustern, Verifikationsprüfungen sowie positiver und negativer Evidenz aus früheren Ausführungen.
Übergangskanten ( $E_G$ ): Verbinden globale Knoten und speichern pfadabhängiges Wissen über den Übergang von einer globalen Richtung zu einer anderen. Diese Kanten protokollieren, wann ein Übergang anwendbar ist, seine Begründung, Risiken sowie Evidenz auf Trajektorienebene (Leistungssteigerungen, Kompilierungsergebnisse).

2.2 Optimierungsprozess

Die Optimierung verläuft als $T$ -stufiger Evolutionsprozess, bei dem die Implementierung $c_t$ iterativ aktualisiert wird:

Zustandszusammenfassung: Die aktuelle Implementierung und der Verlauf werden zu einem kompakten Zustand $s_t$ zusammengefasst, der Aufgabenkontext, Feedback (Kompilierung, Korrektheit, Laufzeit) und kürzliche Optimierungspfade enthält.
Auswahl der globalen Richtung: HTAM wählt eine globale Richtung $g_t$ $g_{t}$ basierend auf dem aktuellen Zustand und dem kürzlichen Verlauf aus. Es verwendet einen übergangsbewussten Bewertungsmechanismus, der Evidenz aus dem kürzlichen Präfix globaler Richtungen ( $p^g_t$ $p_{t}^{g}$ ) unter Verwendung positionsbewusster Gewichte (inspiriert von Aufmerksamkeitsmechanismen) aggregiert. Dies vermeidet die hohen Kosten der Speicherung vollständiger Ketten-Speicher, erfasst jedoch Mehrstufen-Abhängigkeiten.
- Der Score wird als gewichtete Summe der Übergangsnutzen von vorherigen Knoten zum Kandidatenknoten berechnet.
- Eine Softmax-basierte Richtlinie balanciert die Ausbeutung hochbelohnter Übergänge mit der Exploration.
Auswahl der lokalen Strategie: Sobald eine globale Richtung gewählt ist, wird der Suchraum auf die zugehörigen lokalen Knoten eingeschränkt. Ein LLM bewertet Kandidaten lokaler Strategien basierend auf dem Zustand, der aktuellen Implementierung und den spezifischen globalen/lokalen Knotenspeichern, um die beste Strategie $l_t$ auszuwählen.
Implementierung und Aktualisierung: Das LLM generiert die nächste Implementierung $c_{t+1}$ , geleitet von der ausgewählten lokalen Strategie. Evaluierungsfeedback wird dann verwendet, um den HTG zu aktualisieren: Statistiken globaler Knoten, Evidenz lokaler Knoten und Feedback der Übergangskanten werden aktualisiert, um zukünftige Schritte zu informieren.

3. Hauptbeiträge

Grob-zu-Fein-Rahmenwerk: HTAM zerlegt die LLM-basierte Operator-Optimierung in die Auswahl globaler Richtungen und die Verfeinerung lokaler Strategien und adressiert damit die Granularitätsdiskrepanz bestehender Methoden.
Strukturierte Speicherorganisation: Es führt einen Hierarchischen Übergangsgraphen ein, der Erfahrung in globale Knoten, lokale Knoten und Übergangskanten organisiert. Diese Struktur reduziert den Suchraum und die Speicherkomplexität im Vergleich zur Speicherung vollständiger Ketten oder flachem Abruf.
Übergangsbewusste Entscheidungsfindung: Das Rahmenwerk verwendet einen Bewertungsmechanismus, der den kürzlichen Optimierungsverlauf über Kanten-Speicher aggregiert und Mehrstufenentscheidungen leitet, ohne die exponentiellen Kosten der Speicherung vollständiger Trajektorien.
Empirische Validierung: Umfangreiche Experimente zeigen, dass HTAM Baseline-Methoden in Bezug auf Korrektheit, Geschwindigkeit beim Finden schneller Lösungen und dem geometrischen Mittel der Beschleunigung signifikant übertrifft.

4. Experimentelle Ergebnisse

Die Experimente wurden auf KernelBench, einer Suite von 250 Aufgaben (Level-1 einzelne Operatoren, Level-2 fusionierte Operatoren, Level-3 modellbezogene Aufgaben), durchgeführt, wobei DeepSeek-R1 als Backbone-Modell auf NVIDIA A100-GPUs verwendet wurde.

Leistungssteigerungen: HTAM erreichte 98,4 % Korrektheit, 84,0 % Fast@1 (Prozentsatz der Aufgaben mit einer korrekten und schnelleren Lösung beim ersten Versuch) und eine geometrische mittlere Beschleunigung von 1,978×.
Vergleich mit Baselines:
- Im Vergleich zur stärksten reproduzierten Baseline (Flat Memory Retrieval) verbesserte HTAM die Korrektheit um 8,8 Prozentpunkte, Fast@1 um 25,2 Prozentpunkte und die Beschleunigung von 1,464× auf 1,978×.
- Es übertraf die reine LLM-Generierung und Feedback-basierte Verfeinerungsmethoden signifikant.
Generalisierung:
- Backend-Transfer: HTAM blieb über verschiedene LLM-Backends (DeepSeek-R1, GPT-4o, DeepSeek-V4-Flash) hinweg effektiv, obwohl die Leistung mit der Codegenerierungsfähigkeit des Backends korrelierte.
- Cross-Benchmark-Transfer: Bei Initialisierung mit einem aus KernelBench aufgebauten Speicher und Anwendung auf Robust-KBench-Aufgaben zeigte HTAM verbesserte Beschleunigungen (z. B. +1,76× für LayerNorm), was auf übertragbare Optimierungserfahrung hindeutet, wobei die Gewinne operatorabhängig waren.
Ablationsstudien: Das Entfernen der hierarchischen Struktur oder des lokalen Speichers führte zu signifikanten Leistungseinbußen (z. B. fiel die Beschleunigung ohne lokalen Speicher auf 0,974×), was bestätigt, dass strukturierte global-lokale Erfahrung entscheidend ist. Ebenso reduzierte das Entfernen des präfixbewussten Übergangsscorings die Leistung und validierte die Notwendigkeit einer verlaufabhängigen Entscheidungsfindung.

5. Bedeutung und Behauptungen

Die Arbeit behauptet, dass strukturierter Speicher nicht lediglich als Repository vergangener Fälle dient, sondern als Entscheidungs-Prior für die Auswahl der nächsten nützlichen Optimierungsrichtung. Durch die hierarchische Organisation von Erfahrung und die Modellierung von Übergängen zwischen Optimierungsschritten ermöglicht HTAM LLMs, den komplexen, mehrstufigen Suchraum der GPU-Kernel-Optimierung effektiver zu navigieren als flacher Abruf oder einmalige Generierung.

Die Autoren betonen, dass HTAM einen praktischen Weg zur Automatisierung der Entwicklung hochperformanter Operatoren bietet und dabei Korrektheits- und Beschleunigungsniveaus nahe an Experten erreicht, während die Online-API-Kosten überschaubar bleiben (ca. 0,21 $ pro Aufgabe für eine 6-stufige Optimierung mit DeepSeek-R1).

Anerkannte Einschränkungen:

Die Evaluation ist durch Rechen- und API-Kosten begrenzt und konzentriert sich hauptsächlich auf KernelBench und eine Teilmenge von Robust-KBench.
Das System verlässt sich auf ausführbare Validierung (Kompilierung, Korrektheit, Laufzeit) anstelle formaler Verifikation, was bedeutet, dass Korrektheit nicht für alle möglichen Formen oder Datenverteilungen garantiert ist.
Die Generalisierung auf erheblich unterschiedliche Hardware oder Operator-Familien könnte eine Anpassung des Speicherschemas erfordern.