Am 1. August fand in Peking der 9. Meet AI Compiler Technical Salon statt. Im Mittelpunkt dieser Veranstaltung standen die neuesten Fortschritte in der KI-Kompilierungstechnologie. Zahlreiche Experten aus Industrie und Forschungseinrichtungen präsentierten ihre Erkenntnisse zu Programmiersprachen, Operatorentwicklung, Kompilierungsoptimierung und Inferenzausführung und demonstrierten die kollaborative Entwicklung von KI-Compilern von der Hochsprachen-Ausdrucksweise bis zur Hardware-Ausführung.

In,Guo Hui und Xiao Hang, Forscher bei BAAI AI Compiler, präsentierten zum Thema „FlagTree: Triton-TLE Language Extension, Tile IR Backend and Compiler Optimization Practices“ eine Reihe von Untersuchungen, die das FlagTree-Team im Zusammenhang mit Triton durchgeführt hat.

Angesichts immer komplexerer Hardwarearchitekturen und unregelmäßigerer Modelloperatoren schlug das Team die Triton Language Extension (TLE) vor, um einen progressiven Kanal zwischen High-Level-DSL und Low-Level-Hardwaresteuerung herzustellen; gleichzeitig führten sie Kompilierungsoptimierungen am FlagTree-Compiler durch, wobei sie sich auf das TileIR-Backend, automatisches Tuning, Datenlayout, Befehlsplanung und andere Aspekte konzentrierten.

Das Hauptziel dieser Arbeit ist es, Entwicklern zu ermöglichen, bei Bedarf tiefer in Hardware-Details einzutauchen und gleichzeitig die Benutzerfreundlichkeit und das Community-Ökosystem von Triton zu erhalten. Außerdem soll eine schnelle Operatoroptimierung, ein architekturorientiertes Tuning und eine Optimierung auf nativer Codeebene mit einem relativ einheitlichen Entwicklungssystem abgedeckt werden.

HyperAI hat die geteilten Inhalte zusammengestellt und zusammengefasst, ohne deren ursprüngliche Bedeutung zu verändern.

Folge dem offiziellen WeChat-Konto "HyperAI" und antworte im Hintergrund mit dem Schlüsselwort "".0801 KI-CompilerSie können die Präsentations-PPT des autorisierten Referenten durch Anklicken von "..." abrufen.

Aufbauend auf Triton, wird das Gleichgewicht zwischen Abstraktion und Performance neu gewahrt.

In den letzten Jahren sind KI-Compiler mit einem immer deutlicher werdenden Widerspruch konfrontiert:Hardwarearchitekturen und Modelloperatoren werden immer komplexer, aber Entwickler hoffen weiterhin, mit höherwertigen DSLs eine Leistung zu erzielen, die der von Experten geschriebenen Kernels nahekommt.

Ist die Abstraktion zu hoch, kann der Compiler möglicherweise nicht genügend Informationen erhalten, um eine tiefgreifende Optimierung durchzuführen; ist die Abstraktion zu niedrig, greifen die Entwickler wieder auf den komplexen Entwicklungsmodus von CUDA oder herstellereigenen Sprachen zurück.

Der Erfolg von Triton liegt darin, die Entwicklung von GPU-Operatoren von der Thread-Ebene auf die Tile-Ebene zu heben.Die Benutzer verwenden die Python-DSL, um die Rechenbeziehungen zwischen Datenblöcken zu beschreiben, während Aufgaben wie Thread-Mapping, Registerzuweisung, Datenlayout, Pipelining und Synchronisierung hauptsächlich vom Compiler übernommen werden.

Dieser Ansatz senkt die Entwicklungshürde für leistungsstarke Operatoren und hat ein großes Community-Ökosystem gefördert. Mit der kontinuierlichen Weiterentwicklung von GPUs der nächsten Generation, domänenspezifischen Architekturen und inländisch produzierten KI-Chips stößt die ursprüngliche Abstraktion von Triton jedoch zunehmend an ihre Grenzen.

einerseits,Wenn das Compiler-Backend die Speicherstruktur, den Kommunikationsmechanismus oder die Recheneinheit der neuen Hardware noch nicht unterstützt, wird es für Frontend-Entwickler schwierig sein, diese Funktionen selbst zu nutzen.auf der anderen Seite,Immer mehr kritische Operatoren benötigen eine präzise Kontrolle über Speicherebenen, parallele Granularität, CTA-Kollaboration und die Überlappung von Kommunikation und Berechnung, und der ursprüngliche Triton-Code hat manchmal Schwierigkeiten, diese Absichten vollständig auszudrücken.

