Jak okiełznać Blackwella bez szaleństwa od CUDA

Pisząc kernele GPU, większość z nas od razu wyobraża sobie niekończące się strony kodu CUDA C++, żonglowanie wskaźnikami i bolesne debugowanie wycieków pamięci czy warunków wyścigu. Sytuacja staje się jeszcze bardziej złożona, gdy na rynek trafiają potwory takie jak architektura NVIDIA Blackwell (sm_100a). Jest w niej tyle asynchronii i nowych silników danych, że stare podejścia „po prostu uruchom wątek na rdzeniu" nie wycisną już nawet połowy wydajności.

Niedawno natknąłem się na projekt modern-gpu-programming-for-mlsys od zespołu MLC-AI. To w zasadzie nie tylko repozytorium, ale żywy interaktywny tutorial dotyczący nowoczesnego programowania GPU. Facetowie poszli na całość: uczą pisać najnowocześniejsze kernele (SOTA) przy użyciu Pythona i nowego DSL-a o nazwie TIRx.

Dlaczego to nie jest kolejny tutorial

Zwykle kursy GPU utykają przy klasycznym mnożeniu macierzy z lat 2010. Tutaj fokus przesunięto na sprzęt jutra. W centrum uwagi jest architektura Blackwell, jej hierarchia pamięci, rdzenie tensorowe i mechanizmy asynchronicznego przesyłu danych.

Autorzy proponują ścieżkę od zrozumienia sprzętu do pisania kerneli, które faktycznie działają w nowoczesnych systemach ML. Zamiast walczyć z niskopoziomowym C++, piszesz w TIRx (Tensor IR next). To rozszerzenie Apache TVM, które pozwala opisywać logikę kernela w Pythonie, zachowując kontrolę nad rejestrami, cache i silnikami tensorowymi.

Co znajdziesz w tym samouczku

Cała zawartość podzielona jest na logiczne bloki, które stopniowo podnoszą poziom trudności.

Najpierw omawiany jest model wykonania i pamięci. Nie ma tu nudnych definicji, za to mnóstwo praktyki z TMA (Tensor Memory Accelerator) i zrozumienia, jak dane faktycznie docierają do jednostek obliczeniowych. Co ciekawe, autorzy szczegółowo skupiają się na modelu Roofline — pomaga zrozumieć, na czym kod jest wąskim gardłem: przepustowością pamięci czy mocą obliczeniową rdzeni.

Następnie przychodzi czas na głębokie zanurzenie w TIRx. Uczysz się pracować z kaflami (TileLayout), osiami i mechanizmem swizzle (sprytne przetasowanie danych w pamięci, aby uniknąć konfliktów banków). Na końcu masz działający kernel GEMM (mnożenie macierzy).

Potem przychodzi prawdziwe „mięso". Autorzy pokazują, jak zamienić zwykły GEMM w kombajn wysokiej wydajności. W grę wchodzą pipeline'y TMA, specjalizacja warpów i grupowanie CTA w klastry. Jeśli te terminy brzmią jak zaklęcia, tutorial ładnie je rozłoży na czynniki pierwsze.

W finale tworzysz Flash Attention 4. To dokładnie to, na czym działają nowoczesne duże modele językowe. Kernel zawiera dwa etapy MMA (Matrix-Matrix Accumulation), online softmax i obsługę GQA (Grouped Query Attention).

Stos technologiczny i funkcje

Projekt jest ściśle powiązany z ekosystemem Apache TVM. Aby uruchomić przykłady, potrzebujesz:

1. Biblioteka apache-tvm (TIRx jest w niej zawarty).
2. PyTorch do weryfikacji wyników i przygotowywania danych.
3. GPU Blackwell (np. B200), jeśli chcesz uruchomić kod na prawdziwym sprzęcie.

Jeśli nie masz pod ręką Blackwella (co przy obecnym stanie rzeczy jest logiczne dla większości z nas), repozytorium jest nadal niezwykle użyteczne jako baza teoretyczna. Logika optymalizacji pamięci i planowania obliczeń ma zastosowanie również do starszych architektur, choć niektóre funkcje jak TMA są specyficzne dla nowszych generacji NVIDIA.

Aby zbudować książkę lokalnie i czytać ją w wygodnym formacie, wystarczy standardowy Sphinx:

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

Czy warto poświęcić czas

Często spotykam developerów infrastruktury ML, którzy boją się zanurzyć w pisaniu kerneli, uważając to za domenę „C++ bogów". Ten projekt dowodzi, że przy odpowiednich abstrakcjach (jak TIRx) i zrozumieniu architektury można pisać kod światowej klasy, pozostając w ekosystemie Pythona.

Projekt będzie szczególnie przydatny dla osób pracujących nad optymalizacją sieci neuronowych „pod maską" lub chcących zrozumieć, dlaczego Flash Attention jest szybszy od zwykłego podejścia. Nawet jeśli nie planujesz jutro wydać własnej biblioteki Blackwella, przejście przez to, jak działają tensorowa pamięć i asynchroniczne kolejki, naprawdę rozjaśnia umysł i zmienia perspektywę na obliczenia wysokiej wydajności.

Jedyną wadą jest to, że próg wejścia wciąż jest wysoki. To nie jest „GPU dla dummies", ale raczej „GPU dla tych, którzy chcą zostać profesjonalistami". Ale jeśli jesteś gotowy, żeby rozgryźć, jak bajty faktycznie przemieszczają się wewnątrz karty graficznej, trudno teraz znaleźć lepszy materiał na ten temat.