FlashAttention-3: Schnellere und präzisere Attention durch Asynchronität und geringe Präzision

• Der Attention-Engpass von Transformern hat die Trainings- und Inferenzgeschwindigkeit von LLMs mit langem Kontext begrenzt, und FlashAttention-3 ist eine neue Version, die dies mithilfe von Hopper-GPU-Funktionen stärker reduzieren soll
• Der Kern besteht darin, die Asynchronität von Tensor Core und TMA zu nutzen, um Berechnung und Datenbewegung zu überlappen, und matmul und softmax verschränkt anzuordnen, um GPU-Leerlaufzeiten zu verringern
• Auf H100 erreichte FlashAttention-2 nur 35 % der theoretischen maximalen FLOPS, FlashAttention-3 kommt in FP16 jedoch auf bis zu 740 TFLOPS, also bis zu 75 % des H100-Theoriewerts
• Bei niedriger FP8-Präzision vergrößern Ausreißer in LLM-Aktivierungen die Fehler, daher senkt incoherent processing auf Basis der Hadamard-Transformation den Fehler gegenüber einer FP8-Attention-Basislinie um das 2,6-Fache
• FlashAttention-3 wurde auf GitHub veröffentlicht, und je stärker die Attention-Kosten sinken, desto größer wird der Spielraum für längere Kontexte und effizienteres LLM-Training sowie effizientere Inferenz

---

Ziele und Leistung von FlashAttention-3

• Attention ist die Kernschicht von Transformern, wird aber bei großen Sprachmodellen und Anwendungen mit langem Kontext zum zentralen Engpass
• FlashAttention und FlashAttention-2 haben einen Ansatz etabliert, der Attention durch geringere GPU-Speicher-Lese- und Schreibzugriffe beschleunigt; die meisten Bibliotheken nutzen ihn heute zur Beschleunigung von Transformer-Training und -Inferenz
• Dieser Ansatz hat dazu beigetragen, dass die Kontextlänge von LLMs in den letzten zwei Jahren von 2–4K bei GPT-3 und OPT auf 128K bei GPT-4 und 1M bei Llama 3 gestiegen ist
• FlashAttention-2 erreichte auf H100-GPUs nur 35 % der theoretischen maximalen FLOPS, doch FlashAttention-3 steigert dies durch neue Funktionen der Hopper-GPU
• Die veröffentlichten Leistungswerte von FlashAttention-3 sind wie folgt
  • In FP16 1,5- bis 2,0-mal schneller als FlashAttention-2
  • Bis zu 740 TFLOPS in FP16
  • 75 % Auslastung der theoretischen maximalen FLOPS von H100
  • Nahezu 1,2 PFLOPS in FP8
  • 2,6-mal geringerer Fehler gegenüber einer FP8-Attention-Basislinie

Rückblick auf den FlashAttention-Ansatz

• FlashAttention beschleunigt Attention, indem die Reihenfolge der Attention-Berechnung neu angeordnet wird und Tiling sowie Neuberechnung genutzt werden; zugleich sinkt der Speicherbedarf in Bezug auf die Sequenzlänge von quadratisch auf linear
• Eingabeblöcke werden aus HBM in SRAM geladen, Attention wird auf diesen Blöcken ausgeführt, und anschließend wird die Ausgabe in HBM aktualisiert
• Da die große zwischengelagerte Attention-Matrix nicht in HBM geschrieben wird, verringern sich Speicher-Lese- und Schreibzugriffe, was in der tatsächlichen Laufzeit 2- bis 4-fache Beschleunigung ermöglichen kann
• Wenn Tiling und Softmax-Rescaling zusammen eingesetzt werden, lässt sich blockweise verarbeiten und dennoch ohne Approximation die korrekte Ausgabe erhalten

Hopper-GPU-Funktionen: WGMMA, TMA, FP8

