Zwei Zahlen auf einer GPU
Du weißt bereits, wie man c = a + b in Python macht. Du hast es tausendmal gemacht. Hier reden wir darüber in PyTorch-Tensoren. Ein Tensor ist einfach ein Array von Zahlen. Ihn auf eine GPU zu legen bedeutet, dass dieses Array im Speicher der GPU lebt, statt im normalen CPU-Speicher. Wenn a und b zwei Tensoren auf einer GPU sind, ist diese eine Zeile so schnell fertig, dass du nie darüber nachdenkst.
Jetzt mach es kleiner. Sagen wir, a und b sind zwei einzelne Floats, beide auf der GPU. Dieselbe Zeile. c = a + b. Was läuft tatsächlich auf dem Chip?
Die Antwort ist ein Kernel. In dieser Welt ist ein Kernel ein kleines Programm, das die GPU auf einem bestimmten Datenstück ausführt. Nicht der OS-Kernel, von dem dein Laptop bootet. Nicht der Mathe-Kernel aus einem Lineare-Algebra-Lehrbuch. Das Wort wird häufig wiederverwendet, und das ist nicht deine Schuld. In der GPU-Welt bedeutet ein Kernel einfach: eine kleine Funktion, die die GPU angewiesen bekommt, sie jetzt sofort, parallel, auf den Daten auszuführen, die du ihr gegeben hast.
Am Ende wirst du in der Lage sein, auf ein PyTorch-Code-Snippet zu schauen und zu zählen, wie viele Kerne die GPU ausführen wird. Das klingt nach einem kleinen Trick, und das ist es auch, aber es ist der erste Schritt, um das Gefühl der "GPU als Blackbox" loszuwerden – das Gefühl, wenn dein Modell langsam ist und du nicht weißt, warum. Jede Frage, die du zur GPU-Leistung stellen kannst, führt letztlich zurück zu: "Welche Kernel liefen und was haben sie gemacht?" Also fangen wir hier an.
Dein erster Kernel
Lass uns a und b etwas größer machen. Dieses Mal Tensoren der Länge 8. Immer noch eine PyTorch-Zeile: c = a + b.
Wenn du das ausführst, sagt deine CPU (die Maschine, die das eigentliche Python macht) der GPU: Hey, führe dieses Programm auf diesen Daten aus. Diese Anweisung nennt man einen Launch. Was gestartet wird, ist ein Kernel: ein Programm, bereit zur Ausführung. Launches selbst sind billig, Mikrosekunden pro Stück. Das, was um den Launch herum passiert (die Daten, die zur GPU gehen, die Ergebnisse, die zurückkommen), ist das, wo die wirklichen Kosten liegen, und das werden wir im Laufe der Zeit zählen.
Innerhalb des Kerns wird die eigentliche Arbeit von winzigen Arbeitern erledigt, den sogenannten Threads. Eine GPU hat Tausende davon verfügbar. Für unsere Addition der Länge 8 übernehmen 8 Threads die Arbeit: Thread 0 kümmert sich um Element 0, Thread 1 um Element 1 und so weiter bis Thread 7. Jeder Thread führt dasselbe winzige Programm aus: lies ein Element von a, lies das passende Element von b, addiere sie, schreibe das Ergebnis in c.
(In der Praxis startet die GPU Threads in Gruppen fester Größe, sogenannten Warps, bei NVIDIA-Karten immer 32 Threads, und maskiert die überzähligen, wenn dein Array nicht glatt aufgeht. Kannst du vorerst ignorieren.)
Wir haben also eine PyTorch-Zeile, einen Launch, einen Kernel, 8 Threads, die 8 Additionen durchführen. Jetzt zählen wir, was tatsächlich über den Chip gegangen ist. Um die Addition durchzuführen, braucht jeder Thread sein Element von a und sein Element von b. Das sind 8 Lesevorgänge von a und 8 Lesevorgänge von b. Dann schreibt jeder Thread sein Ergebnis in c. Das sind 8 Schreibvorgänge.
Diese Lese- und Schreibvorgänge gehen an den großen Speicher, der direkt neben dem GPU-Chip sitzt. Bei Rechenzentrumskarten (A100, H100) heißt dieser Speicher HBM (High-Bandwidth Memory). Bei Verbraucherkarten (RTX 4090) und den Colab-artigen T4s, die Leute wahrscheinlich tatsächlich ausprobieren werden, heißt er GDDR. Wie auch immer, es ist schneller Speicher, der neben dem Chip sitzt, und wir nennen ihn einfach GPU-Speicher. Er ist schnell, aber er ist nicht umsonst, und jeder Zugriff darauf kostet etwas.
Ein Kernel = ein Launch = ein Durchlauf über die Daten. Was auch immer der Kernel in seinem Inneren tut, die Lese- und Schreibvorgänge an seinen Rändern (die Reisen zum GPU-Speicher, um Eingaben zu holen, die Reise zurück, um Ausgaben zu schreiben) sind der Teil, der kostet. Das ist das ganze Prinzip.
Nichts davon ändert sich, wenn die Tensoren größer werden. Dieselbe PyTorch-Zeile, derselbe eine Kernel, nur mehr Threads. Wenn a und b eine Million Elemente haben, startet die GPU denselben Kernel mit einem größeren Schwarm von Threads. Die Mathematik skaliert, die bewegten Bytes skalieren, das mentale Modell nicht. Eine Zeile, ein Kernel.
Was zwischen zwei Operationen passiert
1c = (a + b).relu()
Du kennst Python gut genug, um zu wissen, dass dies zwei Operationen sind, eine Addition und dann ein ReLU, verkettet. In einem Interpreter sind das zwei Funktionsaufrufe. Auf einer GPU, im Eager-Modus von PyTorch, sind das zwei Kernel-Launches: einer für die Addition, einer für das ReLU. So weit, so wenig überraschend.
Was wirklich interessant ist, ist das, was zwischen den beiden Kernen passiert.
Wenn die Addition fertig ist, muss sie ihr Ergebnis irgendwo ablegen. Dieses Irgendwo ist der GPU-Speicher. Die Addition schreibt ein vollständiges Zwischenarray (nennen wir es tmp) in den Speicher. Einen Moment später startet das ReLU, und seine erste Aufgabe ist es, dasselbe tmp-Array wieder aus dem Speicher zu lesen. Es liest das Ganze, wendet ReLU auf jedes Element an und schreibt das Ergebnis in c.
Zähle den Speicherverkehr für diese beiden Kernel:
- Die Addition: liest
a, liestb, schreibttmp. Drei array-große Transfers. - Das ReLU: liest
tmp, schreibtc. Zwei weitere.
Fünf array-große Transfers insgesamt. Vergleiche das mit der Addition der Länge 8 allein aus dem vorherigen Abschnitt, die drei waren. Das Hinzufügen von .relu() zur Kette hat dich nicht nur die Rechenarbeit des ReLU gekostet. Es hat dich einen vollständigen Roundtrip des Arrays durch den GPU-Speicher gekostet, weil tmp geschrieben werden musste, nur damit der nächste Kernel es wieder einlesen konnte.
Nichts wurde zwischengespeichert. tmp durfte nicht in einem Register oder einem schnellen lokalen Cache bleiben. Es ging hinaus in den GPU-Speicher (der langsam-weiten Art) und kam gleich zurück. Die beiden Kernel sind einander fremd. Sie müssen Daten über das eine Medium übergeben, das beide zu sprechen wissen: den GPU-Speicher.
Warum macht PyTorch das so? Weil im Eager-Modus, wenn du a + b schreibst, PyTorch es sofort ausführt. Es weiß nicht, dass du als Nächstes .relu() aufrufen wolltest. Jede Operation wird in dem Moment ausgeführt, in dem ihre Python-Zeile erreicht wird. Es gibt keinen Plan, keine Vorausschau. Jede Operation steht für sich allein, produziert ein echtes Array und übergibt es über den Speicher an das, was als Nächstes kommt.
Dies ist das Muster, das du dir merken solltest. Jeder Zwischenwert in deinem PyTorch-Code wird physisch in den GPU-Speicher geschrieben und von der nächsten Operation wieder eingelesen. Jeder einzelne. Das ist es, was die "Kernel-Anzahl" wirklich misst. Jeder zusätzliche Kernel ist ein weiterer Roundtrip, den deine Daten durch den GPU-Speicher machen müssen.
Fusion: zwei Operationen, ein Kernel
Stell dir einen Kernel vor, der alles auf einmal erledigt: liest sein Element von a, liest sein Element von b, addiert sie, wendet ReLU auf das Ergebnis an (alles innerhalb des Kernels, auf winzigem, thread-eigenem Arbeitsspeicher, der den Chip nie verlässt), und schreibt erst dann den endgültigen Wert in c. Das Zwischenergebnis (a + b) existiert immer noch, aber nur innerhalb des Kernels, im privaten Arbeitsspeicher jedes Threads. Es wird nie in den GPU-Speicher geschrieben. tmp existiert als echtes Array überhaupt nicht.
Zähle jetzt die Transfers. Lesevorgänge von a: 1 pro Element. Lesevorgänge von b: 1 pro Element. Schreibvorgänge nach c: 1 pro Element. Drei array-große Transfers. Dieselbe Mathematik wie bei der Zwei-Kernel-Version, aber zwei Roundtrips weniger.
Bei Länge 8 ist das ein Rundungsfehler. Das interessiert niemanden. Bei Länge 1 Million oder 100 Millionen werden diese zusätzlichen Roundtrips zu einem großen Teil der Laufzeit, und die Wanduhr spiegelt das wider. Warum der Speicherverkehr letztlich so dominant wird, ist das gesamte Thema von Artikel 2 dieser Serie, also lasse ich das "Warum" hier außen vor. Der Punkt jetzt ist nur: gleiche Mathematik, weniger Fahrten, in der Praxis schneller.
Dieser Trick, Operationen, die separate Kernel gewesen wären, in einem einzigen Kernel zu kombinieren, sodass das Zwischenergebnis nie den GPU-Speicher besuchen muss, hat einen Namen. Er heißt Fusion. Das ist das ganze Wort. Das ist die ganze Idee.
Jetzt der unbequeme Teil. Das manuelle Schreiben dieses kombinierten Kernels sieht für Addition + ReLU einfach aus. Zwei Operationen. Eine Zeile "Rechenarbeit" dazwischen. Aber echter PyTorch-Code hat Dutzende von Operationen, die miteinander verkettet sind, jede mit ihren eigenen Formen, Datentypen und Broadcasting-Regeln. Einen fusionierten Kernel zu schreiben, der das alles korrekt handhabt, ist echte Ingenieursarbeit. Normalerweise würde man diese routinemäßigen elementweisen Kernel nicht von Hand schreiben.
Gute Nachrichten: PyTorch wird mit einem Tool ausgeliefert, das die Umschreibung für dich erledigt, genau für diese Art von Fall, automatisch. Es heißt torch.compile.
Du hast torch.compile fast sicher schon so verwendet:
1model = torch.compile(model)
Eine Zeile. Jemand im Internet hat dir gesagt, dass es Dinge schneller macht. Hier ist, was es tatsächlich tut, in einfachem Deutsch: Anstatt deine Operationen eine nach der anderen auszuführen, wie es der Eager-Modus tut, erfasst torch.compile die Tensor-Operationen, die deine Funktion ausführt, sucht nach Möglichkeiten, sie zu kombinieren, und generiert optimierten Code. Spätere Aufrufe, die dieselben Annahmen erfüllen, können diesen Code wiederverwenden.
Die Fusion, die wir oben auf dem Papier durchgeführt haben (Addition und ReLU teilen sich einen Kernel, tmp berührt nie den Speicher), ist genau die Art von Sache, die torch.compile automatisch mit deinem Code machen wird, solange die Operationen einfach genug sind. Wenn Leute sagen, dass torch.compile "PyTorch schneller macht", ist das ein großer Teil dessen, was sie meinen.
Für die Fälle, die torch.compile nicht von selbst fusionieren kann (benutzerdefinierte Operationen, die es nicht erkennt, ungewöhnliche Reduktionen, seltsame Speicherlayouts), muss trotzdem jemand einen Kernel von Hand schreiben. Dafür sind Tools wie Triton und CUDA da. Ein separater Artikel.
Überzeuge dich selbst
Alles oben war Zählen von Kernen auf dem Papier. Zeit, sie auf einer echten GPU zu zählen. Wenn du irgendeine Maschine mit einer CUDA-GPU zur Hand hast (einen Arbeitsplatzrechner, ein Colab-Notebook, eine Cloud-Instanz), kannst du das in ein paar Minuten selbst ausführen.
Das Tool ist torch.profiler. Es ist in PyTorch eingebaut. Es zeichnet lediglich auf, was die GPU während der Ausführung deines Codes tatsächlich getan hat, und gibt dir eine Tabelle zurück, die du lesen kannst.
Schritt 1: die Eager-Version.
Packe die Zwei-Operationen-Zeile in eine Funktion, damit wir etwas zum Aufrufen haben:
1import torch2from torch.profiler import profile, ProfilerActivity34def add_relu(a, b):5 return (a + b).relu()67a = torch.randn(1_000_000, device="cuda")8b = torch.randn(1_000_000, device="cuda")910with profile(activities=[ProfilerActivity.CUDA]) as prof:11 add_relu(a, b)12 torch.cuda.synchronize()1314print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
torch.cuda.synchronize() ist nur dazu da, um sicherzustellen, dass die GPU fertig ist, bevor wir die Zeitmessungen lesen. GPU-Arbeit läuft asynchron, und ohne die Synchronisation würdest du manchmal den Launch-Overhead anstelle der tatsächlichen Kernel-Arbeit messen.
Schritt 2: lies die Ausgabe.
Deine tatsächliche Profiler-Ausgabe wird wahrscheinlich mehr Zeilen haben, als du erwarten würdest. Eine Reihe von Speicherzuweisungs- und PyTorch-Buchhaltungszeilen werden eingemischt sein. Die Zeilen, die uns interessieren, sind die CUDA-Kernel, die eigentlichen Funktionen, die die GPU ausgeführt hat. Suche nach Zeilen mit kernel im Namen. Die beiden werden ungefähr so aussehen:
1vectorized_elementwise_kernel<...BinaryFunctor_add...> 1 12us2vectorized_elementwise_kernel<...threshold_kernel...> 1 9us
Die genauen Vorlagennamen ändern sich zwischen PyTorch-Versionen (ReLU erscheint oft als threshold, weil das die zugrunde liegende Operation ist, und Addition manchmal als CUDAFunctor_add). Versuche nicht, das Ganze zu parsen. Zähle einfach die Zeilen. Zwei Zeilen. Zwei Kernel. Einen für die Addition, einen für das ReLU. Genau das, von dem wir vor einem Abschnitt gesagt haben, dass es passieren würde.
Schritt 3: die kompilierte Version.
Eine Zeile Änderung. Wickle die Funktion in torch.compile ein:
1compiled = torch.compile(add_relu)
Hier ist eine kleine Falle, die es wert ist, zu wissen, bevor du es ausführst. Der erste Aufruf einer mit torch.compile kompilierten Funktion ist langsam, manchmal seltsam langsam, weil dann torch.compile seine Arbeit erledigt: deinen Code analysieren, herausfinden, was fusioniert werden kann, den fusionierten Kernel generieren. Wenn du den allerersten Aufruf profilierst, misst du den Kompilierungsschritt, nicht den Kernel. Das Muster ist also: rufe es einmal auf, um es aufzuwärmen, verwirf das Ergebnis, dann profiliere.
1compiled(a, b) # Aufwärmen, verwerfen2torch.cuda.synchronize()34with profile(activities=[ProfilerActivity.CUDA]) as prof:5 compiled(a, b)6 torch.cuda.synchronize()78print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
Schritt 4: lies die Ausgabe erneut.
Dieses Mal:
1triton_poi_fused_add_relu_0 1 14us
Eine Zeile. Ein Kernel. Der Name verrät dir sogar, was es getan hat: fusionierte Addition und ReLU. Gleiche Mathematik wie zuvor, ein Launch statt zwei.
Du hast gerade genau das getan, worüber dieser Artikel spricht, in einem Satz: du hast PyTorch gebeten, zwei Operationen zu einem Kernel zu kombinieren, hast den Profiler beobachtet und bestätigt, dass die Anzahl von zwei auf eins gesunken ist. Fusion, in freier Wildbahn, auf deiner Maschine.
Wenn du es dramatischer sehen willst, probiere es mit ein paar verschiedenen Tensor-Größen. Bei Länge 100 laufen beide Versionen so schnell, dass der Unterschied im Rauschen untergeht. Bei Länge 10 Millionen oder 100 Millionen beginnt die kompilierte Version klar die Nase vorn zu haben, weil der Roundtrip, den wir weggeschnitten haben, bei dieser Größenordnung ein echter Brocken Arbeit ist.
Das Zählen von Kernen ist jetzt keine abstrakte Empfehlung mehr. Du hast einen Weg, es zu überprüfen.
Fassen wir zusammen
Hier ist das Ganze auf einen Schlag.
Dein PyTorch-Code, wenn er auf einer GPU läuft, wird in eine Sequenz von Kernen umgewandelt. Jeder Kernel ist ein Launch, ein Durchlauf über deine Daten, ein Roundtrip durch den GPU-Speicher, um Eingaben zu holen und Ausgaben zu schreiben. Einfache Operationen werden zu einem Kernel. Verkettungen von Operationen werden standardmäßig zu einem Kernel pro Operation, wobei Zwischenergebnisse dazwischen Roundtrips durch den Speicher machen. torch.compile kann einfache Verkettungen für dich fusionieren, sodass diese Zwischenergebnisse nie den Speicher berühren. Weniger Kernel bedeuten normalerweise weniger Speicherverkehr. Und weniger Speicherverkehr bedeutet normalerweise schneller.
