FlashAttention (DE)

FlashAttention ist ein revolutionärer Algorithmus zur Berechnung des Aufmerksamkeitsmechanismus (Attention), der entwickelt wurde, um das Training und die Inferenz von großen Sprachmodellen (LLMs) erheblich zu beschleunigen, während die volle Rechengenauigkeit erhalten bleibt. Der Algorithmus wurde erstmals 2022 von einem Forschungsteam der Stanford University unter der Leitung von Tri Dao vorgestellt.

Die Kernidee von FlashAttention besteht darin, die Berechnungen unter Berücksichtigung der GPU-Speicherhierarchie neu zu organisieren. Dies ermöglicht es, die Anzahl der Zugriffe auf den langsamen Speicher zu minimieren und den Hauptengpass des standardmäßigen Aufmerksamkeitsmechanismus zu beseitigen.

Problematik der Standard-Attention

Der standardmäßige Selbst-Aufmerksamkeitsmechanismus in Transformern wird nach folgender Formel berechnet: $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ wobei Q, K, V die Matrizen für Queries, Keys und Values sind.

Das Hauptproblem dieses Ansatzes ist die quadratische Komplexität in Bezug auf Zeit und Speicher (O(N²)) relativ zur Sequenzlänge N. Bei einer naiven Implementierung muss die vollständige Aufmerksamkeitsmatrix S der Größe N×N berechnet und im GPU-Speicher gehalten werden, was zu zwei kritischen Problemen führt:

1. Hoher Speicherverbrauch: Die Speicherung einer N×N-Matrix wird bei der Arbeit mit langen Kontexten unpraktikabel.
2. Ein-/Ausgabeoperationen (IO): Der Hauptengpass ist nicht die Anzahl der arithmetischen Operationen, sondern die ständigen Zugriffe auf den langsamen GPU-Speicher.

GPU-Speicherhierarchie

Um das Problem zu verstehen, ist es wichtig, zwei Arten von Speicher in einer GPU zu unterscheiden (am Beispiel der NVIDIA A100):

• SRAM (statischer Speicher): Schneller On-Chip-Speicher mit geringer Kapazität (~20 MB) und enormer Bandbreite (bis zu 19 TB/s).
• HBM (High-Bandwidth Memory): Langsamerer Speicher mit großer Kapazität (40–80 GB) und deutlich geringerer Bandbreite (ca. 1,5 TB/s).

Diese Asymmetrie macht den Standard-Attention-Algorithmus speicherbandbreitenlimitiert (memory-bound), da er ständig große Matrizen aus dem langsamen HBM liest und schreibt, was die Hauptursache für die Latenz ist.

Kerninnovationen von FlashAttention

FlashAttention ist ein IO-bewusster (IO-aware) Algorithmus, der das Problem durch die Minimierung von HBM-Zugriffen löst. Dies wird durch drei Haupttechniken erreicht.

Tiling und blockweise Verarbeitung

Anstatt die gesamte Matrix auf einmal zu verarbeiten, teilt FlashAttention die Eingabematrizen Q, K und V in kleine Blöcke (Tiles), die in den schnellen SRAM passen. Der Algorithmus lädt diese Blöcke nacheinander, führt alle Attention-Berechnungen für sie durch und aktualisiert das Endergebnis, ohne die vollständige Aufmerksamkeitsmatrix im langsamen HBM zu speichern.

Online-Softmax-Berechnung

Ein zentraler technischer Durchbruch war die „Online“-Berechnung von Softmax. Die Standard-Softmax-Funktion erfordert Kenntnis aller Elemente des Eingabevektors zur Normalisierung. FlashAttention verwendet einen modifizierten Algorithmus, der es ermöglicht, Softmax stückweise zu berechnen. Er verwaltet zwei Zwischenwerte (das aktuelle Maximum und die Summe der Exponentialfunktionen), die bei der Verarbeitung neuer Blöcke aktualisiert werden. Dies ermöglicht ein exaktes Ergebnis ohne den Zugriff auf die gesamte Matrix auf einmal.

Zusammenfassen von Operationen in einem CUDA-Kernel

Alle Attention-Operationen (Matrixmultiplikation QKᵀ, Maskierung, Softmax, Multiplikation mit V) sind in einem einzigen fusionierten CUDA-Kernel (fused kernel) zusammengefasst. Dies reduziert die Anzahl der Lese-/Schreibvorgänge zum HBM drastisch: Anstatt mehrmals über die gesamte Matrix zu iterieren, lädt der Algorithmus einen Block einmal in den SRAM, führt alle Berechnungen durch und schreibt nur das Endergebnis zurück.