• FlashAttention-2 kann auf der Ampere-A100-GPU bis zu 70 % der theoretischen maximalen FLOPS erreichen, nutzt die neuen Funktionen der Hopper-GPU jedoch nicht ausreichend aus
• FlashAttention-3 verwendet drei Hopper-Funktionen
  • WGMMA: die Warpgroup-Matrix-Multiply-Accumulate-Funktion, die die neuen Tensor Cores von Hopper nutzt und einen höheren Durchsatz als mma.sync von Ampere bietet
  • TMA: eine dedizierte Hardwareeinheit zur Beschleunigung des Datentransfers zwischen global memory und shared memory, die Indexberechnung und out-of-bound predication übernimmt und so den Registerverbrauch senkt
  • FP8: kann den Tensor-Core-Durchsatz gegenüber FP16 verdoppeln, bringt aber wegen der geringeren Bitzahl zur Darstellung von Fließkommazahlen einen Genauigkeits-Trade-off mit sich
• FlashAttention-3 nutzt die Abstraktionen von NVIDIA CUTLASS, um Hopper-Funktionen einzusetzen
• Schon das reine Neuschreiben von FlashAttention zur Nutzung dieser neuen Funktionen verbessert die FP16-Forward-Pass-Leistung von etwa 350 TFLOPS bei FlashAttention-2 auf 540–570 TFLOPS

Überlappung von GEMM und Softmax durch Asynchronität

• Die Hauptoperationen von Attention sind GEMM zwischen Q-K und P-V sowie softmax
• Auf modernen Beschleunigern sind nicht-matmul-Operationen viel langsamer als matmul, und Spezialfunktionen wie die Exponentialfunktion in softmax werden auf anderen Einheiten verarbeitet als floating point multiply-add oder matrix multiply-add
• H100 SXM5 bietet bei FP16-Matrix-Multiplikation 989 TFLOPS, während der Durchsatz für Spezialfunktionen bei 3,9 TFLOPS liegt, also 256-mal niedriger
• Bei einer Head-Dimension von 128 kann die Exponentialfunktion trotz eines 512-fach höheren matmul-FLOPS-Werts zeitlich 50 % gegenüber matmul ausmachen
• In FP8 wird matmul doppelt so schnell, die Exponentialgeschwindigkeit bleibt aber gleich, weshalb die parallele Ausführung von matmul und softmax noch wichtiger wird

• Pingpong-Scheduling zwischen Warpgroups

  • Der Warp-Scheduler der GPU führt automatisch einen gewissen Overlap aus, indem er andere Warps ausführt, während einige Warps auf GEMM-Ergebnisse warten
  • FlashAttention-3 nutzt Synchronisationsbarrieren, um GEMM und softmax zweier Warpgroups manuell stärker zu überlappen
  • Warpgroup 1 führt zuerst GEMM1 einer Iteration und GEMM0 der nächsten Iteration aus
  • Danach verarbeitet Warpgroup 1 softmax, während Warpgroup 2 GEMM ausführt
  • Dieses Pingpong-Schedule versteckt softmax hinter der GEMM-Ausführungszeit der anderen Warpgroup
  • Das tatsächliche Scheduling ist nicht ganz so sauber wie im Diagramm, erhöht aber beim FP16-Attention-Forward-Pass mit Head-Dimension 128 und Sequenzlänge 8K den Durchsatz von etwa 570 TFLOPS auf 620 TFLOPS

• Overlap innerhalb einer Warpgroup

  • Auch innerhalb einer einzelnen Warpgroup kann ein Teil von softmax ausgeführt werden, während GEMM dieser Warpgroup läuft
  • Dieses Pipelining steigert den FP16-Attention-Forward-Durchsatz von etwa 620 TFLOPS auf 640–660 TFLOPS
  • Dafür müssen jedoch GEMM-Akkumulatoren und Softmax-Ein- und Ausgaben gemeinsam gehalten werden, was den Registerdruck erhöht
  • Insgesamt bietet diese Technik einen vorteilhaften Trade-off

Niedrige FP8-Präzision und incoherent processing

