Blackwell zähmen, ohne den Verstand an CUDA zu verlieren
Beim Schreiben von GPU-Kernels denken die meisten von uns sofort an endlosen CUDA-C++-Code, das Jonglieren mit Pointern und das qualvolle Debugging von Memory Leaks oder Race Conditions. Wenn dann auch noch Monster wie die NVIDIA-Blackwell-Architektur (sm_100a) auf den Markt kommen, wird die Sache noch komplexer. Es gibt so viel Asynchronität und neue Daten-Engines im Inneren, dass die alten „einfach einen Thread auf einen Core loslassen"-Ansätze nicht mehr einmal die Hälfte der Leistung herausholen.
Kürzlich bin ich auf das Projekt modern-gpu-programming-for-mlsys vom MLC-AI-Team gestoßen. Im Grunde ist es nicht nur ein Repository, sondern ein lebendiges, interaktives Tutorial zur modernen GPU-Programmierung. Die Jungs haben alles gegeben: Sie bringen dir bei, SOTA-Kernels (State-of-the-Art) mit Python und einer neuen DSL namens TIRx zu schreiben.
Warum das nicht nur ein weiteres Tutorial ist
Normalerweise bleiben GPU-Kurse bei der klassischen Matrixmultiplikation aus den 2010er Jahren hängen. Hier liegt der Fokus auf der Hardware von morgen. Im Mittelpunkt steht die Blackwell-Architektur, ihre Memory-Hierarchie, Tensor Cores und asynchrone Datenbewegungsmechanismen.
Die Autoren schlagen einen Weg vor, der vom Verständnis der Hardware zum Schreiben von Kernels führt, die in modernen ML-Systemen tatsächlich funktionieren. Anstatt mit Low-Level-C++ zu kämpfen, schreibst du in TIRx (Tensor IR next). Das ist eine Apache-TVM-Erweiterung, mit der du die Kernel-Logik in Python beschreiben kannst, während du die Kontrolle über Register, Cache und Tensor-Engines behältst.
Was steckt in diesem Tutorial
Der gesamte Inhalt ist in logische Blöcke unterteilt, die den Schwierigkeitsgrad schrittweise erhöhen.
Zuerst werden das Ausführungs- und das Memory-Modell behandelt. Keine langweiligen Definitionen, sondern viel Praxis mit TMA (Tensor Memory Accelerator) und dem Verständnis, wie Daten tatsächlich zu den Compute-Units gelangen. Interessanterweise gehen die Autoren detailliert auf das Roofline-Modell ein – es hilft zu verstehen, worauf dein Code beschränkt ist: Memory-Bandwidth oder die rohe Rechenleistung der Cores.
Dann folgt ein tiefer Einstieg in TIRx. Du lernst, mit Tiles (TileLayout), Achsen und dem Swizzle-Mechanismus zu arbeiten (cleveres Daten-Shuffling im Memory, um Bank-Konflikte zu vermeiden). Am Ende hast du einen funktionierenden GEMM-Kernel (Matrixmultiplikation).
Danach kommt das echte „Fleisch". Die Autoren zeigen, wie man einen normalen GEMM in ein Hochleistungsmonster verwandelt. TMA-Pipelines, Warp-Specialization und das Gruppieren von CTAs in Cluster kommen zum Einsatz. Wenn diese Begriffe wie Beschwörungsformeln klingen, wird das Tutorial sie übersichtlich aufbereiten.
Im Finale erstellst du Flash Attention 4. Genau das, worauf moderne Large Language Models laufen. Der Kernel umfasst zwei MMA-Stages (Matrix-Matrix-Accumulation), Online-Softmax und GQA-Support (Grouped Query Attention).
Technischer Stack und Features
Das Projekt ist eng mit dem Apache-TVM-Ökosystem verknüpft. Um die Beispiele auszuführen, brauchst du:
- Bibliothek
apache-tvm(TIRx ist darin enthalten). - PyTorch zum Verifizieren der Ergebnisse und Vorbereiten der Daten.
- Eine Blackwell-GPU (z. B. B200), wenn du den Code auf echter Hardware ausführen möchtest.
Wenn du keine Blackwell zur Hand hast (was für die meisten von uns gerade logisch ist), ist das Repository trotzdem extrem nützlich als theoretische Grundlage. Die Logik der Memory-Optimierung und der Compute-Scheduling gilt auch für ältere Architekturen, auch wenn einige Features wie TMA spezifisch für neuere NVIDIA-Generationen sind.
Um das Buch lokal zu bauen und in einem komfortablen Format zu lesen, reicht das Standard-Sphinx:
pip install -r requirements-docs.txt
sphinx-build -b html . _build/html
Lohnt es sich?
Ich stoße oft auf ML-Infrastructure-Entwickler, die sich nicht an Kernel-Programmierung heranwagen und sie als Vorrecht der „C++-Götter" betrachten. Dieses Projekt beweist, dass man mit den richtigen Abstraktionen (wie TIRx) und dem Verständnis der Architektur erstklassigen Code schreiben kann, während man im Python-Ökosystem bleibt.
Das Projekt ist besonders nützlich für diejenigen, die neuronale Netze „unter der Haube" optimieren oder verstehen wollen, warum Flash Attention schneller ist als der übliche Ansatz. Selbst wenn du nicht vorhast, morgen deine eigene Blackwell-Bibliothek zu veröffentlichen, räumt das Durcharbeiten, wie Tensor Memory und asynchrone Queues funktionieren, richtig auf und verändert deine Perspektive auf High-Performance-Computing.
Der einzige Nachteil ist, dass die Einstiegshürde immer noch hoch ist. Das ist kein „GPU für Dummies", sondern eher „GPU für die, die Profis werden wollen". Aber wenn du bereit bist herauszufinden, wie Bytes tatsächlich in einer Grafikkarte bewegt werden, ist es gerade schwer, besseres Material zu diesem Thema zu finden.
Ähnliche Projekte