Hoe je Blackwell temt zonder gek te worden van CUDA

Als het gaat om het schrijven van GPU-kernels, denken de meesten van ons meteen aan eindeloze pagina's CUDA C++-code, het managen van pointers en het frustrerende debuggen van memory leaks of race conditions. De situatie wordt nog complexer wanneer monsters zoals de NVIDIA Blackwell-architectuur (sm_100a) op de markt verschijnen. Er is zoveel asynchronie en nieuwe data engines binnenin dat de oude "start gewoon een thread naar een core"-benaderingen niet eens de helft van de prestaties meer uit de hardware persen.

Onlangs stuitte ik op het modern-gpu-programming-for-mlsys-project van het MLC-AI-team. In wezen is het niet zomaar een repository, maar een levend interactief tutorial over moderne GPU-programmering. De jongens hebben er alles aan gedaan: ze leren je SOTA (state-of-the-art) kernels schrijven met Python en een nieuwe DSL genaamd TIRx.

Waarom dit niet zomaar een andere tutorial is

Meestal raken GPU-cursussen vast aan de klassieke matrixvermenigvuldiging uit de jaren 2010. Hier is de focus verschoven naar de hardware van morgen. Het middelpunt van de aandacht is de Blackwell-architectuur, de memory hierarchy, tensor cores en asynchrone data movement-mechanismen.

De auteurs stellen een pad voor van het begrijpen van de hardware tot het schrijven van kernels die daadwerkelijk werken in moderne ML-systemen. In plaats van te vechten met low-level C++, schrijf je in TIRx (Tensor IR next). Dit is een Apache TVM-extensie waarmee je kernel-logica in Python kunt beschrijven terwijl je controle behoudt over registers, cache en tensor engines.

Wat zit er in deze tutorial

Alle inhoud is opgedeeld in logische blokken die de moeilijkheidsgraad geleidelijk verhogen.

Ten eerste behandelen ze het execution en memory model. Geen saaie definities hier, maar volop praktijkervaring met TMA (Tensor Memory Accelerator) en het begrijpen hoe data daadwerkelijk de compute units bereikt. Interessant is dat de auteurs uitgebreid stil staan bij het Roofline model — het helpt begrijpen waar je code door beperkt wordt: memory bandwidth of de pure kracht van de cores.

Daarna volgt een diepgaande verkenning van TIRx. Je leert werken met tiles (TileLayout), assen en het swizzle-mechanisme (slimme data-shuffling in het geheugen om bank conflicts te voorkomen). Aan het einde krijg je een werkende GEMM (matrixvermenigvuldiging) kernel.

Daarna komt het echte "werk". De auteurs laten zien hoe je een gewone GEMM omzet in een high-performance monster. TMA pipelines, warp specialization en het groeperen van CTA's in clusters komen aan bod. Als deze termen klinken als toverformules, legt de tutorial ze netjes uit.

In de finale creëer je Flash Attention 4. Dit is precies waar moderne large language models op draaien. De kernel omvat twee MMA (Matrix-Matrix Accumulation) stages, online softmax en GQA (Grouped Query Attention) ondersteuning.

Technische stack en features

Het project is nauw gekoppeld aan het Apache TVM-ecosysteem. Om de voorbeelden te draaien heb je nodig:

1. Library apache-tvm (TIRx zit hierin).
2. PyTorch voor het verifiëren van resultaten en het voorbereiden van data.
3. Een Blackwell GPU (bijv. B200) als je de code op echte hardware wilt draaien.

Als je geen Blackwell bij de hand hebt (wat logisch is voor de meesten van ons nu), is de repository nog steeds uiterst nuttig als theoretische basis. De logica van geheugenoptimalisatie en compute scheduling is ook van toepassing op oudere architecturen, hoewel sommige features zoals TMA specifiek zijn voor nieuwere NVIDIA-generaties.

Om het boek lokaal te bouwen en in een handig formaat te lezen, is standaard Sphinx voldoende:

pip install -r requirements-docs.txt
sphinx-build -b html . _build/html

Is het de moeite waard

Ik kom vaak ML-infrastructuurontwikkelaars tegen die bang zijn om zich te verdiepen in kernel-schrijven, omdat ze dit beschouwen als het domein van "C++-goden". Dit project bewijst dat je met de juiste abstracties (zoals TIRx) en begrip van de architectuur wereldklasse code kunt schrijven terwijl je in het Python-ecosysteem blijft.

Het project is vooral nuttig voor degenen die werken aan het optimaliseren van neurale netwerken "onder de motorkap" of willen begrijpen waarom Flash Attention sneller is dan de gebruikelijke aanpak. Zelfs als je niet van plan bent om morgen je eigen Blackwell-bibliotheek uit te brengen, helpt het doorgronden van hoe tensor memory en asynchrone queues werken echt om je mindset te veranderen en je perspectief op high-performance computing te verbreden.

Het enige nadeel is dat de instapdrempel nog steeds hoog is. Dit is niet "GPU voor dummies", maar eerder "GPU voor degenen die profs willen worden." Maar als je bereid bent om uit te zoeken hoe bytes daadwerkelijk bewegen binnen een grafische kaart, is het moeilijk om beter materiaal over dit onderwerp te vinden.