• LLM-Aktivierungen können Ausreißer enthalten, deren Magnitude deutlich größer ist als die der übrigen Features
• Ausreißer erschweren die Quantisierung und vergrößern den Quantisierungsfehler stark
• FlashAttention-3 nutzt incoherent processing, wie es in Quantisierungsarbeiten wie QuIP verwendet wird
• Dazu werden query und key mit einer zufälligen orthogonalen Matrix multipliziert, um Ausreißer zu verteilen und Quantisierungsfehler zu verringern
• In der Implementierung wird eine Hadamard-Transformation mit zufälligem Vorzeichen verwendet
  • Wenn die Head-Dimension d ist, kann dies pro Attention-Head in O(d log d) statt O(d²) ausgeführt werden
  • Die Hadamard-Transformation ist memory-bandwidth-bound; wird sie mit einer vorausgehenden Operation fusioniert, die ebenfalls memory-bandwidth-bound ist, etwa Rotary Embedding, lässt sie sich ohne Zusatzkosten ausführen
• In einem Experiment, in dem Q, K und V aus einer Standardnormalverteilung erzeugt und 0,1 % der Einträge mit großer Magnitude versehen wurden, um Ausreißer zu simulieren, verringerte incoherent processing den Quantisierungsfehler um das 2,6-Fache

Benchmarks und Veröffentlichungsstatus

• FlashAttention-3 wurde nicht nur mit FlashAttention-2 verglichen, sondern auch mit Triton- und cuDNN-Implementierungen, die bereits neue Hardwarefunktionen der Hopper-GPU nutzen
• In FP16 zeigt sich gegenüber FlashAttention-2 eine Beschleunigung von etwa 1,6- bis 1,8-mal
• In FP8 werden nahezu 1,2 PFLOPS erreicht
• Das FlashAttention-3-GitHub-Repository wurde veröffentlicht
• Das Paper ist ebenfalls im selben flash-attention-Repository verfügbar

Weitere Optimierungen und künftige Integration

• Das Paper enthält neben den im Blog behandelten Inhalten auch Optimierungen wie variable length sequence, persistent kernel und in-kernel transpose für FP8
• Wenn Algorithmen passend zur Ausführungshardware entworfen werden, können große Effizienzgewinne und neue Modellfunktionen wie längere Kontexte erschlossen werden
• Zu den künftigen Arbeiten gehören die Optimierung der LLM-Inferenz und die Verallgemeinerung der Techniken auf andere Hardwarearchitekturen
• Es wird erwartet, dass FlashAttention-3 in zukünftige PyTorch-Releases integriert wird

[GN⁺]

Hacker-News-Kommentare

• Wenn man sich die Code-Kommentare ansieht, wirkt es so, als hätte Tri Dao bereits im April 2022, direkt nach der Ankündigung von Hopper/H100, an FA3 gearbeitet.
  Es ist etwas interessant, dass es bis zur heutigen Veröffentlichung des Codes mehr als zwei Jahre gedauert hat; möglicherweise liegt das daran, dass noch bessere Lösungen in Vorbereitung sind.
  Tris jüngere Paper-Historie neigt sich eher zu SSM- und Mamba-artigen Architekturen. FlashAttention hat in Bezug auf die Sequenzlänge quadratische Zeitkomplexität, während neuere Algorithmen quasisubquadratisch sind, also nicht nur dieselbe Berechnung effizienter ausführen, sondern den Rechenaufwand selbst deutlich reduzieren.
  Dao und Gu zeigten dieses Jahr in einem langen Paper, dass sich Mamba/SSM ebenfalls so formulieren lassen, dass sie gut durch dieselben Hardware-Primitivoperationen beschleunigt werden können, von denen Transformer profitieren.

  • Bis die Strong Exponential Time Hypothesis (SETH) bewiesen oder widerlegt ist, braucht man quadratische Kosten oder muss auf etwas verzichten. Letztlich sind das die Kosten einer vollständigen Suche.
    Würde man SETH beweisen oder widerlegen, wäre damit auch das P-vs-NP-Problem gelöst, daher ist kaum zu erwarten, dass das bald geschieht.
    Entscheidend ist, ob ein bestimmter Anwendungsfall diese Kosten tragen kann.

