IBMs Spyre-Beschleuniger wechselt zu PyTorch Inductor und kommt mit einem rund siebenmal kleineren Compiler aus

Vom Eigenbau-Stack zu Inductor

Bisher setzte torch-spyre auf sendnn, einen separaten, komplett von Grund auf entwickelten Compiler-Stack. Nun kommt ein Open-Source-Backend zum Einsatz, das Inductor erweitert, den Compiler hinter PyTorchs torch.compile. Die Autoren berichten, dass der größte Teil von Inductor unverändert übernommen werden konnte. Die eigentliche Arbeit konzentrierte sich auf drei Stellen, an denen die Hardware von Spyre dem Compiler grundlegend andere Fragen stellt als eine GPU.

Zugleich betonen sie, dass es sich nicht um einen simplen Austausch handelte. Auf den unteren Ebenen von Inductor und bei der Codegenerierung waren weiterhin umfangreiche Erweiterungen außerhalb des PyTorch-Codebaums (out-of-tree) nötig.

IBM hat Spyre als komplett neu konzipierten KI-Beschleuniger für Inferenz-Workloads in Unternehmen entworfen, mit Fokus auf Energieeffizienz. Dem Beitrag zufolge handelt es sich um einen 75-Watt-Beschleuniger mit PCIe-Anbindung, der weder eine spezielle Stromversorgung noch eine besondere Kühlung benötigt. Er ist ein Dataflow-Beschleuniger aus 32 kleinen Kernen, die über einen Ring verbunden sind und jeweils einen eigenen lokalen Scratchpad-Speicher von 2 MB statt eines gemeinsamen Caches besitzen. Daten werden in Blöcken fester Größe bewegt, den sogenannten Sticks: 128 Byte, das entspricht 64 Werten in fp16.

Wie der Compiler aufgebaut ist

Der Compiler arbeitet in zwei Stufen. In der ersten übersetzen Inductor-Erweiterungen, die im torch-spyre-Repository liegen, einen per torch.compile erzeugten Graphen in eine Zwischendarstellung namens SuperDSC. In der zweiten wandelt ein proprietärer Compiler namens DeepTools SuperDSC in Instruktionen um, die auf der Spyre-Hardware laufen. Der Beitrag konzentriert sich auf die erste Stufe.

Die Erweiterungen klinken sich in die Standard-Pipeline ein: Dynamo verfolgt das Modell, AOTAutograd partitioniert es, und Inductor überführt den FX-Graphen in seine LoopLevelIR. Torch-spyre ergänzt drei Erweiterungspunkte:

• Decompositions: Spyre-spezifische Operatorzerlegungen auf der FX-Graph-Ebene, etwa die Abbildung von layernorm auf Folgen primitiver Spyre-Operationen.
• Lowerings: Spyre-spezifische Operator-Lowerings vom FX-Graphen in die LoopLevelIR.
• CustomPreSchedulingPasses: der Großteil des Spyre-spezifischen Compilers, darunter Layout-Zuweisung, Tiling, Arbeitsaufteilung und Scratchpad-Allokation.

Letzterer ist derzeit über einen Monkey-Patch auf Inductors GraphLowering.codegen() eingebunden und kein vollwertiger Hook. Inductors eigener _pre_fusion_custom_pass läuft nach Angaben der Autoren zu spät, da er erst ausgeführt wird, nachdem die LoopLevelIR in SchedulerNodes umgewandelt wurde. Danach unterstützt Inductor kein umfassendes Umschreiben des Graphen mehr. Das Team arbeitet mit dem PyTorch-Upstream daran, diese Erweiterungspunkte zu formalisieren, damit auch andere Out-of-tree-Hardware-Backends sie nutzen können.

Drei Hardware-Unterschiede hinter dem Design

Laut dem Beitrag ist Inductors Vorstellung von Hardware von GPUs geprägt, und drei Elemente davon lassen sich nicht auf Spyre übertragen.

Sticks und Layouts. Spyre kann Speicher nicht auf Ebene einzelner Elemente adressieren. Der Stick, ein ausgerichteter 128-Byte-Block, ist die Einheit für Datenbewegung und Berechnung, weshalb Inductors Indexarithmetik in Sticks denken muss. Eine punktweise Operation verlangt, dass alle Eingaben und die Ausgabe dieselbe Stick-Dimension teilen, nämlich die innerste. Kommen Eingaben mit inkompatiblen Sticks an, muss der Compiler ein Restickify einfügen, eine Datenbewegungsoperation, die die Sticks selbst neu formt. Die Wahl der Stick-Dimensionen über einen ganzen Graphen hinweg, sodass möglichst wenig Daten bewegt werden, ist ein globales Optimierungsproblem. Auch auf die Geschwindigkeit wirkt sich das Layout aus: Bei einer gebatchten Matrixmultiplikation, bei der alle Tensoren bereits im Gerätespeicher liegen, senkte eine andere Aufteilung (Tiling) der Eingabe- und Ausgabetensoren die Kernel-Zeit laut den Autoren über eine Reihe von Batchgrößen und Formen hinweg um 29 bis 69 Prozent.

Scratchpad-Speicher. Jeder Kern besitzt einen kleinen lokalen Scratchpad, dahinter liegt nur der globale Gerätespeicher. Jeder Transfer findet statt, weil der Compiler eine explizite Anweisung erzeugt hat. Ohne hardwareverwaltete Zwischenschicht muss ein Wert, der nicht im Scratchpad bleiben soll, zurückgeschrieben und erneut geladen werden. Da Workloads auf Spyre häufig speichergebunden statt rechengebunden sind, geht an dieser Stelle laut dem Beitrag oft die meiste Zeit verloren.

Kernzuweisung. Spyres Kerne arbeiten als statische Dataflow-Maschine, ohne einen Laufzeit-Scheduler nach GPU-Vorbild, der die Last auf Streaming-Multiprozessoren verteilt. Die Aufteilung der Arbeit auf die Kerne muss daher zur Compile-Zeit festgelegt werden.

Was das bedeutet

Für Teams, die Beschleuniger außerhalb des PyTorch-Codebaums entwickeln, bietet der Beitrag eine Fallstudie: Ein neuer Chip mit einem stark abweichenden Ausführungsmodell könnte den Großteil von Inductor wiederverwenden und den Aufwand auf wenige hardwarespezifische Passes konzentrieren. Die Autoren nennen Inductor „genau das richtige Fundament“ für die Arbeit, und der etwa siebenmal kleinere Compiler ist ihr Maß für den gewonnenen Hebel. Die geplante Upstream-Arbeit an formalen Erweiterungspunkten soll, falls sie abgeschlossen wird, diesen Ansatz für andere Hardware-Backends erleichtern.