FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao; Daniel Y. Fu; Stefano Ermon; Atri Rudra; Christopher Ré

Paper · arXiv 2205.14135

FlashAttention: Schnelle und speichereffiziente exakte Attention mit IO-Bewusstsein

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

Veröffentlicht 2022-05-27cs.LG arXiv PDF

Mit Übersetzung und KI lesen

Zusammenfassung in drei Zeilen

🎯 Standardmäßige exakte Attention materialisiert eine $O (N^{2})$ -Attention-Matrix im langsamen GPU-HBM, wodurch sie speichergebunden wird und sich nur schwer auf lange Kontexte skalieren lässt.
🔧 FlashAttention berechnet exakte Attention innerhalb eines einzigen gefusten CUDA-Kernels durch Tiling der Eingaben über schnelles On-Chip-SRAM, inkrementelle Softmax-Aktualisierungen und Recomputation der Attention auf dem Chip während des Backward-Passes.
📈 FlashAttention bringt bis zu 3-fache End-to-End-Trainingsbeschleunigung bei GPT-2, trainiert 15 % schneller als der MLPerf-1.1-Rekord für BERT-large und skaliert den Kontext auf bis zu 64K Token.

Abstract

Transformers sind bei langen Sequenzen langsam und speicherhungrig, da die Zeit- und Speicherkomplexität der Self-Attention quadratisch mit der Sequenzlänge wächst. Approximative Attention-Methoden haben versucht, dieses Problem anzugehen, indem sie Modellqualität für eine reduzierte Rechenkomplexität opferten, erzielen in der Praxis jedoch oft keinen Wall-Clock-Geschwindigkeitsvorteil. Wir argumentieren, dass ein fehlendes Prinzip darin besteht, Attention-Algorithmen IO-bewusst zu gestalten – also Lese- und Schreibzugriffe zwischen den Speicherebenen der GPU zu berücksichtigen. Wir stellen FlashAttention vor, einen IO-bewussten exakten Attention-Algorithmus, der Tiling nutzt, um die Anzahl der Speicherzugriffe zwischen dem GPU High Bandwidth Memory (HBM) und dem GPU On-Chip-SRAM zu reduzieren. Wir analysieren die IO-Komplexität von FlashAttention und zeigen, dass er weniger HBM-Zugriffe als Standard-Attention erfordert und für eine Spanne von SRAM-Größen optimal ist. Zudem erweitern wir FlashAttention auf Block-Sparse Attention, was einen approximativen Attention-Algorithmus liefert, der schneller ist als jede existierende approximative Attention-Methode. FlashAttention trainiert Transformers schneller als bestehende Baselines: 15 % End-to-End-Wall-Clock-Beschleunigung bei BERT-large (Sequenzlänge 512) im Vergleich zum MLPerf-1.1-Trainingsgeschwindigkeitsrekord, 3 $\times$ Beschleunigung bei GPT-2 (Sequenzlänge 1K) und 2,4 $\times$ Beschleunigung in der Long-Range Arena (Sequenzlänge 1K–4K). FlashAttention und Block-Sparse FlashAttention ermöglichen längere Kontexte in Transformers, was zu Modellen höherer Qualität führt (0,7 bessere Perplexität bei GPT-2 und 6,4 Punkte Zuwachs bei der Klassifikation langer Dokumente) sowie völlig neue Fähigkeiten eröffnet: die ersten Transformers, die bei den Aufgaben Path-X (Sequenzlänge 16K, 61,4 % Genauigkeit) und Path-256 (Sequenzlänge 64K, 63,1 % Genauigkeit) Ergebnisse oberhalb des Zufallsniveaus erzielen.

Original-Abstract (Englisch)

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3 $\times$ speedup on GPT-2 (seq. length 1K), and 2.4 $\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

Leseleitfaden

Problem und Motivation

Transformers weisen eine quadratische Zeit- und Speicherkomplexität bezüglich der Sequenzlänge $N$ auf. Obwohl approximative Attention-Methoden die theoretische FLOP-Anzahl verringern, führen sie oft nicht zu tatsächlichen Wall-Clock-Beschleunigungen, da Berechnungen auf modernen GPUs primär durch Datentransfers (IO) statt durch arithmetische Operationen limitiert werden.

Bei Standard-Attention werden bei der Berechnung von $S = Q K^{⊤}$ , $P = softmax (S)$ und $O = PV$ Zwischenmatrizen der Größe $N \times N$ in das langsame High Bandwidth Memory (HBM) geschrieben und mehrfach wieder ausgelesen. Dies verursacht $Θ (N d + N^{2})$ HBM-Zugriffe und einen Speicherbedarf von $O (N^{2})$ , was praktische Kontextfenster stark einschränkt.

Kernideen