• Ich frage mich, wie stark der FlashAttention-Algorithmus an die Hardware gebunden ist.
  In dieser Ankündigung heißt es zum Beispiel, dass die asynchronen Funktionen der H100-GPU genutzt werden; das scheint also zu bedeuten, dass man diesen Geschwindigkeitsgewinn auf Karten außerhalb der H-Serie nicht bekommt.
  Außerdem benötigt die eigentliche FlashAttention-Bibliothek CUDA, aber der Algorithmus scheint nach Metal portiert worden zu sein[^0]. Wenn der Algorithmus eher etwas wie eine reine Funktion ist, müsste er sich dann nicht auf jeder GPU bzw. in jedem Machine-Learning-Framework implementieren lassen?
  [0]: https://github.com/philipturner/metal-flash-attention

  • Es gibt viele gute Antworten, aber kurz gesagt ist er „in der Praxis ziemlich stark“ an Hardware gebunden. Das hier ist ein brauchbares Beispiel.

    https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
    Karpathys nanoGPT prüft, ob torch.nn.functional.scaled_dot_product_attention vorhanden ist, und ruft dann FlashAttention auf.
    https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
    Laut der Dokumentation möchte man in der Praxis meist FA2 aufrufen; FA2 optimiert die Kernel des Geräts, um die Softmax-Operation auf einer Dreiecksmatrix aufzuteilen und unnötige Rundreisen von Floating-Point-Batches zwischen GPU und CPU zu verringern.
    https://arxiv.org/pdf/2307.08691
    Das FA2-Paper ist fast vollständig aus der Perspektive der Hardware geschrieben, auf der es ausgeführt wird.

  • Die algorithmische Verbesserung von FlashAttention besteht hauptsächlich darin, den Softmax-Teil der Attention aufzuteilen und wieder zusammenzuführen; das ist an sich keine völlig neue Idee. Der überwältigende Beitrag liegt darin, diese Methode und ihre Details auf Nvidia-Hardware effizient umgesetzt zu haben.
  • Das ursprüngliche FlashAttention war kaum hardwareabhängig.
    Bei neueren Versionen hängt es vom Abstraktionsniveau ab. ThunderKittens[0] liefert, ähnlich wie im Beitrag beschrieben, gegenüber FA2 etwa einen 1,3- bis 2-fachen Geschwindigkeitsgewinn und ist dabei relativ allgemein auf GPUs anwendbar.
    Mit jeder neuen Hardware kann es hardwarespezifische Funktionen geben, die zusätzliche Leistung herausholen. Üblicherweise übernehmen Vendoren Funktionen, die anderen einen Vorsprung verschaffen, aber wie schon bei CUDA werden APIs und Bibliotheken dabei fragmentiert.
    [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
  • Konzeptionell ein wenig, aus Sicht der praktischen Implementierung aber stark gebunden. Selbst eine Standard-Python-Implementierung kompiliert intern Kernel, die auf bestimmte Hardware zugeschnitten sind.
  • Aus praktischer Sicht ergänzend: Für AMD-Hardware fehlt bislang noch eine wirklich gute Implementierung von flash-attention-2. ROCm wird langsam brauchbar, ist aber noch nicht auf dem Niveau von CUDA.

• Ich würde gern die Compiler-Leute fragen: Besteht die Möglichkeit, dass ein Compiler Optimierungen wie FlashAttention selbst findet? TVM und tinygrad scheinen in diese Richtung zu gehen, aber ich finde es schwer zu glauben, dass das möglich ist.

  • Theoretisch ist es möglich. Dank der algebraischen Eigenschaften der Mathematik ist eine weitreichende Umordnung möglich, und dazu kommt noch relativ standardmäßiges polyedrisches Loop-Tiling.
    Da die Kosten dafür jedoch hoch sind, muss man die Ergebnisse dieser Suche cachen.
    e-graph-Optimierung scheint für diesen Bereich gut geeignet zu sein. Dafür wäre allerdings ein großer Paradigmenwechsel bei der Verarbeitung von Optimierungs-Pässen nötig, weshalb es abgesehen von einigen Nischen-Tools kaum verbreitet wurde. Es passt zum Beispiel nicht gut zu traditionellen Aufrufgraphen; um e-graphs über und zwischen Basic Blocks und for-Schleifen hinweg einzusetzen, müsste man den Kontrollfluss stark verändern, und break sowie return werden ebenfalls nicht unterstützt.
  • Es wirkt wie ein extrem schwieriges, aber nicht unmögliches Problem.
    Ich weiß allerdings nicht genau, wo der aktuelle Stand der Compiler-Optimierung im Hinblick auf Datenlayout und die Maximierung der Prozessorauslastung liegt.
    Ich habe einmal einen Vortrag über Optimierung gesehen, in dem kleine Optimierungen zwar die Geschwindigkeit erhöhten, ihr Effekt aber gering war im Vergleich zu den Geschwindigkeitsschwankungen durch Unterschiede im Speicherlayout, die durch diese Optimierungen oder sogar durch zufällige Änderungen entstanden.
    Der Vortrag konzentrierte sich eher darauf, das Signal im Rauschen zu erkennen, aber dieses Rauschen selbst ist ein Hinweis darauf, dass Compiler hier selbst mit viel einfacheren Formen des angesprochenen Problems nicht besonders gut umgehen können.
    Schon die CPU-/Speicherarchitektur ist komplex, wenn es darum geht, wie Cache und Zugriffsmuster die Geschwindigkeit beeinflussen; nimmt man noch die GPU-Architektur dazu, ist das Gebiet wohl noch weitgehend unerforscht.
    Vielleicht wird es eines Tages möglich sein. Da es um KI geht, stellt sich auch die Frage, ob eine hinreichend kluge KI das erledigen könnte, aber das hängt davon ab, was „hinreichend“ bedeutet.
    Als Test auf sehr hohem Niveau für ein KI-Modell könnte man sich vorstellen, ihm etwas wie micrograd zu geben und zu verlangen, bei gleicher Schnittstelle etwas zu bauen, das schneller ist als torch. Davon sind wir noch nicht einmal annähernd dort, aber falls es möglich wird, wäre das interessant.
  • Ich denke nein. Man sollte das als einen anderen Algorithmus betrachten. Statt nur die Mathematik zu berücksichtigen, entwirft man den Algorithmus mit Blick auf die Form der Hardware.
    TVM erscheint mir nachvollziehbar. Genau genommen macht es etwas anderes, ist aber ziemlich nah dran.
    Bei tinygrad weiß ich allerdings nicht, warum das diesen Eindruck vermittelt hat.
  • https://github.com/uwplse/tensat
  • In einer umhüllenden Sprache wie Python High-Level-Operatoren aufzurufen, ist ziemlich knifflig.

• Falls jemand das auf ROCm / AMD MI300x portieren möchte, bitte bei [e-mail] melden. Ich sende garantiert keinen Spam.
  Ich kann Rechenzeit für diese Arbeit spenden.

  • Also eine Firma für AMD-Beschleuniger-Server! Tolle Arbeit, und ich hoffe, jemand nimmt das an :)
  • Nicht unhöflich gemeint, aber mich interessiert die Absicht hinter diesem Angebot. Bekommt jemand nur Hardware-Zugang und soll dann dieses Porting kostenlos machen? Was hat die Person davon?

• FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
  Wie ist die Performance von FA3 auf Consumer-GPUs wie 3090 und 4090?

  • Es ist nur für Hopper gedacht. Die Verbesserungen sind stark an Hopper-Features wie Warp Groups und TMA gebunden.
    Auf einer 4090 könnte man mit der Triton-Implementierung von FP8-Attention trotzdem einen Geschwindigkeitsschub bekommen: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html

• TMA (Tensor Memory Accelerator). This is a special hardware unit that accelerates the transfer of data between global memory and shared memory, taking care of all index calculation and out-of-bound predication. This frees up registers, which is a valuable resource to increase tile size and efficiency.
  So wie ich es verstehe, reduziert TMA zwar auch den Registerbedarf, wichtiger ist aber, dass die Adressgenerierung von der Hardware übernommen wird. Je schneller die umliegenden Operationen werden, desto eher kann die Adressgenerierung zum Flaschenhals werden.

• Das ist eine der wichtigsten Verbesserungen in der gesamten KI. Sie ermöglicht es, auf derselben Hardware mehr und schneller zu tun, und bringt den meisten KI-Nutzern fast ohne Kompromisse Vorteile.

  • Für Nutzer mit einer H100 stimmt das wohl.

• Ich frage mich, warum FlashAttention mit variabler Maskierung ungefähr 5-mal langsamer ist als ohne. Wenn gute Unterstützung für Maskierung fehlt, verschwindet der Optimierungseffekt fast vollständig.

  • Wo sehen Sie diesen Benchmark?