Das Aufkommen neuer Sprachen und domänenspezifischer Sprachen wie Gluon, TLX und TileLang spiegelt denselben Trend wider: Bei der Entwicklung von KI-Operatoren geht es nicht mehr nur um das Schreiben eines einzelnen Kernels, sondern auch um die Beschreibung von Datenlayout, paralleler Hierarchie, Pipeline, Kommunikationstopologie und Hardwareeigenschaften.

TLE versucht nicht, Triton zu ersetzen, sondern vielmehr dessen Syntax und Ökosystem schichtweise zu erweitern.Es besteht aus drei Schichten: TLE-Lite, TLE-Struct und TLE-Raw, die jeweils für Lightweight Semantic Hints, Architektur-Aware Control und Native Code-Level Optimization stehen.

TLE-Lite ist für Algorithmenentwickler und Szenarien mit schneller Optimierung konzipiert.Die Entwickler müssen sich nicht direkt mit der zugrundeliegenden Hardware auseinandersetzen; stattdessen können sie den Compiler mit expliziteren Strukturinformationen ergänzen.Beispielsweise kann der Zugriff auf einen Tensor über Subtiles erforderlich sein, eine Berechnung auf einem verteilten Mesh ausgeführt werden oder ein CTA eine Producer-Consumer-Pipeline nutzen. Am Beispiel von Subtile-Operationen lässt sich zeigen, dass Entwickler direkt einen logischen Subblock aus einem größeren Tensor extrahieren, Aktivierungs-, Normalisierungs- oder statistische Berechnungen durchführen und ihn anschließend zurückschreiben können, ohne Offsets manuell berechnen, Masken erstellen oder Grenzen behandeln zu müssen.

Da der Compiler erkennt, dass es sich um einen regulären Kachelzugriff handelt, kann er die Datenstruktur, Vektorisierung, Bankkonflikte und Registerwiederverwendung weiter optimieren. Dieser Ansatz eignet sich besonders für Sparse Attention, lokale Normalisierung, Blockstatistiken und Routing-Operatoren. In verteilten Szenarien verwendet TLE Device Mesh, um verschiedene Ebenen wie Knoten, GPUs, Blockcluster und Blöcke zu beschreiben und sie in einer einheitlichen, mehrdimensionalen Topologie zu organisieren.

Entwickler schreiben Programme mithilfe eines Mesh-basierten Ansatzes. Compiler und Laufzeitumgebung bilden die logische Topologie anschließend auf die tatsächliche Hardware und die Kommunikationsmechanismen ab. Dadurch sind Ringkommunikation, Barrier-Synchronisation und Sharded Access nicht mehr nur verstreute Ränge und Kommunikationsgruppen im Code, sondern werden zu analysierbaren, strukturierten Semantiken. Sobald Kommunikationsbeziehungen explizit ausgedrückt sind, kann der Compiler topologiebewusstes Scheduling, Kommunikations- und Berechnungsüberlappung, Barrier-Merging und Deadlock-Prüfung durchführen.

TLE-Lite abstrahiert die interne Zusammenarbeit der CTA mithilfe von Pipeline-Primitiven in ein Produzent-Konsument-Modell. Entwickler beschreiben hauptsächlich, wer die Daten erzeugt und wer sie verbraucht, während die zugrundeliegenden Barrier-, Pufferwiederverwendungs- und Synchronisierungsmechanismen vom Compiler übernommen werden.Dies bedeutet nicht, die zugrundeliegenden Fähigkeiten zu schützen, sondern sie in eine Programmstruktur umzuwandeln, die analysiert, verifiziert und optimiert werden kann.

Von leichtgewichtiger Semantik bis hin zu nativem Passthrough, wobei verschiedene Optimierungstiefen abgedeckt werden.

Wenn TLE-Lite sich primär mit plattformübergreifender semantischer Repräsentation befasst,TLE-Struct hingegen konzentriert sich mehr auf Architekturbewusstsein und Feinabstimmung.

