Entwurf von hardwarebewussten Algorithmen: FlashAttention
Das „T“ in chatGPT steht für Transformer, die Schlüsselarchitektur moderner KI-Entwicklung. Ursprünglich entwickelten Forscher den Transformer für maschinelle Übersetzung. Diese Architektur führte die Selbstaufmerksamkeit ein, wie im Papier Attention is all you need beschrieben. Über Schichten miteinander verbundener Knoten entsteht eine mathematische Darstellung, die Beziehungen und Relevanzen identifiziert, um eine Eingabesequenz in eine Ausgabesequenz zu transformieren.
Wenn Aufmerksamkeit alles ist, was Sie brauchen, machen wir es besser…
Die Transformer-Architektur hat eine neue Ära der KI-Forschung eingeleitet. Diese Forschung steigert die Effizienz des Kernmechanismus, der Aufmerksamkeit. Da die Zeit- und Speicherkomplexität der Aufmerksamkeit quadratisch (O(n²)) mit der Sequenzlänge (n) skaliert, erschwert sie die Modellierung langer Sequenzen. Lange Texte, Codebasen oder hochauflösende Bilder erfordern jedoch, dass Algorithmen Abhängigkeiten über große Distanzen erfassen. Um diese Herausforderung zu lösen, entwickelten Forscher hardwarebewusste und speichereffiziente Algorithmen wie FlashAttention.
Einleitung
Dieser Artikel zeigt die Konzepte, die FlashAttention (2022) erfolgreich gemacht haben. Die Techniken der zweiten (2023) und dritten (2024) Iteration behandeln wir in späteren Beiträgen.
Voraussetzungen
Folgende Kenntnisse erleichtern das Verständnis der Themen dieses Artikels:
- Der Transformer und der Aufmerksamkeitsmechanismus
- Matrizenmultiplikation
- Softmax-Operation
- Vorwärts-/Rückwärtsausbreitung (Forward-/Backward-Pass)
- Die GPU-Speicherhierarchie
- GPU-Leistungsoptimierung
- CUDA-Programmierungskonzepte (Thread-Blocks, Warps, Kerne)
- Gleitpunktformate (FP16, BF16, FP8)
Entwurf von hardwarebewussten und speichereffizienten Algorithmen
Moderne GPUs wie Hopper und Ampere bieten eine enorme Anzahl an Gleitpunktoperationen pro Sekunde (FLOPS), die ihre theoretische Rechenleistung darstellen. Gleichzeitig begrenzt jedoch die Speicherbandbreite die Geschwindigkeit, mit der Daten zwischen dem GPU-Speicher und den Verarbeitungseinheiten übertragen werden. Forscher entwerfen hardwarebewusste und speichereffiziente Algorithmen, indem sie die Speicherhierarchie optimal ausnutzen und möglichst viele der theoretischen FLOPS verwenden.
FlashAttention ist ein Paradebeispiel für einen Algorithmus, der längere Kontexte in Transformern ermöglicht. Er optimiert den Aufmerksamkeitsmechanismus für die zugrunde liegende Hardware.
FlashAttention (2022)
FlashAttention reduziert die Anzahl der Lese-/Schreibvorgänge zwischen GPU-Hochgeschwindigkeitsspeicher (HBM) und GPU-On-Chip-SRAM durch Tiling. Die Forscher beschreiben ihn als einen „IO-bewussten genauen Aufmerksamkeitsalgorithmus“.
GPU-Speicher: HBM & SRAM
Begriffe rund um GPU-Speichertypen variieren, obwohl sie oft ähnliche Konzepte beschreiben. FlashAttention nutzt zwei Speichertypen: HBM und SRAM.
| Speicher | AKA | Hauptmerkmale |
|---|---|---|
| HBM (High Bandwidth Memory) | GPU-Speicher, globaler Speicher | Langsam, große Speicherkapazität |
| SRAM (Static Random-Access Memory) | L1-Cache, Shared Memory | Schnell, kleine Speicherkapazität, On-Chip |
GPU-Computemodell
Aufmerksamkeitsberechnung
Die Aufmerksamkeitsaufstellung
Eine Auffrischung der Variablen zur Berechnung der Selbstaufmerksamkeitsschicht im Transformer:
- Abfrage (Q): Der Abfragevektor ist die aktuelle Eingabe oder das Element, für das die Aufmerksamkeit berechnet wird. Er gehört zu einer Matrix der Größe Nxd, wobei N die Sequenzlänge und d die Kopfdimension darstellt.
- Schlüssel (K): Die Schlüsselmatrix hat dieselben Dimensionen wie die Abfragematrix. Die Multiplikation von Abfrage- und Schlüsselvektoren ergibt Ähnlichkeitswerte.
- Ähnlichkeitswert (S): S misst, wie ähnlich die Abfrage jedem Element in der Sequenz ist. Die Multiplikation der Abfragematrix mit der transponierten Schlüsselmatrix liefert eine NxN-Matrix der Ähnlichkeitswerte.
- Aufmerksamkeitswahrscheinlichkeit (P): Die Softmax-Operation normalisiert S und erzeugt Wahrscheinlichkeiten, die sich zu 1 summieren.
- Wert (V): Die Wertvektoren enthalten Informationen über jedes Element und multiplizieren sich mit den Wahrscheinlichkeiten, um die Ausgabe zu erzeugen.
FlashAttention optimiert die Aufmerksamkeit
FlashAttention reduziert redundante Lese- und Schreiboperationen. Durch Techniken wie Tiling und Kernfusion bleibt der Speicherbedarf niedrig, ohne die Genauigkeit zu beeinträchtigen. Zusätzlich optimiert ein Online-Trick die Softmax-Operation.
Schlussfolgerung
FlashAttention zeigt eindrucksvoll, wie hardwarebewusste Algorithmen die Effizienz in KI-Systemen steigern. Optimierungen bei der Speicherhierarchie und Nutzung von FLOPS machen den Aufmerksamkeitsmechanismus schneller und senken den Speicherbedarf.
