Dieser Vortrag präsentiert max_evol, ein agentenbasiertes System für autonomes Maschinelles Lernen, das das Portieren und Optimieren von Deep-Learning-Kerneln (speziell von Triton zu JAX Palas) als ingenieurwissenschaftliches Suchproblem anstelle einer simplen Einmal-LLM-Übersetzung behandelt. Das von Alai entwickelte System kombiniert max_kernel und alpha_evolve zu einer geschlossenen Feedbackschleife mit Hardware-in-the-Loop-Validierung und evolutionärer Suche. Da direkte Übersetzungen zwischen Triton (pointerbasiert) und JAX Palas (abstrahierte ref-Objekte) zu massiven Leistungsabfällen führen, automatisiert der Agent Syntaxprüfung, Kompilierung, Korrektheitsvalidierung via Pytest und evolutionäre Leistungsoptimierung. In einer Fallstudie zur block-sparsamen Aufmerksamkeit (block sparse attention) erzielte das System einen 1,78- bis 1,8-fachen Speedup gegenüber der Baseline, während es Ausnahmen wie Out-of-Memory-Fehler autonom behandelte.
Wichtigste Punkte & Zeitstempel
0:00 Autonome ML-Entwicklung: Vorstellung von is_agent und max_evol als agentenbasierte Systeme zur Automatisierung komplexer ML-Ingenieuraufgaben wie der Kernel-Portierung.
0:29 Das Portierungsproblem: Die Notwendigkeit, PyTorch-Forschungscode für effiziente Inferenz in optimierten JAX-Code zu portieren, um XLA-Funktionen und benutzerdefinierte Kernel voll auszuschöpfen.
0:32 Die Rolle von Custom Kernels: Kompiliereinheiten verarbeiten dichte Matrixmultiplikationen exzellent, doch neuartige sparse Operationen (z. B. Mixture of Experts, block sparse attention) erfordern benutzerdefinierte Kernel, um I/O-Arbeiten (HBM-Datentransfers) mit Rechenoperationen zu überlappen.
6:00 Dachlinienmodell (Roofline Model): Nutzung evolutionärer Fitnessfunktionen (basiert auf Latenz und Tile-Volume), um die tatsächlich erreichbare Leistung näher an die theoretische Hardware-Obergrenze heranzuführen.
7:42 Triton versus JAX Palas: Triton fungiert als dünner CUDA-Wrapper mit expliziter Pointer-Arithmetik, während Palas abstrakte ref-Objekte und Block-Spezifikationen verwendet, weshalb direkte Übersetzungen zu 5- bis 10-fachen Leistungseinbrüchen führen.
9:53 Engpass manueller Portierung: Mangelnde Spezialisten, die sowohl ML-Wissen als auch Low-Level-Hardware-Kompetenz besitzen, machen manuelle Portierungen extrem zeitaufwendig und fehleranfällig.
11:54 LLM-Einsatzstufen: Abgrenzung von simplen Copilots (Level 1) und naiven Prompt-Agenten (Level 2) hin zu geschlossenen Agenten-Suchschleifen (max_evol).
13:11 Architektur von max_evol: Ein geschlossener Sub-Agent mit integrierter Dokumentationsindexierung, Ausführungsumgebung und Hardware-in-the-Loop-Validierung.
15:37 12-Schritte-Implementierungszyklus: Strukturierter Workflow mit strikten Verträgen zwischen den Schritten, um Fehler automatisiert zu erkennen und auf vorherige Stufen zurückzuspringen.
17:04 Validierungs-Pipeline: Stufenweise Überprüfung von Syntax, Linting, Kompilierung, Formkonsistenz und numerischer Genauigkeit via Pytest auf CPU.
18:37 Cloud-Ausführung via CLI: Nutzung der offiziellen Google Colab CLI und geplante Erweiterung auf diverse Cloud-GPU-Anbieter.
19:24 Rust-basierte Werkzeuge: Implementierung des Hauptagenten in Rust mit Tools wie kernel_db zur Verfolgung des Optimierungsfortschritts.
20:30 Evolutionäre Hardware-Optimierung: Mutation von Kernel-Implementierungen, Ausführung auf tatsächlichen Beschleunigern, Erfassung von Metriken und Iteration über eine Kernel-Datenbank.
23:06 Fallstudie und Fehlerbehebung: Portierung von block-sparsamer Aufmerksamkeit mit Erreichung eines 1,8-fachen Speedups bei gleichzeitiger autonomer Bewältigung von Out-of-Memory- und Slice-Fehlern.
27:47 Roadmap: Geplante Bereitstellung umfangreicher Benchmarks und Ausweitung auf hochgradig parallelisierte Experimente.
Glossar
Triton: Eine Programmiersprache und ein Compilier-Framework, das als dünner Wrapper um CUDA dient und die Entwicklung von Hochleistungskerneln mit expliziter Pointer-Arithmetik ermöglicht.
JAX Palas: Ein Programmier-Framework in JAX, das das Schreiben von High-Performance-Kerneln über abstrakte Referenzobjekte (ref) und Block-Spezifikationen anstelle direkter Pointer-Operationen erlaubt.
XLA (Accelerated Linear Algebra): Ein Domain-Specific Compiler, der lineare Algebra-Graphen optimiert, um die Ausführungsgeschwindigkeit auf Beschleunigern wie GPUs und TPUs zu maximieren.
HBM (High Bandwidth Memory): Ein speicherintensiver, vertikal gestapelter Arbeitsspeicher für Grafikprozessoren, der extrem hohe Datenübertragungsraten bei geringer Latenz bietet.
DMA (Direct Memory Access): Eine Hardware-Funktion, die es Peripheriegeräten ermöglicht, Daten direkt in den Hauptspeicher zu schreiben oder daraus zu lesen, ohne die CPU zu belasten.
Roofline-Modell (Roofline Model): Ein visuelles Leistungsmodell zur Bewertung der erreichbaren Rechenleistung von Prozessoren unter Berücksichtigung von Rechenintensität und Speicherbandbreite.
Mixture of Experts (MoE): Eine Deep-Learning-Architektur, bei der je nach Eingabe nur bestimmte neuronale Teilnetzwerke (Experten) aktiviert werden, wodurch spärliche (sparse) Aktivierungsmuster entstehen.
Tile-Volume: Die Datenmenge (in Bytes), die bei einer rechentechnischen Operation in einem Datenblock gleichzeitig verarbeitet und im schnellen SRAM gehalten wird.
Pytest: Ein verbreitetes Test-Framework für Python zur automatisierten Validierung der numerischen Korrektheit von Softwarekomponenten.
Doom Loop: Ein unerwünschter algorithmischer Zustand, bei dem ein Sprachmodell oder Agent wiederholt dieselben fehlerhaften Werkzeuge oder Eingaben anwendet, ohne Fortschritte zu erzielen.