Unterschiedliche GPUs, DSAs und KI-Beschleuniger unterscheiden sich erheblich in ihren Speicherebenen, Ausführungseinheiten, Synchronisationsmechanismen und On-Chip-Netzwerken. TLE-Struct stellt Entwicklern daher eine hierarchische Parallel- und Speicherstruktur zur Verfügung, die es ihnen ermöglicht, Datenlayout, Berechnungsabbildung und Speicherhierarchie explizit zu definieren, ohne sich direkt an die proprietäre Schnittstelle eines Anbieters binden zu müssen.

Beispielsweise kann derselbe lokale Puffer auf dem GPU-Backend dem gemeinsam genutzten Speicher und auf dem DSA-Backend dem Scratchpad oder dem On-Chip-SRAM zugeordnet werden. Der Benutzer gibt eine strukturierte Speicherabsicht an, und der Compiler ist dafür verantwortlich, diese in einen Adressraum und Speicherzugriffsanweisungen zu übersetzen, die für die Zielhardware geeignet sind.

Die Zählung von Experten in MoE ist ein typisches Szenario. Dieser Operator zählt im Wesentlichen die Anzahl der an verschiedene Experten weitergeleiteten Token und wird leicht von Faktoren wie dem Layout des gemeinsamen Speichers, gleichzeitigen Aktualisierungen, Bankkonflikten und blockübergreifender Aggregation beeinflusst.

Mit TLE-Struct können Entwickler das Layout der Zähler explizit organisieren und verschiedene Experten oder Token verschiedenen Bereichen des lokalen Speichers zuordnen; nachdem der Compiler diese Strukturinformationen erhalten hat, generiert er entsprechende Zugriffs- und Synchronisierungsmethoden.

TLE-Raw richtet sich an Experten für Leistungsoptimierung und erhält die Schnittstellen für herstellereigenen Code.

Manche Performance-Pfade erfordern die direkte Verwendung von CUDA, Assembler oder dedizierten intrinsischen Funktionen. Die erzwungene Umstrukturierung in eine übergeordnete DSL kann die Performance beeinträchtigen oder die Migrationskosten erhöhen. TLE-Raw ermöglicht es Entwicklern, nativen Code in die Triton/TLE-Architektur einzubetten oder direkt auf die Kompilierungspipelines der Hersteller zuzugreifen.

Nehmen wir All-Gather GEMM als Beispiel,Dieser Operator umfasst Kommunikation, Matrixberechnungen, Pufferung und Synchronisierungsmanagement. Entwickler können die zugrundeliegenden Kommunikationsfunktionen wiederverwenden und gleichzeitig strukturierte Tensoren und Tiles zur Darstellung von Berechnungen nutzen. Schließlich organisiert eine einheitliche Kompilierungspipeline die verschiedenen Teile zu aufrufbaren Operatoren.

Dieses geschichtete Design ermöglicht es Algorithmenentwicklern, Operatorentwicklern und Leistungsexperten, innerhalb desselben Systems unterschiedliche Optimierungstiefen zu wählen, ohne mit der niedrigsten Programmierebene beginnen zu müssen.

Leistungstests zeigen auch, dassDer gesamte Abstraktionsaufwand von TLE ist kontrollierbar.

Im Radix-Select-Test replizierte TLE den TensorRT-LLM-Algorithmus und erreichte dabei eine Leistung von ca. 85%–97% über verschiedene Datenstrukturen hinweg. Für Teams, die plattformübergreifende Wartung und schnelle Iterationen benötigen, bietet die Erzielung einer nahezu Expertenleistung mit einheitlicherem Code einen erheblichen technischen Mehrwert.

In einem SparseMLA-Szenario mit einem 128K-Kontext,TLE nutzt Pipeline-Primitive, um die Zusammenarbeit zwischen verschiedenen Ausführungsrollen auszudrücken und erreicht eine Leistung von ungefähr 90%, was der Leistung der FlashMLA-Baseline entspricht.

Das Team testete All-Gather auch auf einem einzelnen Knoten mit acht NVIDIA H100-Prozessoren und untersuchte die Integration von GEMM und All-Gather weiter. Der Fokus liegt nicht nur auf dem Austausch der Kommunikationsbibliothek, sondern auf der Einbindung der Kommunikation in Operator-Ausdrücke und Kompilierungsoptimierungen. Dadurch werden Datenübertragung, lokale Berechnung, Synchronisierung und die anschließende Verarbeitung zu einer Pipeline zusammengeführt.

