Meta stellt TLX vor, einen hardwarenahen Compiler für GPU-Workloads

Eine Zwischenschicht zwischen CUDA und Triton

Meta beschreibt TLX als Antwort auf die wachsende Lücke zwischen hardwarenaher GPU-Programmierung und höher angesiedelten Werkzeugen, die auf eine hohe Entwicklerproduktivität ausgelegt sind. Moderne GPUs sind zunehmend auf Funktionen wie Warp-Spezialisierung, Tensor Memory, asynchrone Speicherbewegungen und Ausführung auf Cluster-Ebene angewiesen. Ihre effiziente Nutzung erfordert, dass Entwickler mehrere Ausführungsinstanzen koordinieren, statt sich nur auf einzelne Threads oder Tensor-Operationen zu konzentrieren.

Nach Angaben von Meta kann CUDA diese Strukturen abbilden. Dafür ist jedoch mitunter umfangreicher, architekturspezifischer Orchestrierungscode erforderlich. Triton bietet einen produktiveren Ansatz, doch einige der für aktuelle Hardware nötigen Ausführungsbeziehungen bleiben hinter Compiler-Abstraktionen verborgen. TLX soll die fehlende Schicht zwischen den beiden Modellen bereitstellen.

Meta bezeichnet diesen Ansatz als Multiple Instructions, Multiple Warps oder MIMW. Das Modell behandelt Warp-Gruppen als eigenständige Ausführungsinstanzen und ermöglicht es Programmierern zu beschreiben, wie diese zusammenarbeiten, kommunizieren, synchronisieren und Ressourcen teilen.

Für neuere GPU-Architekturen entwickelt

Nach Angaben von Meta wurde der Bedarf für diese Abstraktion deutlicher, als das Unternehmen NVIDIA-Hopper- und -Blackwell-GPUs in seiner Infrastruktur für maschinelles Lernen einführte. Hopper brachte Warp-Spezialisierung und den Tensor Memory Accelerator, während Blackwell unter anderem Tensor Memory, Ausführung auf Cluster-Ebene, gepaarte CTA-Ausführung und weitere asynchrone Mechanismen ergänzte.

Diese Funktionen beeinflussen den Aufbau von Hochleistungskernels. Bei einem FlashAttention-Kernel kann beispielsweise eine Warp-Gruppe Speicherübertragungen anstoßen, während eine andere Berechnungen auf Tensor Cores ausführt. Ein Blackwell-GEMM-Kernel muss möglicherweise gleichzeitig Warp-Gruppen, Tensor-Memory-Zuweisungen und Cluster-Ressourcen koordinieren.

TLX ergänzt Triton um eine explizite Orchestrierung auf Warp-Gruppen- und Cluster-Ebene. Meta zufolge können Entwickler dadurch neue Hardwarefunktionen schneller nutzen, ohne bei der Entwicklungsgeschwindigkeit Abstriche machen zu müssen.

Produktiver Einsatz im gesamten ML-Stack von Meta

Metas Software-Stack für maschinelles Lernen basiert auf PyTorch. Der Modellcode durchläuft zunächst die Aufzeichnung des Rechengraphen, Optimierung, Kompilierung und Absenkung, bevor daraus GPU-Kernels werden. Nach Angaben des Unternehmens hat sich Triton von einer experimentellen Kernelsprache zu einem grundlegenden Bestandteil dieses Stacks entwickelt. Ein erheblicher Anteil der produktiven GPU-Kernels, darunter vom Compiler erzeugte und manuell optimierte Kernels, wird in Triton implementiert.

Meta zufolge treibt TLX inzwischen produktive Workloads im gesamten Unternehmen an, vom Training für KI und RecSys über die RecSys-Inferenz bis zum Vortraining groß angelegter LLMs. Das Unternehmen erklärt außerdem, TLX habe bei diesen Workloads Leistungssteigerungen ermöglicht und werde auf NVIDIA- und AMD-GPUs eingesetzt.

Für Nutzer von Metas Infrastruktur für maschinelles Lernen signalisiert die Ankündigung den Versuch, hardwarespezifische GPU-Optimierung verfügbar zu machen, ohne dass jeder Entwickler direkt auf der Detailstufe von CUDA arbeiten muss. Für den breiteren Markt für GPU-Software spiegelt TLX Metas Einschätzung wider, dass neuere Beschleunigerfunktionen Programmierabstraktionen erfordern, die auf der Zusammenarbeit zwischen Warp-Gruppen basieren, statt sich allein auf die Automatisierung durch den Compiler zu verlassen.