• Es wäre schön, wenn ein Experte ein paar Fragen beantworten könnte :)
  Ist FlashAttention einfach ein Drop-in-Ersatz für die Attention-Berechnung in LLMs? Kann man es überall verwenden, wo die „Attention“-Operation eingesetzt wird, oder muss ein LLM eigens darauf trainiert werden, FA zu verwenden?
  In welchem Verhältnis steht FA zu Strategien wie GQA oder Sliding-Window-Attention? Sind das zueinander orthogonale Konzepte, oder braucht jede Strategie eine eigene FA-Implementierung?
  Kürzlich hat llama.cpp Unterstützung für FlashAttention hinzugefügt — heißt das, dass dort nun etwas wie die von FlashAttention bereitgestellten CUDA-Kernels verwendet werden?
  Und zuletzt: In diesem Artikel werden FlashAttention und Triton verglichen. Ist Triton nicht eher so etwas wie eine Abstraktionsschicht? Kann man FA nicht in Triton implementieren? Die Formulierung „FlashAttention vs. Triton“ ist mir nicht ganz klar

  • 1. Fast richtig. Mathematisch ist es äquivalent. Auf Software-Seite gibt es nur Themen wie Abhängigkeits-Versionierung oder Datenformate im Speicher, und FlashAttention 2 ist bereits in HuggingFace und mehreren populären Bibliotheken enthalten. Es ist gut möglich, dass FlashAttention 3 bald ebenfalls aufgenommen wird, aber zum Ausführen wird eine H100-GPU benötigt.
    2. FlashAttention 2 hat in einem früheren Versionsupdate Unterstützung für GQA hinzugefügt:
       https://github.com/Dao-AILab/flash-attention
    3. Hier wird diese FlashAttention-Implementierung in reinem CUDA C++ mit der Triton-Implementierung eines ähnlichen Algorithmus verglichen, die in Triton geschrieben ist: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
  • FlashAttention kann die Attention-Berechnung in LLMs als Drop-in ersetzen.
    FlashAttention ist eine Methode zur Berechnung des Softmax(QK^T)V-Teils der Attention, während GQA eine Methode zur Berechnung der Q-, K- und V-Matrizen ist. Sliding-Window-Attention ist, soweit ich nicht ganz sicher bin, eine Änderung der Attention-Maske, die steuert, welche Queries welche Keys sehen können.
    Ich habe llama.cpp zwar nicht verwendet, aber die Beschreibung, dass dort nun CUDA-Kernels übernommen werden, scheint im Großen und Ganzen zuzutreffen.
    Zur letzten Frage: Gemeint ist eine frühere FlashAttention-Implementierung, die in Triton geschrieben wurde.

• Weil in diesem Artikel gesagt wird, dass Operationen wie sigmoid sehr langsam sind, frage ich mich Folgendes:
  Moderne LLMs verwenden oft Aktivierungsfunktionen wie SiLU, Swish und SoLU, in denen sigmoid oder Softmax enthalten sind.
  Hat ReLU geringere Performance-Einbußen? Falls ja, wäre es dann vielleicht besser, wieder zum guten alten ReLU zurückzukehren?

  • ReLU ist buchstäblich eine lineare Funktion, die an einem bestimmten Punkt bei 0 abgeschnitten wird, und daher rechnerisch viel günstiger als Funktionen, die Exponentialfunktionen enthalten. Mit einer so einfachen Aktivierungsfunktion dürfte es allerdings schwierig sein, konkurrenzfähige Ergebnisse zu erzielen.