CUDA-Stream-Interleaving steigert Token-Generierung in PyTorch
Autoregressive Decoder-Modelle wie GPT-2 sind zentral für moderne generative KI-Anwendungen und werden in nahezu allen Bereichen des täglichen Lebens eingesetzt. Ihre Token-Generierung erfolgt sequenziell, was zwar effizient erscheint, aber erhebliche Rechenlast mit sich bringt – insbesondere bei großen Modellen und hohen Sequenzlängen. Um Latenz und Kosten zu minimieren, werden kontinuierlich Optimierungen entwickelt, von spezialisierten CUDA-Kernen über CUDA Graphs bis hin zu spekulativer Sampling-Technik. In diesem Beitrag stellen wir eine effektive Methode vor, die Token-Generierung in PyTorch durch CUDA-Stream-Interleaving optimiert. Obwohl die Implementierung einfach ist, behebt sie eine oft übersehene Leistungsbremse: die Blockierung durch Host-Device-Synchronisation beim frühen Beenden von Sequenzen (Early Stopping). Wir nutzen ein einfaches GPT-2-Modell aus der HuggingFace-Transformers-Bibliothek (v5.1.0), laufen auf einer NVIDIA L40S-GPU mit PyTorch 2.10.0. Zunächst zeigen wir eine naive Implementierung, die bei Verdopplung der Sequenzlänge die Laufzeit etwa vervierfacht – ein klassisches O(N²)-Verhalten. Dies resultiert aus der wiederholten Berechnung des gesamten Eingabesequenzverlaufs bei jedem Token. Die Lösung: KV-Caching, das die Key- und Value-Tensoren aus früheren Schritten speichert und wiederverwendet. Dadurch sinkt die Komplexität von O(N²) auf O(N), was zu einer drastischen Geschwindigkeitssteigerung führt. Trotzdem bleibt die CUDA-Speicherfragmentierung hoch, da bei jedem Schritt leicht größere Speicherblöcke angefordert werden. Um dies zu beheben, testen wir zwei Ansätze: Erstens expandable memory segments (experimentell, über PYTORCH_ALLOC_CONF="expandable_segments:True"), was die Fragmentierung reduziert und eine geringfügige Performance-Verbesserung bringt. Zweitens StaticCache, ein festes Cachegröße-Modell, das den Speicherbedarf vorab festlegt. Dies senkt die Fragmentierung erheblich, führt aber zu unnötiger Berechnung von Maskierung für nicht genutzte Tokens – eine Leistungskosten, die mit steigender Sequenzlänge wächst. Eine weitere Optimierung ist die Modellkompilierung mit torch.compile, die besonders bei statischen Eingabegrößen (wie bei StaticCache) maximale Performance erzielt. Die Kombination aus StaticCache und torch.compile führt bereits zu einem signifikanten Geschwindigkeitsgewinn. Der entscheidende Schritt ist jedoch die Behandlung des Early-Stopping-Overheads. Die Prüfung, ob alle Sequenzen beendet sind (torch.all(finished).item()), verursacht eine blockierende Synchronisation zwischen CPU und GPU. Dies führt zu erheblichem GPU-Idle-Zeit – wie in Nsight™ Systems-Profilen sichtbar: zwischen den Schritten ruht die GPU etwa 110 Mikrosekunden, was in Hochleistungsworkloads katastrophal ist. Unsere Lösung: CUDA-Stream-Interleaving im Ping-Pong-Verfahren. Wir verwenden zwei CUDA-Streams, um die Generierung von Token i-1 und i zu pipelinen. Während der CPU die EOS-Prüfung für Schritt i-1 durchführt (in der Host-Thread), läuft bereits die Berechnung für Schritt i auf der GPU. Die Ergebnisse werden über nicht-blockierende Kopien in gepinnter Speicherzone übertragen. Dadurch wird die Synchronisation verdeckt, die GPU bleibt kontinuierlich beschäftigt – die Leistungsauslastung steigt stabil an. Die Ergebnisse zeigen: Bei einer Batch-Größe von 32 und Sequenzlänge 100 erreicht die Kombination aus StaticCache, torch.compile und Stream-Interleaving eine Geschwindigkeitssteigerung um fast das Fünffache gegenüber der Basisimplementierung. Der Nutzen ist besonders bei kleinen Batch-Größen maximal (bis zu 11,6 %), da dort die Kernel-Ladezeit relativ zur Berechnungszeit höher ist. Bewertung durch Experten: Diese Methode ist ein Paradebeispiel für feinabgestimmte Low-Level-Optimierung in PyTorch. Während vLLM oder TensorRT-LLM für Produktionsszenarien oft bessere Ergebnisse liefern, ist das Stream-Interleaving für Entwicklungs- und Testumgebungen äußerst wertvoll. Es zeigt, dass selbst kleine, scheinbar unbedeutende Synchronisationspunkte massive Leistungseinbußen verursachen können. Die Methode erfordert jedoch Sorgfalt: Falsche Synchronisation kann zu CUDA-Fehlern oder Datenkorruption führen. Zudem kann der Mehraufwand durch mehrere Streams zu erhöhter Speicherreservierung führen. Insgesamt ist die Technik eine starke Ergänzung zu bestehenden Optimierungen – besonders dort, wo keine dedizierte Inference-Engine zur Verfügung steht. Für Entwickler, die PyTorch-native Inferenz maximieren wollen, ist sie ein essenzieller Bestandteil des Optimierungswerkzeugs.
