Nach der Präsentation des Jalapeño MLA-Kernels auf der HotChips und den darauffolgenden Kommentaren von SemiAnalysis gab es viele Diskussionen. Als Hardware-Team von OpenAI haben wir dieses kleine Goldnugget gerade eben angerissen: die Tatsache, dass KI unsere Kernel schreibt, und dass wir, wenn sie das tut, nicht wirklich jede Zeile des Kernels verstehen müssen. Wir haben aber deutlich ausgelassen: Wie ist so etwas möglich? Wie sollte man dies im Vergleich zu einer traditionelleren Methode der Codegenerierung betrachten? Ist der optimierte Kernel genauso fehlerfrei wie der nicht optimierte?
Was meinen Hintergrund betrifft, so habe ich über ein Jahrzehnt an Compilern für Beschleuniger gearbeitet. Ich habe XLA ins Leben gerufen, eine hervorragende Compiler-Infrastruktur mit einem erstklassigen, unternehmensübergreifenden Team und Engagement. In den letzten mehr als zwei Jahren bei OpenAI habe ich versucht, das Konzept zu überdenken, wie Compiler im Zeitalter der KI funktionieren sollten. Neue Compiler-Formulierungen werden auf bestehenden Stärken aufbauen, aber es ist nicht zu leugnen, dass es ein mächtiges neues Werkzeug gibt, das es im Werkzeugkasten zu nutzen gilt.
Dies wird eine kleine Reise werden, aber ich hoffe, zu erhellen, wie KI für die Automatisierung der Verbesserung von Computerprogrammen eingesetzt wird; d.h. die optimierende Kompilierung. Ich glaube tatsächlich, dass wir durch KI etwas erleben könnten, das wir als "Compiler 2.0" bezeichnen. KI ist in ihren Vorschlägen grundsätzlich weniger eingeschränkt, und was sie vorschlägt, ist ein Ergebnis des Trainings und des Kontexts des Modells. Dies führt mich dazu, sie als "stochastischen Optimierer" zu klassifizieren – dies kann Herausforderungen mit sich bringen, ist aber, wie wir sehen werden, auch eine Quelle großer Stärken…
Eine Vielzahl akademischer Forschung und industrieller Anwendungen ist bereits auf diesem Weg und deckt rasch das Potenzial für die Beteiligung von KI im Bereich der optimierenden Compiler auf, aber wir sind an einem Punkt angelangt, an dem eine grobe Erklärung angebracht ist.
Hintergrund
Compiler nehmen Programme entgegen und geben übersetzte oder verbesserte Versionen dieser Programme aus.
Programme haben sowohl auf der Eingabe- als auch auf der Ausgabeseite eine Semantik, die uns sagt, was die Programme bedeuten, was sie möglicherweise tun können und wie wir über diese Dinge, die sie tun können, nachdenken können.
Diejenigen von uns, die an Compilern arbeiten, betrachten sie oft wie reine Funktionen – sie nehmen eine Datenstruktur entgegen und geben eine Datenstruktur aus, die eine entsprechende Semantik haben sollte.
Manchmal konzentrieren sich unsere Compiler auf das "Herunterbrechen" oder "Übersetzen". Zum Beispiel können sie C entgegennehmen und x86-64-Assembly ausgeben, was wir oft als "niedrigere Ebene" betrachten. Aber oft tun sie im Rahmen dieses Prozesses mehr als nur Übersetzung…
Unsere Compiler konzentrieren sich in der Praxis auf das "Optimieren". Sie können eine Datenstruktur, die das Programm repräsentiert – in unserer Fachsprache eine "Zwischendarstellung" (Intermediate Representation, IR) – entgegennehmen und versuchen, eine bessere Version dieses Programms zu erzeugen. Manchmal bedeutet "besser", dass es weniger Taktzyklen zur Ausführung benötigt, manchmal, dass es weniger unnötigen Code enthält, manchmal bedeutet es die Spezialisierung auf Dinge, von denen wir beweisen können, dass sie "wahr sein müssen" für das Programm (partielle Auswertung).
Betrachten wir nun kurz, dass LLMs ursprünglich entwickelt wurden, um menschlichen Text von einer Sprache in eine andere zu übersetzen. Übersetzung liegt eindeutig in ihrem Fachgebiet. Und wir können durch unsere tägliche Nutzung von LLMs sehen, dass sie auch neue Lösungen schreiben und bestehende Lösungen verbessern können. Viele von uns Programmierern haben auch Erfahrung damit, ein LLM zu bitten: "Optimiere diesen Code-Schnipsel", und sie können es bemerkenswert gut. (Allerdings müssen wir wissen, dass sie den Code korrekt optimiert haben, worauf wir noch zu sprechen kommen!) Dies soll lediglich hervorheben, dass LLMs die Fähigkeiten besitzen, die wir in einem optimierenden Compiler suchen.
Optimierung und Optimalität
Optimierende Compiler versuchen, wenig überraschend, die Optimalität des Programms, an dem sie arbeiten, nach einem bestimmten Ziel (normalerweise Ausführungszeit) zu steigern. Dies ist im allgemeinen Fall für ein beliebiges Programm so schwierig, dass es ein Theorem gibt, das als Vollbeschäftigungstheorem für Compiler-Ingenieure bekannt ist. (Ich habe das erst herausgefunden, nachdem ich mich für den Beruf des Compiler-Ingenieurs entschieden hatte, aber es hat mich dennoch getröstet!)
"Superoptimierer" sind ein erstaunliches kleines Teilgebiet der optimierenden Compiler. Stellen Sie sich ein gegebenes Programm vor, und wir können durch seine Semantik sagen, was es tut. Was ist das optimalste Programm, das dieselbe Semantik hat? Das ist es, was Superoptimierer zu bewältigen versuchen, und es ist im Wesentlichen ein Suchproblem…
Stellen Sie sich vor, ich versuche, das kürzeste Programm zu finden, das dieselbe Semantik hat, und ich hätte eine Möglichkeit zu fragen, ob ein Kandidatenprogramm dieselbe Semantik hat. Ich könnte hypothetisch jedes Programm in objektiver Reihenfolge aufzählen und das kleinste auswählen, das dieselbe Semantik hat.
Allerdings klingt das Aufzählen jedes Programms in objektiver Reihenfolge ziemlich unlösbar. Eines meiner Lieblingsfachartikel, veröffentlicht 2013 unter dem Titel "STOKE" (Stochastic Superoptimization), fragte: "Nun, was ist, wenn wir Programme einfach immer wieder zufällig verändern, beobachten wir dann irgendwann das beste Programm?" Sie schlugen vor, dass man durch einen Random Walk (und mit unserem Ur-Maschinenlern-Freund Markov Chain Monte Carlo / Metropolis-Hastings) irgendwann dieses optimale Programm sehen würde.
Monte-Carlo-Veränderungen sind typischerweise dumm (man wählt zufällig eine Veränderung), aber auch schnell. LLMs sind sehr intelligent (viele Reasoning-Token), aber vergleichsweise langsam.
Was wäre, wenn wir anstelle der dummen/schnellen Monte-Carlo-Veränderungen LLMs die Richtungen herausfinden ließen, in die die Programme entwickelt werden sollen? Wir hätten einen stochastischen Optimierer, der sehr intelligent wäre und unser Programm durch den optimierten Programmsraum führen würde.
Intuitionen für Optimierung
Treten wir einen Schritt zurück. Denken Sie an die Person, die Sie kennen und die am besten "optimiert den Haufen aus Code-Schnipseln" verkörpert. Nennen wir sie kurz "Optimier-Olli". Olli hat wahrscheinlich eine Bauchintuition dafür, welche Arten von Code-Änderungen Früchte tragen könnten. Olli probiert wahrscheinlich einige Dinge aus, um zu sehen, ob sie funktionieren, und wenn sie nicht funktionieren, macht er sie rückgängig und probiert etwas anderes. Aber er hat eine gewisse Intuition dafür, welche Arten von Dingen möglich sind und wie man vielleicht den Compiler schlagen kann.
Diese Intuitionen, die Olli hat, gehen oft über das hinaus, was Compiler tun. Obwohl moderne optimierende Compiler in ihren Ergebnissen ziemlich beeindruckend sind, basieren sie auf relativ einfachen Regeln und Heuristiken. In der Fachsprache basieren sie auf der Idee einer lokalen Datenflusstransformation, die bis zu einem Fixpunkt ausgeführt wird. Wir ordnen die Überlegungen auch phasenweise an; d.h. wir bauen Compiler-Pipelines, um zuerst A und dann B zu berücksichtigen, aber nicht das zusammengesetzte AB-Problem. Scheduler und Registerallokatoren sind ein berüchtigtes Beispiel dafür; viele Doktorarbeiten wurden zum zusammengesetzten Scheduler-Register-Allokator versucht (um die Vorteile der Aufhebung der Phasenordnung zu nutzen), aber es hat sich als schwierig erwiesen, sie in der Praxis zum Funktionieren zu bringen.
Deshalb ist Ollis Fachwissen wertvoll. Oft weiß Olli, wie man mehrere NP-vollständige Probleme mit Heuristiken ausbalanciert, die auf die Situation zugeschnitten sind. Es gibt also mehr maßgeschneiderte Kontextwahrnehmung und -sensitivität. Olli ist auch in der Lage, Techniken einzusetzen, die optimierende Compiler möglicherweise nicht gewinnbringend anwenden können, insbesondere in Kombination, wie z.B. Auslagerungen oder die Erstellung benutzerdefinierter ABIs oder Transformationen zur Ermöglichung von Vektorisierung oder die anderen Dinge, die uns murmeln lassen: "Ich wünschte, der Compiler hätte eine Möglichkeit, das einfach zu tun…"
Betrachten Sie nun, dass KI durch ihre Reasoning-Fähigkeiten in der Lage sein könnte, als eine Art Mini-Olli zu agieren. Sie hat vielleicht nicht die passende Intuition, was sich auszahlen wird, aber sie hat eine Ahnung davon, was profitabel sein könnte, und sie kann viele, viele Versuche unternehmen.
Mit diesem Ansatz können wir, anders als im STOKE-Papier, nicht garantieren, dass wir mit fortschreitender Zeit das optimale Programm sehen, aber weil die KI "menschlichere" Reasoning-Fähigkeiten hat, kann sie tatsächlich pro Zeiteinheit eine signifikante, menschenähnliche Wirkung erzielen.
Die Verbindung herstellen: MLA-Kernel
Lassen Sie mich zunächst sagen: Ich weiß nicht, welchen Low-Level-Code die KI für unseren Jalapeño MLA-Kernel ausgegeben hat, aber ich weiß, wie man das numpy für MLA eintippt.
In dem XLA-Compiler, an dem ich zuvor gearbeitet habe, würden wir diese numpy-Operationen zu Klumpen zusammenfassen und dann ein Metaprogramm namens "Emitter" verwenden, um es auf Schleifen, Anweisungen und Low-Level-Primitive herunterzubrechen.
Als das XLA-Compiler-/Emitter-Programm dies tat, musste ich mich nicht darum kümmern, welches Assembly hinten herauskam. Für unseren stochastischen Optimierer übernimmt die KI konzeptionell die Stelle des Emitter-Metaprogramms – sie bricht sowohl herunter als auch optimiert, und wir können sie bitten, weiter und weiter in Richtung Roofline zu optimieren.
Ich hoffe, dies macht deutlich, wo die KI einzuordnen ist und wie sie analog zu einer Komponente in einem bestehenden optimierenden Compiler-System ist. Es ist auch hilfreich, darüber nachzudenken: Welche Schicht wir als "Assembler-Code" betrachten, verschiebt sich nun nach oben. Wenn Sie normalen C++-Code eingeben und ihn mit -O3 (der höchsten typischen Optimierungsstufe) kompilieren, erwarten Sie nicht, das resultierende Assembly zu verstehen, selbst wenn Sie den eingegebenen C++-Code verstanden haben. Wir machen hier das Analoge, aber mit einer höherwertigen und mathematischeren Eingabespezifikation.
Eine entscheidende Frage ist nun, wie wir überprüfen, dass das Programm, das wir von der KI erhalten, tatsächlich äquivalent zur Beschreibung auf höherer Ebene / numpy ist. Dieser Überprüfungsmechanismus etabliert die Korrektheit des stochastischen KI-Optimierungsprozesses. Ich erwarte, dass ein zukünftiger Blogbeitrag näher darauf eingehen wird, aber für den Moment sei gesagt, dass das Testen auf semantische Äquivalenz möglich ist und wir es tun. Beschleunigerprogramme eignen sich besonders gut für starke, vollständige Verträge, die wir verifizieren können, "ob sie genau das tun, was das KI-optimierte Programm tut", da sie in ihrem breiten Kontext recht mathematisch und datenflussorientiert sind.
Beachten Sie, dass viele relevante Techniken in diesem Bereich durch Bemühungen im Teilgebiet der Programmsynthese Pionierarbeit geleistet wurden. Während optimierende Compiler sagen: "Hier ist ein Programm mit Semantik, mach es besser, aber mit äquivalenter Semantik!", sagt die Programmsynthese: "Ich glaube, es gibt ein Programm mit dieser Semantik, bitte versuche, das Beste zu finden, das du kannst." Programmsynthese ist ein schwierigeres Problem als optimierende Kompilierung, aber sie ist auch grundsätzlich weniger eingeschränkt. Es ist im Wesentlichen das, was Menschen wie Olli tun, wenn sie besser sind als der optimierende Compiler, und es ist etwas, bei dessen Automatisierung uns KI jetzt helfen kann. Die KI kann "Inspiration" aus dem ursprünglichen Programm ziehen, aber sie muss nicht nur kleine lokale Transformationen daran vornehmen. Klassische optimierende Compiler werden nicht sehen: "Oh, du hast einen Bubble-Sort geschrieben" und ihn, durch das Verständnis des Vertrags, gegen einen Quick-Sort austauschen, aber sowohl Olli als auch die KI sind dazu in der Lage. Dies bringt uns mehr in das Regime der Programmsynthese mit stochastischer Optimierung als in das klassische Regime der optimierenden Compiler.
Dies alles fügt sich in der Tatsache zusammen, dass man mit etwas beginnen kann, das "nicht weit vom numpy entfernt ist", 48 Stunden wartet und einen optimierten Kernel mit derselben Semantik hat, wie wir in unserem HotChips-Vortrag gezeigt haben:
Wie die Folie ebenfalls anmerkt, können wir auf unserer Maschine oft beobachten, dass die KI in ihrer Leistung unsere menschlichen Experten überholt, sogar bei den Kerneln, von denen wir dachten, sie seien recht gut abgestimmt. Oft ist noch ein beträchtlicher erreichbarer Prozentsatz übrig, einfach aufgrund der vielen Varianten von Kombinationen / Permutationen, die möglicherweise untersucht werden müssen. Diese sind für einen menschlichen Performance-Ingenieur oft unlösbar mühsam.
Zusammenfassung & Schlussfolgerung
Ein Compiler ist letztendlich nur eine Funktion. Wir geben unser Programm an diese Funktion und erhalten eine bessere Version unseres Programms zurück. Von dem Programm, das wir erhalten, und dem Programm, das wir eingegeben haben, erwarten wir, dass sie dieselbe Semantik haben.
Traditionelle optimierende Compiler verbessern Programme durch Datenflussregeln und Heuristiken. Diese sind in ihrer Herkunft vollständig verständlich, können aber auch in den möglichen Schritten eingeschränkter sein.
Im Gegensatz dazu muss die KI als stochastischer Optimierer nur "hart nachdenken" und etwas ausspucken. Ihre Schritte sind nicht grundsätzlich eingeschränkt, was sie analoger zu unserem menschlichen Experten-Optimierer macht. Wir brauchen Wege, um zu überprüfen, ob die Programme, die sie ausspuckt, korrekt sind und dieselbe Semantik implementieren, die wir eingegeben haben, und diese haben wir auch eingerichtet. Und diese Art der KI-Optimierung eignet sich besonders gut für mathematische Operationen, die sehr starke Verträge haben. Die Verträge vermeiden die Notwendigkeit, den Kernel Zeile für Zeile zu verstehen.
So haben wir den KI-generierten MLA-Kernel erhalten!