IO-Bewusstsein (IO-Awareness): Explizite Berücksichtigung der GPU-Speicherhierarchie: schneller, kleiner On-Chip-SRAM (z. B. 192 KB/SM mit ~19 TB/s Bandbreite auf der A100) gegenüber großem, langsamem HBM (40–80 GB mit 1,5–2,0 TB/s Bandbreite).
Tiling (Blockweise Softmax-Berechnung): Aufteilung der Eingaben $Q, K, V$ in Blöcke, die vollständig in das SRAM passen. Durch das fortlaufende Nachführen von Zeilenmaxima und Normalisierungsstatistiken der Softmax wird die Softmax-Reduktion inkrementell über Blöcke hinweg berechnet, ohne die vollständige $N \times N$ -Matrix im HBM zu materialisieren.
Recomputation im Backward-Pass: Vermeidung des Abspeicherns der intermediären $N \times N$ -Attention-Matrix für die Backpropagation. Stattdessen werden nur die Ausgabe $O$ und die Softmax-Statistiken $(ℓ, m)$ gespeichert und $S$ sowie $P$ während des Backward-Passes blockweise direkt im SRAM auf dem Chip neu berechnet.
Kernel-Fusion: Fusion von Matrixmultiplikationen, Maskierung, Softmax und Dropout in einen einzigen CUDA-Kernel, um redundante HBM-Roundtrips zu eliminieren.

Methode (mit den zentralen Gleichungen)

Gegeben $Q, K, V \in R^{N \times d}$ , setzt FlashAttention die Blockgrößen auf $B_{c} = ⌈ M / (4 d)⌉$ und $B_{r} = min (⌈ M / (4 d)⌉, d)$ , wobei $M$ die SRAM-Kapazität darstellt.

Inkrementelles Softmax-Tracking

Für einen Vektor, der in Teilvektoren $x^{(1)}$ und $x^{(2)}$ zerlegt ist, aktualisieren sich die fortlaufenden Statistiken über: $m (x) = max (m (x^{(1)}), m (x^{(2)})),$ $ℓ (x) = e^{m (x^{(1)}) - m (x)} ℓ (x^{(1)}) + e^{m (x^{(2)}) - m (x)} ℓ (x^{(2)}) .$

Für Block $j$ der Keys/Values und Block $i$ der Queries:

Berechne unnormalisierte Scores im SRAM: $S_{ij} = Q_{i} K_{j}^{⊤} \in R^{B_{r} \times B_{c}}$ .
Berechne lokale Block-Statistiken: $\tilde{m}_{ij} = rowmax (S_{ij})$ , $\tilde{P}_{ij} = exp (S_{ij} - \tilde{m}_{ij})$ und $\tilde{ℓ}_{ij} = rowsum (\tilde{P}_{ij})$ .
Aktualisiere die fortlaufenden globalen Normalisierer: $m_{i}^{new} = max (m_{i}, \tilde{m}_{ij}), ℓ_{i}^{new} = e^{m_{i} - m_{i}^{new}} ℓ_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{ℓ}_{ij} .$
Aktualisiere den fortlaufenden Ausgabeblock $O_{i} \in R^{B_{r} \times d}$ : $O_{i} \leftarrow diag (ℓ_{i}^{new})^{- 1} (diag (ℓ_{i}) e^{m_{i} - m_{i}^{new}} O_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{P}_{ij} V_{j}) .$

Komplexitätsanalyse

IO-Komplexität: Standard-Attention benötigt $Θ (N d + N^{2})$ HBM-Zugriffe. FlashAttention benötigt $Θ (N^{2} d^{2} M^{- 1})$ HBM-Zugriffe. Die Autoren beweisen, dass kein exakter Attention-Algorithmus asymptotisch $o (N^{2} d^{2} M^{- 1})$ Zugriffe über alle $M \in [d, N d]$ hinweg erreichen kann.
FLOPs und Speicher: FlashAttention benötigt $O (N^{2} d)$ FLOPs (mit leicht mehr FLOPs im Backward-Pass durch die Recomputation) und erfordert neben Ein- und Ausgaben lediglich $O (N)$ Zusatzspeicher.

Erweiterung: Block-Sparse FlashAttention

Wird die Berechnung gemäß einer vordefinierten Sparsity-Maske $M \in {0, 1}^{N / B_{r} \times N / B_{c}}$ auf Nicht-Null-Blöcke beschränkt, ergibt sich eine IO-Komplexität von $Θ (N d + N^{2} d^{2} M^{- 1} s)$ , wobei $s$ der Anteil der Nicht-Null-Blöcke ist.

Ergebnisse (mit den wichtigsten Kennzahlen)

Trainingsbeschleunigung:
- BERT-large (Seq.-Länge 512): Übertrifft den MLPerf-1.1-Geschwindigkeitsrekord um 15 % (17,4 min gegenüber 20,0 min auf $8 \times$ A100s).
- GPT-2 (Seq.-Länge 1K): Erreicht bis zu 3,5-fache Beschleunigung gegenüber HuggingFace und bis zu 1,8-fache Beschleunigung gegenüber Megatron-LM.
- Long-Range Arena (LRA): Bringt eine 2,4-fache Beschleunigung gegenüber Standard-Attention; Block-Sparse FlashAttention erreicht 2,8-fache Beschleunigung.
Fähigkeiten bei längeren Kontexten:
- GPT-2 mit 4K Kontext: Trainiert schneller als Megatron-LM bei 1K Kontext und verbessert gleichzeitig die Validierungs-Perplexität um 0,7 Punkte (17,5 vs. 18,2).
- Klassifikation langer Dokumente: Das Skalieren des Kontexts von 512 auf 16K führt zu einem Zugewinn von 4,3 Micro- $F_{1}$ bei MIMIC-III (52,8 auf 57,1) und bis zu 8,5 Zuwachs bei ECtHR (72,2 auf 80,7 bei 8K).
- Path-X und Path-256: FlashAttention ist der erste Transformer, der Path-X (Seq.-Länge 16K) mit 61,4 % Genauigkeit löst. Block-Sparse FlashAttention erreicht 63,1 % Genauigkeit bei Path-256 (Seq.-Länge 64K).

