FlashAttention: exakte Attention-Kernel für CUDA und ROCm
Schnelle und speichereffiziente exakte Aufmerksamkeit. FlashAttention Dieses Repository stellt die offizielle Implementierung von FlashAttention und FlashAttention-2 aus den folgenden Artikeln bereit.
Auf einen Blick
- Was ist das?
- Die offizielle Implementierung von FlashAttention und FlashAttention-2 samt Betaversionen für Hopper- und Blackwell-GPUs, mit Unterstützung für Ampere, Ada, H100 und AMD MI-Serien.
- Für wen ist es gedacht?
- FlashAttention richtet sich an ML-Teams, die lange Sequenzen trainieren oder Inferenz mit KV-Cache fahren und dabei Speicher sparen wollen; die Anforderungen sind konkret zu prüfen, denn Linux, PyTorch 2.2 oder neuer und CUDA 12 beziehungsweise ROCm 6 sind Voraussetzung, und ohne korrekt installiertes ninja zieht sich die Kompilierung laut README über Stunden.
- Darf ich es kommerziell nutzen?
- Ja. BSD-3-Clause ist eine freizügige Lizenz: Sie dürfen darauf aufbauende Software nutzen, verändern und verkaufen, solange Sie die Urheberrechts- und Lizenzhinweise beibehalten.
- Wird es noch gepflegt?
- Ja. Die letzten Commits kamen vor 1 Tag.
- In welcher Sprache ist es geschrieben?
- Hauptsächlich Python, laut der Sprachstatistik von GitHub.
Die Antworten beruhen auf den GitHub-Daten des Projekts (zuletzt abgeglichen am 15. September 2026) und auf unserer Analyse. Sie sind keine Rechtsberatung.
TIEFGEHENDE OPEN-SOURCE-ANALYSE
FlashAttention: exakte Aufmerksamkeit mit IO-Bewusstsein
Das Repository Dao-AILab/flash-attention ist die offizielle Implementierung von FlashAttention und FlashAttention-2 aus den im README zitierten Papieren. Das erste Papier stammt von Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra und Christopher Ré und beschreibt schnelle, speichereffiziente und exakte Aufmerksamkeit mit IO-Bewusstsein. FlashAttention-2 von Tri Dao setzt auf bessere Parallelität und Arbeitsaufteilung.
Daneben gibt es Beta-Stände für neuere Hardware: FlashAttention-3 ist für Hopper-GPUs wie die H100 optimiert und unterstützt FP16- und BF16-Durchläufe in beide Richtungen sowie FP8 vorwärts. FlashAttention-4 ist in CuTeDSL geschrieben und zielt auf Hopper und Blackwell, etwa H100 und B200. Die Metadaten weisen rund 24.800 Sterne und über 3.000 Forks aus, bei 1.275 offenen Issues; die Lizenz ist BSD-3-Clause.
Installation, ninja und die MAX_JOBS-Bremse
Die Anforderungen listet das README klar: CUDA- oder ROCm-Toolkit, PyTorch 2.2 oder neuer, die Python-Pakete packaging, psutil und ninja, sowie Linux. Windows funktioniert möglicherweise ab v2.3.2, wird laut README aber noch nicht als getestet geführt. Die Standardinstallation läuft über pip install flash-attn --no-build-isolation oder python setup.py install aus dem Quellcode.
Das Werkzeug ninja verdient besondere Aufmerksamkeit: Ohne es kann die Kompilierung laut README zwei Stunden dauern, weil sie keine CPU-Kerne parallel nutzt; mit ninja dauert sie auf einer 64-Kern-Maschine drei bis fünf Minuten. Das README empfiehlt, mit ninja --version und dem Rückgabecode die korrekte Installation zu prüfen und ninja notfalls neu zu installieren. Auf Maschinen mit weniger als 96 GB RAM und vielen Kernen kann die Umgebungsvariable MAX_JOBS die parallelen Kompilierjobs begrenzen, etwa mit MAX_JOBS=4 vor dem pip-Aufruf.
GPU-Abdeckung: Ampere, Hopper und AMD-Backends
Auf NVIDIA-Seite unterstützt FlashAttention-2 mit CUDA die GPU-Familien Ampere, Ada und Hopper, also etwa A100, RTX 3090, RTX 4090 und H100, mit den Datentypen fp16 und bf16 und Kopfdimensionen bis 256. Ab flash-attn 2.5.5 läuft die Rückwärtsberechnung für Kopfdimension 256 auch auf Consumer-GPUs, sofern kein Dropout verwendet wird. Turing-GPUs wie T4 und RTX 2080 sind hier nicht enthalten; das README verweist dafür auf das separate Repository flash-attention-turing.
Für AMD stehen unter ROCm 6.0 oder neuer zwei Backends bereit. Das composable_kernel-Backend ist der Standard und deckt MI200x, MI250x, MI300x, MI355x sowie RDNA 3/4 ab. Das Triton-Backend, dessen Kernel aus dem aiter-Paket als Git-Submodul unter third_party/aiter stammen, unterstützt CDNA- und RDNA-GPUs mit fp16, bf16 und fp32 und bietet MQA/GQA, Rotary-Embeddings, ALiBi, paged attention und FP8 über die FlashAttention-v3-Schnittstelle. Sliding-Window-Attention ist dort als laufende Arbeit markiert. Die Installation aktiviert man mit FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE".
Die Schnittstellen von flash_attn_func bis KV-Cache
Die Python-Schnittstelle stellt flash_attn_func, flash_attn_qkvpacked_func und flash_attn_with_kvcache bereit. Alle berechnen skalierte Punktprodukt-Aufmerksamkeit in der Form softmax(Q @ K^T * softmax_scale) @ V. Die Funktionen nehmen Parameter wie dropout_p, softmax_scale, causal, window_size, alibi_slopes und deterministic entgegen.
Die qkvpacked-Variante erwartet Q, K und V in einem gestapelten Tensor und vermeidet im Rückwärtsdurchlauf das explizite Verketten der Gradienten. Multi-Query- und Grouped-Query-Attention laufen, indem die Funktionen weniger KV-Köpfe als Q-Köpfe akzeptieren, sofern die Q-Kopfanzahl teilbar ist. Für die Inferenz aktualisiert flash_attn_with_kvcache die Caches k_cache und v_cache direkt im Speicher, wendet bei gesetzten rotary_cos und rotary_sin Rotationseinbettungen an und unterstützt keinen Rückwärtsdurchlauf. Ein Multi-Head-Attention-Modul liegt in flash_attn/modules/mha.py, dazu kommen laut README eine vollständige GPT-Modellimplementierung und Trainingsskripte.
Leistungs- und Speicherangaben im README
Das README enthält Benchmark-Diagramme für A100 und H100, die FlashAttention mit der Standard-Attention von PyTorch vergleichen. Die dokumentierte A100-Konfiguration nutzt Kopfdimensionen 64 und 128, verborgene Dimension 2048, Sequenzlängen von 512 bis 16k und eine Batchgröße von 16k geteilt durch die Sequenzlänge. Die Diagramme selbst sind Bilder; Zahlen im Text beziehen sich auf den Speicher.
Dort ist von zehnfacher Speicherersparnis bei Sequenzlänge 2k und zwanzigfacher bei 4k die Rede, weil der Speicherbedarf von FlashAttention linear mit der Sequenzlänge wächst, der der Standardattention aber quadratisch. Beim Training berichtet das README von einer vollständigen GPT-Implementierung, die gegenüber der Huggingface-Baseline um das Drei- bis Fünffache schneller sein soll und bis zu 225 TFLOPs pro Sekunde auf einer A100 erreicht, was 72 Prozent Modell-FLOPs-Auslastung entspricht, ohne Aktivierungs-Checkpointing. Diese Angaben stammen aus dem README selbst und sind nicht unabhängig verifiziert.
Von 2.0 bis 2.7: die Änderungshistorie im Überblick
Das Änderungsprotokoll im README dokumentiert die wichtigsten Schritte. Version 2.0 benannte die unpadded-Funktionen in varlen-Funktionen um und wurde als Neufassung beschrieben, die doppelt so schnell sein soll. Version 2.1 richtete die kausale Maske bei unterschiedlichen Query- und Key-Längen unten rechts aus. Version 2.2 optimierte die Inferenz für kurze Query-Sequenzen und führte flash_attn_with_kvcache ein.
Version 2.3 brachte Sliding-Window-Attention, Version 2.4 ergänzte ALiBi und einen deterministischen Rückwärtsdurchlauf, Version 2.5 folgte mit paged KV-Cache. Version 2.6 fügte Softcapping hinzu, wie es in Gemma-2 und Grok verwendet wird, und Version 2.7 brachte Kompatibilität mit torch.compile. Für Upgrader aus der 1.x-Serie sind diese Einträge die verlässliche Karte der Verhaltensänderungen.
Tests, BSD-3-Clause und Zitierpraxis
Die Tests prüfen laut README, ob FlashAttention dieselben Ausgaben und Gradienten wie eine Referenzimplementierung innerhalb einer numerischen Toleranz erzeugt, wobei der maximale Fehler höchstens doppelt so groß wie der der Baseline sein darf. Zu den Befehlen gehören pytest -q -s tests/test_flash_attn.py für CUDA und pytest tests/test_flash_attn_ck.py für das composable_kernel-Backend; die AMD-Triton-Suite unter tests/test_flash_attn_triton_amd.py dauert laut README mehrere Stunden.
Die Lizenz ist BSD-3-Clause, erlaubt also Weiterverbreitung und Änderung unter Bedingungen, und enthält einen Haftungsausschluss. Das README bittet ausdrücklich darum, die FlashAttention-Papiere zu zitieren, wenn der Code verwendet wird. Einen formellen Support-Prozess nennt es nicht; Fragen laufen über GitHub-Issues. Die neuesten Releases tragen Tags wie fa4-v4.0.0.beta28 vom 26. August 2026 und folgen einem kurzen Beta-Rhythmus.
Redaktionelles Fazit
FlashAttention richtet sich an ML-Teams, die lange Sequenzen trainieren oder Inferenz mit KV-Cache fahren und dabei Speicher sparen wollen; die Anforderungen sind konkret zu prüfen, denn Linux, PyTorch 2.2 oder neuer und CUDA 12 beziehungsweise ROCm 6 sind Voraussetzung, und ohne korrekt installiertes ninja zieht sich die Kompilierung laut README über Stunden.
Community-Notizen