Für Inferenzsysteme ist diese Art der Kommunikations- und Berechnungsfusion oft sinnvoller als die Verbesserung der Spitzenleistung eines einzelnen GEMM isoliert betrachtet, da die Benutzer letztendlich die End-to-End-Latenz wahrnehmen.

Optimierungspraktiken für die Flagtree-Kompilierung

Zusätzlich zu den TLE-Spracherweiterungen übernahm das FlagTree-Team auch zwei weitere Aufgaben:Zunächst integrierten wir das CUDA Tile IR-Backend; anschließend implementierten wir mehrere Triton-Kompilierungsoptimierungen für reale Modell-Workloads.

Das Kernkonzept von CUDA Tile IR besteht darin, Programmen die Möglichkeit zu geben, Tiles auszudrücken, anstatt die zugrunde liegende Thread-Zuordnung vorab festzulegen.Die Eingabe besteht aus einem Tile-Programm und einem dreidimensionalen Tile-Grid. Innerhalb von Dialect werden Berechnungen, Datenansichten und notwendige Abhängigkeiten durch Tile Compute, View und Token-Ordered Operation (TKO) ausgedrückt.

Unter ihnen beschreibt TensorView den globalen Zeiger, Shape und Stride; PartitionView fügt darüber hinaus Block-Mapping-Funktionen hinzu; Token wird verwendet, um die Abhängigkeiten zwischen verwandten TKOs einzuschränken; und Memory Model spezifiziert Speichersemantik und Gültigkeitsbereich separat.

FlagTree ersetzt das native Triton-NVIDIA-Backend nicht vollständig. Stattdessen fügt es dem bestehenden System einen TileIR-Pfad hinzu und überschreibt die TileIR-View- und Token-Schnittstellen mithilfe von TLE-Primitiven. Kernel, die die Anforderungen erfüllen, können auf das TileIR-Backend zugreifen; nicht unterstützte Kernel greifen auf natives CUDA zurück. Beide Pfade können im Compiler parallel verwendet werden.

Hinsichtlich der automatischen OptimierungDas Team schlug FlagOSTune vor, um den Konflikt zwischen Abdeckung und Kosten in realen Modellen von Triton Autotune zu lösen.

Bei der tatsächlichen Modellinferenz kann derselbe Operator einer Vielzahl unterschiedlicher Formen entsprechen. Die Statistiken des Teams zeigen, dass es in sechs Modellen und vier Arten von Inferenzszenarien 1994 einzigartige MM-Formen gibt, was weit über die Abdeckung täglicher Benchmarks hinausgeht. Wird die Kandidatenkonfiguration direkt erweitert, wächst der theoretische Suchumfang rasch auf Millionen von Mengen an.

FlagOSTune verbessert die Leistung durch die Erweiterung des Suchraums und reduziert die Suchkosten durch eine Kombination aus Modellvorhersage und begrenzten Tests in realen Umgebungen. Das System verwendet zunächst XGBoost, um Kandidatenkonfigurationen zu bewerten. Anschließend werden nur wenige der vielversprechendsten Konfigurationen zur Kompilierung und zum GPU-Test weitergeleitet, bevor die Suche mithilfe eines genetischen Algorithmus fortgesetzt wird.

Die Leistung mehrerer Operatoren wurde auf verschiedenen Rechenleistungssystemen von NVIDIA, Moore Threads und Muxi verbessert, mit Beschleunigungen im Bereich von 1,21 bis 7,35x.In Experimenten mit verschiedenen Körperformen unter Verwendung des NVIDIA H20 MM-Operators wurde die Anzahl der Konfigurationen im Suchraum von über 620.000 auf 4.070 reduziert, was zu einer 120-fachen Steigerung der Effizienz bei der Optimierung führte, und zwar nahezu ohne Leistungsverlust.

Im Hinblick auf das Datenlayout konzentrierte sich das Team auf die Optimierung des Leistungsaufwands bei der Datenlayouttransformation.

In Triton beinhaltet `convert_layout` typischerweise eine threadübergreifende Datenumstrukturierung, die das Schreiben in den gemeinsamen Speicher, die Synchronisierung und das anschließende Auslesen erfordert – ein Vorgang, der Kosten verursacht. Bei kleinen oder speicherintensiven Operatoren können selbst wenige Layout-Transformationen zu einem Leistungsengpass führen.