Limitationen und offene Fragen

CUDA-Implementierungsaufwand: Das Implementieren IO-bewusster Kernel erfordert hardwarenahe CUDA-Programmierung statt Standard-PyTorch- oder Compiler-Abstraktionen, was die Portabilität über verschiedene Hardwarearchitekturen erschwert.
Multi-GPU Attention: Der optimale IO-Ablaufplan für einen einzelnen Knoten überträgt sich nicht automatisch auf verteilte Multi-GPU-Attention über Netzwerke (NVLink, InfiniBand).
Untere Schranke für exakte Attention: Obwohl der Ansatz über allgemeine $M$ optimal ist, bleibt der Beweis feingranularer, parametrisierter Komplexitätsuntergrenzen in Abhängigkeit von $M$ eine offene Frage.

Schlüsselbegriffe

IO-awareness: Das algorithmische Prinzip, Operationen unter expliziter Berücksichtigung der Lese- und Schreibzugriffe zwischen verschiedenen Ebenen des Hardwarespeichers (z. B. schnelles On-Chip-SRAM vs. langsameres HBM) zu entwerfen.
High Bandwidth Memory (HBM): Der große primäre GPU-Hauptspeicher, der eine hohe Gesamtkapazität, aber eine deutlich geringere Bandbreite als der On-Chip-Cache bietet.
SRAM: Schneller On-Chip-Scratchpad-Speicher auf modernen GPUs, der innerhalb jedes Streaming Multiprocessors (SM) angesiedelt ist.
Tiling: Restrukturierung einer Berechnung, indem Eingabematrizen in kleinere Blöcke aufgeteilt werden, die vollständig in das SRAM passen, und Zwischenergebnisse über mehrere Durchläufe hinweg akkumuliert werden.
Selective gradient checkpointing: Eine Technik, die das Speichern der großen intermediären Attention-Matrix während des Forward-Passes vermeidet, indem sie während des Backward-Passes mithilfe zwischengespeicherter Softmax-Statistiken direkt auf dem Chip neu berechnet wird.
Kernel fusion: Zusammenfassung mehrerer separater Operationen (Matrixmultiplikation, Maskierung, Softmax, Dropout) in eine einzige GPU-Kernel-Ausführung, um wiederholte Lese- und Schreibzugriffe auf das HBM zu eliminieren.
Block-sparse FlashAttention: Eine Erweiterung von FlashAttention, die nur Nicht-Null-Kachelblöcke anhand einer vorgegebenen Block-Sparsity-Maske berechnet und so die IO-Komplexität proportional zum Besetzungsgrad senkt.
Path-X / Path-256: Aufgaben zur visuellen Inferenz über lange Sequenzen aus der Long-Range Arena (16K und 64K Sequenzlänge), bei denen bisherige Transformers nicht über Zufallsniveau hinauskamen.

Fragen

Approximiert FlashAttention die Attention-Matrix?

Nein. FlashAttention berechnet exakte Self-Attention mit mathematisch identischen Ausgaben wie die Standard-Attention, bis auf numerische Fließkommagenauigkeit.

Warum ist FlashAttention schneller, obwohl es im Backward-Pass mehr FLOPs berechnet?

Weil Standard-Attention speichergebunden ist: Das Hin- und Hertransferieren der $N \times N$ -Matrix zwischen HBM und SRAM kostet weit mehr reale Laufzeit (Wall-Clock Time), als zusätzliche arithmetische Operationen im schnellen SRAM durchzuführen.

Welchen Zusatzspeicher legt FlashAttention für den Backward-Pass ab?

FlashAttention speichert lediglich die Forward-Ausgabe $O$ , die Softmax-Normalisierungsfaktoren $ℓ$ und $m$ (was nur $O (N)$ Zusatzspeicher benötigt) sowie den Zustand des Pseudozufallszahlengenerators zur Rekonstruktion der Dropout-Masken auf dem Chip.

Wie unterscheidet sich FlashAttention in der IO-Komplexität von Standard-Attention?

Standard-Attention erfordert $Θ (N d + N^{2})$ HBM-Zugriffe, wohingegen FlashAttention $Θ (N^{2} d^{2} M^{- 1})$ Zugriffe benötigt, wobei $M$ die Größe des SRAM und $d$ die Head-Dimension ist.

Die Übersicht auf dieser Seite wurde von einer KI aus dem Paper erstellt und kann Fehler enthalten. Details stehen im Paper.

Mit Übersetzung und KI lesen