Schnelle Attention-Kernel und Operator-Fusion reorganisieren die Berechnung von Transformatoren, um Datenbewegungen und Kernel-Starts auf dem Beschleuniger zu begrenzen. Diese Optimierungen lassen grundsätzlich die Gewichte des Modells und seine mathematische Funktion unverändert, auch wenn die Reihenfolge der Operationen und die Präzision kleine numerische Unterschiede erzeugen können. Das von Machine Learning Mastery veröffentlichte Tutorial erläutert diese Techniken mit Codebeispielen.
Eine Transformer-Schicht basiert auf gängigen Operationen: lineare Projektionen, Attention, Normalisierung, Aktivierungsfunktionen und Residualadditionen. Eine direkte Implementierung startet einen separaten Rechenkern, genannt « Kernel », für jede Operation und schreibt jeden Zwischentensor in den Speicher des Beschleunigers. Die Berechnung bleibt korrekt, aber die Laufzeit kann mehr für das Verschieben von Werten und das Starten von Kernen aufgewendet werden als für die eigentlichen Berechnungen.
Operator-Fusion
Der SwiGLU-Block, der von vielen reinen Decoder-Modellen verwendet wird, verdeutlicht dieses Problem. Eine direkte Implementierung startet Matrixmultiplikationen für die Projektionen, einen Kernel für die SiLU-Funktion, einen weiteren für die elementweise Multiplikation und dann eine letzte Matrixmultiplikation. An jeder Grenze kann ein Zwischentensor in den Hochbandbreitenspeicher (HBM) geschrieben und dann wieder gelesen werden. Bei der Fusion werden diese Schritte in einem einzigen Kernel kombiniert. Sie spart Zeit, wenn Speicherbewegungen und Kernel-Starts stark ins Gewicht fallen, und deutlich weniger, wenn bereits eine große Matrixmultiplikation die Ausführung dominiert. Die Wahl des Kernels bleibt empirisch: Die gesamte Arbeitslast muss profiliert werden.
Andere Fusionen sind in Transformatoren üblich, wie das Hinzufügen einer Residualverbindung gefolgt von einer RMSNorm- oder LayerNorm-Normalisierung oder die Kombination von Dequantisierung mit einer Matrixmultiplikation. Die Fusion hat jedoch ihre Grenzen: Operationen mit inkompatiblen Datenorganisationen oder ein sehr großer Zwischenzustand lassen sich schlecht kombinieren. Ein zu großer Kernel kann den Druck auf die Register erhöhen, die Auslastungsrate verringern und langsamer werden als mehrere gut optimierte Kernel.
FlashAttention und Page-Attention
Eine klassische Attention berechnet eine Score-Matrix der Größe T² pro Kopf, wobei T die Sequenzlänge ist. Auf einer GPU verbraucht das Erzeugen, Schreiben in den globalen Speicher und erneute Lesen viel Bandbreite, was bei langen Prefill-Phasen zum Engpass werden kann. FlashAttention verarbeitet Queries, Keys und Values blockweise. Für jeden Block berechnet es einen Teil der Scores, aktualisiert die Softmax-Statistiken und akkumuliert die entsprechende Ausgabe im On-Chip-Speicher, ohne die vollständige Matrix zu speichern. PyTorch macht diese Optimierung über die Funktion scaled_dot_product_attention() verfügbar. Das Tutorial weist darauf hin, dass das Erzwingen eines festen Backends den Code scheitern lassen kann, wenn kein Backend die Eingaben oder die Hardware unterstützt.
Der KV-Cache, der bereits berechnete Keys und Values speichert, wächst Token für Token, bei einer unvorhersehbaren Endlänge. Eine zusammenhängende Zuweisung pro Anfrage lässt daher reservierten und ungenutzten Platz. Der Page-Cache unterteilt den Speicher in Blöcke fester Größe und stützt sich auf eine Blocktabelle, die logische Positionen mit physischen Blöcken verknüpft. Der Attention-Kernel muss dann jede Position über diese Tabelle übersetzen. Die Paging-Methode reduziert nicht die Anzahl der logischen Elemente, begrenzt aber die Fragmentierung und die reservierte Kapazität.
Im bereitgestellten Beispiel verwendet vLLM das Modell Qwen/Qwen2.5-1.5B-Instruct mit Blöcken von 16 Tokens (block_size=16) und dem Backend « TRITON_ATTN ». Laut Tutorial prüft vLLM die Kompatibilität des Backends mit Modell und Hardware bei der Initialisierung und meldet einen Fehler bei Inkompatibilität.
Fusion und CUDA Graphs mit torch.compile
Die Funktion torch.compile() erfasst mit TorchDynamo eine Folge von Operationen und kompiliert sie dann mit TorchInductor, seinem Standard-Backend. TorchInductor kann einen einzigen Triton-Kernel für kompatible Operationen erzeugen, die der Standardmodus separat starten würde. Graph-Unterbrechungen, verursacht etwa durch datenabhängige Kontrollflussoperationen oder die Konvertierung von Tensoren in Python-Skalare, begrenzen die Bereiche, die der Compiler optimieren kann.
Der Modus « reduce-overhead » weist TorchInductor an, CUDA Graphs zu verwenden, wenn der erfasste Graph dafür geeignet ist. Ein CUDA Graph zeichnet die Folge der GPU-Operationen auf und spielt sie mit einer einzigen Einreichung an die CPU ab. Er erzeugt keine zusätzliche Fusion, sondern ergänzt die von TorchInductor. Er erfordert stabile Speicheradressen und kompatible Formen und kann mehr Speicher verbrauchen. Das Tutorial empfiehlt, die Ergebnisse mit toleranzangepassten Toleranzen für den Datentyp zu überprüfen, die Ausführung zu profilieren und die End-to-End-Leistung zu messen.
