FlashAttention

Attention vergleicht viele Paare von Textfragmenten. Wenn alle Zwischenergebnisse in den großen Speicher der Rechenkarte geschrieben werden, kann die Datenübertragung teuer werden. Du möchtest dieselbe Berechnung ausführen und dabei die Übertragung und Speicherung einer riesigen Tabelle begrenzen.

FlashAttention organisiert die Attention-Berechnung in Blöcken im kleinen, schnellen Speicher des Chips. Es schreibt nicht die gesamte Matrix der Vergleiche und Gewichte in den großen GPU-Speicher, also den Speicher der für parallele Berechnungen verwendeten Karte.

Jede Portion trägt einen Teil zum Ergebnis bei, und der Algorithmus aktualisiert die Normalisierung und die Summe korrekt. Man kann nicht einfach getrennte Mittelwerte der Portionen berechnen und diese mitteln: Die Portionen können unterschiedliche Anteile haben.

Die mathematische Operation der exakten Attention bleibt erhalten, mit den üblichen Unterschieden der Computerarithmetik. Bei dichter Attention wächst die Zahl der Vergleiche weiterhin quadratisch. Der Gewinn entsteht durch Speicherorganisation und Ausführung, nicht durch das Entfernen von Abhängigkeiten zwischen Paaren.

Quelle des Mechanismus: Dao et al., §2–3.2, Algorithmus 1.

Mechanismus und Details

Bei dichter Attention wächst die Anzahl der Operationen weiterhin quadratisch mit der Sequenzlänge. Der Gewinn stammt aus der Organisation von Berechnung und Speicher. Beim Training werden manche Größen im Rückwärtslauf neu berechnet statt gespeichert. „Exakt“ bedeutet, dass die Operation selbst nicht angenähert wird; eine andere Reihenfolge von Gleitkommaoperationen kann kleine numerische Unterschiede ergeben.

Blockergebnisse einfach zu mitteln reicht nicht

Softmax normalisiert die Gewichte über alle erlaubten Keys. Blöcke können sehr unterschiedliche Summen unnormalisierter Gewichte haben. Jeder muss daher den richtigen Anteil am Ergebnis erhalten.

Dies lässt sich schrittweise berechnen, indem man das Scoremaximum $m$, die Summe $l=\sum_j e^{s_j-m}$ und den Zähler $u=\sum_j e^{s_j-m}v_j$ führt. Das Ergebnis ist $u/l$. Erhöht ein neuer Block das Maximum auf $m'$, müssen die bisherigen $l$ und $u$ vor dem Addieren neuer Terme mit $e^{m-m'}$ skaliert werden. Die stabile Normalisierung erklären Milakov und Gimelshein, §3, Algorithmus 3.

Das eigene Beispiel verwendet vorgegebene Scores $[0,0,\ln 3,\ln 3]$ und Values $[0,0,1,1]$. Die Gewichte vor der Normalisierung haben das Verhältnis $[1,1,3,3]$. Das Ergebnis ist daher $6/8=0{,}75$. Die Ergebnisse der beiden Blöcke sind 0 und 1. Ihr gewöhnlicher Mittelwert liefert die falsche 0,5.

Die Demonstration isoliert die Aggregation für ein Query und skalare Values. Der vollständige Kernel berechnet Scoreblöcke aus Q und K auf der GPU; hier ermöglichen vorgegebene Zahlen die Prüfung von Hand. Der Wechsel der Blockgröße ist kein Benchmark.

Grouped-Query Attention begrenzt die Anzahl der K/V-Köpfe, FlashAttention die Art von Datentransfer und Verarbeitung. Beide Techniken lassen sich kombinieren. FlashAttention 2 ist eine spätere Implementierungsweiterentwicklung; der hier beschriebene Mechanismus stammt aus der ersten Arbeit.

Ich verwende KI-generierte Inhalte als Teil meines täglichen Lernprozesses.