Axon: A Synthesizing Superoptimizer for Tensor Programs
Axon ist ein synthetisierender Superoptimierer, der durch die Verwendung von Programmsynthese und SMT-basierter semantischer Verifizierung automatisch hochperformante Tensorprogramme für KI-Beschleuniger generiert, um algebraische Transformationen zu entdecken, Tiling und Instruktionsselektion zu optimieren sowie Operatoren zu fusionieren, ohne sich auf handgefertigte Rewrite-Regeln zu verlassen.
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 möchten den perfekten Kuchen backen, aber Sie arbeiten in einer High-Tech-Küche mit sehr speziellen, eigenwilligen Öfen und Mixern. Sie wissen genau, was Sie backen wollen (das Rezept), aber um die Maschine so schnell wie möglich dazu zu bringen, es zu tun, müssen Sie genau wissen, wie Sie die Zutaten einfüllen, wann Sie rühren und welche spezifischen Knöpfe am Ofen Sie drücken müssen.
Lange Zeit war das Schreiben von Code für KI-Chips (die „Öfen“) so, als würde man diesen Kuchen von Hand backen, Krümel für Krümel. Man musste ein Meisterkoch sein, der die exakte Mechanik des Ofens versteht, um ein gutes Ergebnis zu erzielen. Wenn man einen kleinen Fehler beim Einfüllen der Zutaten macht, kommt der Kuchen langsam und trocken aus dem Ofen.
Axon ist ein neues Werkzeug, das wie ein automatisierter, super-intelligenter Sous-Chef fungiert, der nicht nur ein Rezept befolgt – er erfindet bessere Wege, es zu backen.
So funktioniert Axon, in einfachen Schritten erklärt:
1. Das „Was“ vs. das „Wie“
Normalerweise schreibt ein Programmierer ein Rezept wie: „Mische Mehl und Eier, dann backe es.“ Er macht sich keine Gedanken darüber, wie der Ofen sie mischt.
- Das Problem: Der Ofen (der KI-Chip) ist wählerisch. Er hat verschiedene Fächer (wie eine große Rührschüssel, ein kleiner Schneebesen und ein Heizelement). Wenn man ihm sagt, er solle alles in der großen Schüssel mischen, obwohl er auch den Schneebesen hätte benutzen können, verschwendet das Zeit.
- Axons Aufgabe: Axon nimmt Ihr einfaches Rezept und sagt: „Warte, ich kann das schneller machen.“ Es folgt nicht nur dem Rezept; es schreibt die Anweisungen um, um sie an die spezifischen Stärken des Ofens anzupassen.
2. Das „Magische Umstellen“ (Algebraische Transformationen)
Stellen Sie sich vor, Ihr Rezept sagt: „Zuerst schneide alle Karotten klein, dann mische sie in die Suppe, dann füge Salz hinzu.“
- Der alte Weg: Man schneidet, dann mischt, dann fügt Salz hinzu. Man muss die Schüssel ständig zwischen dem Schneidebrett und dem Herd hin- und hertragen.
- Axons Weg: Axon betrachtet die Mathematik und erkennt: „Hey, ich kann das Salz hinzufügen, während die Karotten gerade geschnitten werden, oder ich kann die Karotten in die Suppe mischen, bevor ich sie schneide, und es wird exakt denselben Geschmack haben!“
- Das Ergebnis: Durch das Umstellen der Reihenfolge der Operationen ermöglicht Axon dem Ofen, zwei Dinge gleichzeitig zu tun (wie Schneiden und Erhitzen simultan), anstatt eines nach dem anderen. Es beweist mathematisch, dass der Kuchen immer noch genauso schmeckt, auch wenn die Schritte anders sind.
3. Das „Packpuzzle“ (Tiling)
Stellen Sie sich vor, Sie haben einen riesigen Haufen Zutaten, aber Ihre Rührschüssel ist klein. Sie müssen Zutaten vom Vorratsschrank (langsam) in die Schüssel (schnell) bewegen, ein wenig mischen, sie wieder zurückbewegen und dies wiederholen.
- Das Problem: Wenn Sie zu viel auf einmal bewegen, läuft die Schüssel über. Wenn Sie zu wenig bewegen, verbringen Sie die ganze Zeit damit, zwischen Vorratsschrank und Schüssel hin- und herzulaufen.
- Axons Weg: Axon ist ein Meister im Packen. Es probiert Millionen verschiedener Möglichkeiten aus, die Zutaten in „Chargen“ (Tiles) zu gruppieren. Es fragt: „Was wäre, wenn ich 100 Karotten auf einmal bewege? Was wäre, wenn ich 50 bewege?“ Es testet diese verschiedenen Gruppierungen, um die perfekte Größe zu finden, die die Schüssel voll hält, ohne überzulaufen, und so die Wege zum Vorratsschrank minimiert.
4. Das „Bedienungshandbuch“ (ISA-Synthese)
Der Ofen hat ein spezifisches Bedienfeld mit Knöpfen wie „Hohe Hitze“, „Schneebesen“ und „Unterheben“.
- Das Problem: Manchmal sagt das Rezept „Mischen“, aber der Ofen hat keinen „Mischen“-Knopf. Er hat nur „Schneebesen“ und „Unterheben“.
- Axons Weg: Axon übersetzt „Mischen“ in die exakte Sequenz aus „Schneebesen und dann Unterheben“, die der Ofen versteht. Es findet sogar Abkürzungen, indem es erkennt, dass „Unterheben während der Bewegung zum Herd“ auf diesem speziellen Ofen ein einziger Tastendruck ist, was Zeit spart.
5. Der „Geschmackstest“ (Empirische Auswahl)
Dies ist der wichtigste Teil. Axon rät nicht einfach, welcher Weg der beste ist.
- Der Prozess: Es generiert tausende verschiedene Versionen des Rezepts (einige mit anderen Abläufen, einige mit anderen Chargengrößen, einige mit anderen Tastendrücken).
- Der Test: Es führt all diese Versionen tatsächlich auf dem echten Ofen (dem KI-Chip) aus und misst die Zeit. Es wählt die Version aus, die am schnellsten fertig ist. Es ist, als würde man 1.000 winzige Kuchen backen, um zu sehen, welcher in 5 Minuten statt in 10 Minuten aus dem Ofen kommt.
Die Ergebnisse: Wie viel schneller?
Das Paper testete Axon auf Amazon Trainium-Chips (eine Art KI-Ofen) mit 20 verschiedenen Arten von „Kuchen“ (KI-Aufgaben).
- Gegenüber dem Standard-Ofen: Im Vergleich zur Standard-Art des Backens (Amazon's Neuron Compiler) machte Axon einzelne Aufgaben bis zu 3,7-mal schneller und komplexe Aufgaben bis zu 19-mal schneller.
- Gegenüber Experten-Köchen: Selbst im Vergleich zu Rezepten, die von menschlichen Experten geschrieben wurden, die den Ofen in- und auswendig kennen, war Axon manchmal 1,35-mal schneller.
- Gegenüber anderen Robotern: Es schlug ein anderes automatisiertes Tool namens „Mirage“ um etwa 10 % im Durchschnitt und konnte 16 Arten von Kuchen backen, die das andere Tool überhaupt nicht bewältigen konnte.
Zusammenfassung
Axon ist ein Werkzeug, das eine einfache Beschreibung dessen, was eine KI tun muss, nimmt und automatisch den absolut schnellsten Weg findet, dies auf einem spezifischen Computerchip auszuführen. Dies geschieht, indem es mathematisch beweist, dass es Schritte umstellen kann, Millionen von Wegen zum Packen von Daten testet und dann den Code tatsächlich ausführt, um den Gewinner zu ermitteln. Es nimmt den Menschen die Notwendigkeit, Experten für die winzigen, langweiligen Details der Hardware zu sein, und lässt sie einfach das Rezept schreiben, während Axon das Backen übernimmt.
Ertrinken Sie in Arbeiten in Ihrem Fachgebiet?
Erhalten Sie tägliche Digests der neuesten Arbeiten passend zu Ihren Forschungsbegriffen — mit technischen Zusammenfassungen, in Ihrer Sprache.