Theoretische und praktische Effizienz

Komplexität und Optimalität

FlashAttention reduziert den Speicherverbrauch von O(N²) auf O(N), was eine lineare Skalierung ermöglicht. Es wurde nachgewiesen, dass die IO-Komplexität des Algorithmus für die Berechnung von Attention in einer zweistufigen Speicherhierarchie theoretisch optimal ist. Das bedeutet, eine exakte Attention-Berechnung kann ohne Hardwareänderungen nicht schneller durchgeführt werden.

Empirische Ergebnisse

Die erste Version von FlashAttention zeigte signifikante Verbesserungen:

• Beschleunigung:
  • BERT-large (Sequenzlänge 512): 15 % Beschleunigung beim Training.
  • GPT-2 (Sequenzlänge 1K): 3-fache Beschleunigung.
  • Aufgaben der Long-Range Arena (1K-4K): 2,4-fache Beschleunigung.
• Speichereinsparung: Bis zu 20-fache Speichereinsparung im Vergleich zu exakten Basisimplementierungen.
• Verbesserung der Modellqualität: Durch die Fähigkeit, mit längeren Kontexten zu arbeiten, verschlechtert FlashAttention die Modellqualität nicht, sondern verbessert sie sogar. Beispielsweise verbesserte sich die Perplexität von GPT-2 um 0,7 Punkte, und die Genauigkeit bei Klassifikationsaufgaben für lange Dokumente stieg um 6,4 Punkte.

Evolution und Weiterentwicklungen

Der Erfolg von FlashAttention legte den Grundstein für eine ganze Reihe von hardware-orientierten Algorithmen.

FlashAttention-2 (2023)

Die zweite Version zielte darauf ab, die GPU-Ressourcen besser auszunutzen. Im ursprünglichen FlashAttention betrug die Effizienz auf einer NVIDIA A100 nur 25–40 % des Maximums. FlashAttention-2 führte Verbesserungen bei der Parallelisierung der Berechnungen ein, was Folgendes ermöglichte:

• Eine zweifache Beschleunigung im Vergleich zur ersten Version zu erreichen.
• Die GPU-Auslastung auf 50–73 % des theoretischen Maximums zu steigern.
• Die Unterstützung auf Attention-Heads der Größe 256 sowie auf Architekturen wie Multi-Query Attention (MQA) zu erweitern.

FlashAttention-3 (2024)

Die dritte Version wurde speziell für die GPU-Architektur NVIDIA Hopper (H100) optimiert. Sie nutzt neue Hardwarefähigkeiten wie die Asynchronität der Tensor Cores und die Unterstützung für FP8, was ermöglichte:

• Eine weitere 1,5- bis 2-fache Beschleunigung im Vergleich zu FlashAttention-2 zu erzielen.
• Eine Leistung von bis zu 740 TFLOPS bei FP16 und nahezu 1,2 PFLOPS bei FP8 zu erreichen.

Spezialisierte Lösungen

Die Ideen von FlashAttention wurden in anderen Projekten weiterentwickelt:

• FlashInfer (2025): Eine anpassbare Attention-Engine, die speziell für LLM-Inferenzaufgaben optimiert ist. Sie konzentriert sich auf die effiziente Verarbeitung des KV-Caches im Streaming-Generierungsmodus.
• FlashMLA (2024): Eine Implementierung von Attention mit komprimiertem Kontext-Cache (latent attention), die es ermöglicht, bei sehr langen Sequenzen Speicher zu sparen bei minimalem Informationsverlust.

Einfluss auf die Industrie und das Ökosystem

FlashAttention wurde zu einem fundamentalen Durchbruch und entwickelte sich schnell zum Industriestandard für effizientes Training und Inferenz von LLMs. Er wurde in wichtige Bibliotheken wie PyTorch und Hugging Face integriert und wird in den meisten großen Sprachmodellen (LLaMA, MPT, Falcon, Claude usw.) verwendet.

Gerade FlashAttention und seine Nachfolgeversionen spielten eine entscheidende Rolle bei der Vergrößerung der Kontextfenster von Sprachmodellen: von 2.000–4.000 Token (GPT-3) auf 128.000 Token (GPT-4) und sogar bis zu Millionen von Token in experimentellen Modellen. Der Algorithmus beseitigte eines der Haupthindernisse für die Skalierung von Transformern und eröffnete neue Möglichkeiten für KI-Anwendungen, von der Analyse langer Dokumente bis hin zum multimodalen Verständnis.

Weblinks

• Offizielles FlashAttention-Repository auf GitHub

Literatur

• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. arXiv:2407.08608.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
• Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.