FlagTree erweitert den Mechanismus zum Entfernen von Layoutkonvertierungen durch die Verwendung von Techniken wie Kostenmodellierung, Rückpropagation und lokaler gemeinsamer Lösung, um unnötige Datenlayoutkonvertierungen zu reduzieren.In Tests mit mehr als 100 Bedienern lag die Netto-Konversionseliminierungsrate bei etwa 68%–79%, mit einer Leistungsverbesserung von bis zu 71%.

Eine weitere Optimierung ist die Befehlsneuanordnung. Das Abrollen von Schleifen kopiert lediglich den Schleifenkörper und führt nicht automatisch mehrere Ladevorgänge im Voraus aus. Mit aktivierter Compiler-Neuanordnung können unabhängige Ladevorgänge im Voraus ausgeführt werden, wodurch nachfolgende Berechnungen verschachtelt werden und die Hardware-Pipeline genutzt wird, um die Speicherzugriffslatenz zu verbergen.Bei den drei typischen Operatoren führt diese Optimierung zu einer durchschnittlichen Beschleunigung um das 1,19- bis 1,61-Fache, mit einer maximalen Beschleunigung um das 2-Fache.

Das Team optimierte außerdem die Leistung des Fused Marlin MoE-Operators für reelle Datenformen in DeepSeek-V4-Flash und reduzierte die Datenladezeit durch Fragmentzusammenführung. In Kombination mit Optimierungen auf Algorithmusebene,Im Vergleich zu vLLM CUDA erzielte FlagGems bei allen 53 getesteten Formen eine Beschleunigung, mit einer durchschnittlichen Verbesserung von etwa dem 1,208-Fachen, wenn man die Häufigkeit der Formen berücksichtigt.

Im DeepSeek-V4-Flash-End-to-End-Test auf der NVIDIA H20,Bei TP=4 verringert sich die Latenz des ersten Tokens um 20,221 TP3T und der Gesamtdurchsatz erhöht sich um 12,631 TP3T; bei TP=8 verringert sich die Latenz des ersten Tokens um 12,871 TP3T und der Gesamtdurchsatz erhöht sich um 7,651 TP3T.

Zukünftige Berufsaussichten

Die nächste Phase,FlagTree wird sich auf die Weiterentwicklung des NUMA-Programmiermodells und des MegaKernel-Compilers konzentrieren.

Bei Systemen mit mehreren Knoten und GPUs sowie GPU-internen Clustern und lokalem SRAM variiert die Nähe der Datenzugriffe erheblich. TLE beschreibt explizit, wie Tensoren entlang des Gitters segmentiert und zwischen verschiedenen Verteilungsmethoden umgeschaltet werden. Dadurch kann der Compiler bestimmen, welche Daten für die Berechnung nahe beieinander liegen und welche Kommunikationsvorgänge mit der Berechnung zusammengeführt oder überlappt werden können.

Der MegaKernel-Compiler erweitert den Optimierungsbereich von einem einzelnen Kernel auf einen Modellausführungsprozess. Er identifiziert Berechnungen, die gemeinsam optimiert werden können, zerlegt sie in spezifische Aufgaben und entscheidet dann einheitlich über deren Zuweisung, Parallelisierung und Verarbeitung in einer Pipeline. Dadurch generiert er letztendlich einen oder mehrere MegaKernel.

TLE fungiert als Brücke zwischen Frontend-Semantik und Hardwarearchitektur. Schnittstellen für Tile, Mesh, Pipeline, Memory Layout und native Funktionen ermöglichen es dem Compiler, nicht nur eine Reihe unabhängiger Operatoren zu sehen, sondern einen Aufgabengraphen, der geplant, platziert, synchronisiert und zusammengeführt werden kann.

Von den leichten semantischen Hinweisen von TLE-Lite über die architekturorientierte Steuerung von TLE-Struct bis hin zur nativen Funktionsweiterleitung von TLE-Raw, kombiniert mit verschiedenen Kompilierungsoptimierungstechniken, versucht FlagTree, ein flexibleres, geschichtetes System zwischen der Benutzerfreundlichkeit, der plattformübergreifenden Migration und der maximalen Leistung von Triton zu etablieren und neue Hardwarefunktionen schneller in die tatsächliche Leistung von Modellen und Operatoren umzuwandeln.