Wie reduziert Flash Attention den Speicherverkehr während der lokalen Inferenz?

Flash Attention reduziert den Speicherverkehr, indem es die exakte Attention in Kacheln aufteilt, sodass Zwischenwerte direkt auf dem Chip verarbeitet werden können, anstatt wiederholt in den GPU-HBM geschrieben und daraus gelesen zu werden.

Auf einem lokalen KI-Server wird dieser Effekt mit zunehmender Prompt-Länge wichtiger, da Standard-Attention immer größere temporäre Arbeitsdaten erzeugt. Flash Attention verändert weder das Wissen des Modells noch entfernt es seinen KV-Cache. Stattdessen ordnet es die Attention-Berechnung so um, dass weniger Daten durch die langsameren Ebenen der GPU-Speicherhierarchie bewegt werden.

Standard-Attention bewegt große temporäre Tensoren durch den HBM

Attention beginnt mit Query-, Key- und Value-Tensoren. Anschließend müssen daraus Scores berechnet, normalisiert und mit den Values kombiniert werden. Eine einfache Implementierung legt zwischen diesen Schritten große temporäre Matrizen im High-Bandwidth Memory ab.

Das ursprüngliche FlashAttention-Paper argumentiert, dass die Leistung von Attention nicht nur durch Rechenoperationen, sondern auch durch Lese- und Schreibvorgänge zwischen HBM und On-Chip-SRAM begrenzt wird. Lange Sequenzen verstärken diesen Datenverkehr, weil die Score-Matrix mit der Anzahl der Tokenpaare wächst.

Eine Heim-GPU kann daher erheblich viel Zeit mit dem Verschieben von Daten verbringen, statt nützliche Matrixberechnungen auszuführen – insbesondere bei zunehmender Kontextlänge.

Flash Attention teilt das Problem in On-Chip-Arbeitssätze auf

Flash Attention verarbeitet Blöcke aus Queries, Keys und Values, die in den schnellen On-Chip-Speicher passen. Jede Kachel führt mehrere Attention-Schritte aus, bevor temporäre Daten den Chip verlassen.

FlashAttention-3 beschreibt den zentralen Ansatz dieser Familie als eine Kachelstrategie, die temporären Datenverkehr im globalen Speicher vermeidet. Der Algorithmus erzeugt für die unterstützte Formulierung weiterhin exakte Attention.

Die Kachelgröße hängt von den Hardware-Ressourcen und der Kernel-Implementierung ab. Der entscheidende Mechanismus ist die Datenlokalität: Geladene Values werden wiederverwendet, solange sie sich in der Nähe der Recheneinheiten befinden.

Online-Softmax vermeidet das Speichern der vollständigen Score-Matrix

Bei einer gewöhnlichen Softmax-Funktion scheint es erforderlich zu sein, zunächst alle Scores einer Zeile zu sammeln, bevor die Normalisierung erfolgt. Dadurch müsste die Implementierung eine große Matrix materialisieren. Flash Attention verwaltet stattdessen laufende Normalisierungsstatistiken über mehrere Kacheln hinweg.

Das Triton-Tutorial zu fusionierter Attention zeigt fusionierte Attention mit blockweisem Softmax-Zustand. Dadurch können partielle Score-Blöcke verarbeitet und verworfen werden, ohne jeden temporären Wert zurück in den HBM zu schreiben.

Hier treffen sich die Reduzierung des Speicherverkehrs und die Verringerung des temporären Speicherbedarfs: Der Algorithmus ordnet die Berechnung so um, dass ein großes temporäres Objekt nie im globalen Speicher existieren muss.

Featured

ZimaCube 2 Persönliche Cloud Home NAS

ZimaCube2

$799.00 - $2,499.00

Buy Now

Kernel-Fusion hält mehr Schritte der Attention-Kette in einem Ausführungspfad

Score-Berechnung, Skalierung, Maskierung, Normalisierung und Value-Aggregation können in einem koordinierten Kernel fusioniert werden, anstatt separate Kernel zu starten, die vollständige Zwischenwerte über den globalen Speicher austauschen.

PyTorchs Schnittstelle für Scaled-Dot-Product-Attention kann fusionierte Attention-Implementierungen auswählen, sofern Gerät, Datentyp, Form und angeforderte Funktionen unterstützt werden.

Das Ergebnis kann eine schnellere Prompt-Verarbeitung und einen geringeren Spitzenbedarf an temporärem Speicher sein. Der genaue Gewinn hängt jedoch davon ab, ob die Laufzeit tatsächlich einen Flash-kompatiblen Kernel verwendet.

Weniger IO bedeutet nicht, dass der KV-Cache entfällt

Bei der autoregressiven Inferenz müssen zuvor erzeugte Keys und Values weiterhin für spätere Token verfügbar bleiben, sofern das Modell oder die Laufzeit keine separate Strategie zur Entfernung, Komprimierung, Fensterung oder Auslagerung verwendet.

Die ROCm-Dokumentation von AMD beschreibt Flash Attention als Verbesserung der Speicherlokalität zwischen On-Chip-Speicher und HBM. Diese Optimierung betrifft die Art und Weise, wie Attention auf aktive Zustände zugreift und sie verarbeitet – nicht die Frage, ob historische KV-Zustände existieren.

Diese Unterscheidung verhindert einen häufigen Planungsfehler: Flash Attention kann den Arbeitsspeicherbedarf und den IO-Aufwand reduzieren, während ein Server mit langem Kontext trotzdem an seine Kapazitätsgrenze stößt, weil der vorgehaltene KV-Cache weiter wächst.

Hardware- und Formunterstützung entscheiden über die Verwendung des schnellen Pfads

Flash-Kernel sind spezialisiert. Head-Dimension, Datentyp, Attention-Maske, GPU-Architektur, Laufzeitversion und weitere Funktionen können ein Fallback auf eine andere Implementierung erzwingen.

Der Artikel von ZimaSpace zu Attention-Speicher im Verhältnis zu Modellparametern behandelt die Kapazitätsseite. Flash Attention sollte separat als Optimierung des Ausführungspfads gemessen werden.

Vergleiche kurze und lange Prompts, bestätige das ausgewählte Backend und erfasse Prefill-Latenz, maximal belegten Speicher, Speicherbandbreite und die Rate ausgegebener Token. Der Leitfaden von ZimaSpace zu lokaler KI-Hardware liefert die umfassendere Kapazitätsgrundlage, die durch Kernel-Optimierung nicht ersetzt werden kann.

Die Transformer Engine von NVIDIA behandelt optimierte Attention-Backends und hardwareabhängige Ausführung als Teil eines umfassenderen Stacks zur Kernel-Auswahl.