FlashAttention-3 (PL)

FlashAttention-3 — to algorytm optymalizacji mechanizmu uwagi (attention) w sieciach neuronowych opartych na transformerach, opracowany w celu maksymalnego wykorzystania możliwości sprzętowych GPU architektury NVIDIA Hopper (H100). Algorytm został zaprezentowany w 2024 roku przez grupę badaczy z firm Colfax Research, Meta, NVIDIA, Georgia Tech, Uniwersytetu Princeton oraz Together AI. Praca została przyjęta na konferencję NeurIPS 2024 i wyróżniona jako spotlight.

FlashAttention-3 jest trzecią iteracją w rodzinie algorytmów, następującą po FlashAttention (2022) i FlashAttention-2 (2023). Jego głównym celem jest znaczące przyspieszenie trenowania i inferencji dużych modeli językowych (LLM) przy jednoczesnym zachowaniu dokładności obliczeń.

Wprowadzenie i historia

Problem mechanizmu uwagi

Kluczowym komponentem transformerów jest mechanizm samouwagi (self-attention), jednak jego złożoność obliczeniowa i zapotrzebowanie na pamięć rosną kwadratowo (O(n²)) wraz z wydłużaniem sekwencji wejściowej (n). Stanowi to poważne „wąskie gardło", ponieważ współczesne GPU są zoptymalizowane pod kątem szybkich mnożeń macierzowych, natomiast obliczanie funkcji wykładniczych (np. w Softmax) przebiega wielokrotnie wolniej. Ponadto w naiwnej implementacji pamięć GPU musi przechowywać duży pośredni tensor uwagi, co ogranicza skalowalność modeli.

FlashAttention i FlashAttention-2

Aby rozwiązać ten problem, w 2022 roku zaproponowano FlashAttention, który ograniczył liczbę odwołań do wolnej pamięci globalnej (HBM) dzięki dwóm technikom:

• Przetwarzanie blokowe (tiling): Obliczenia są dzielone na bloki (kafelki), przetwarzane w szybkiej pamięci on-chip (SRAM).
• Łączenie operacji: Wszystkie operacje (mnożenie macierzowe, Softmax) są wykonywane w jednym jądrze GPU bez zapisywania wyników pośrednich do pamięci globalnej.

Pozwoliło to obniżyć złożoność pamięciową z kwadratowej do liniowej i przyspieszyć obliczenia 2–4-krotnie.

W 2023 roku przedstawiono ulepszoną wersję — FlashAttention-2, która zoptymalizowała zrównoleglenie obliczeń. Na GPU architektury NVIDIA Ampere (A100) osiągnęła ona ~70% szczytowej teoretycznej wydajności. Jednak na nowszej architekturze NVIDIA Hopper (H100) jej efektywność okazała się znacznie niższa — około 35%. Wynikało to z tego, że algorytm nie wykorzystywał nowych możliwości sprzętowych Hopper, co stało się impulsem do stworzenia FlashAttention-3.

Nowe możliwości sprzętowe GPU Hopper (H100)

Architektura NVIDIA Hopper udostępniła szereg nowych funkcji, które FlashAttention-3 wykorzystuje do osiągnięcia maksymalnej wydajności:

• WGMMA (Warpgroup Matrix Multiply-Accumulate): Nowy typ instrukcji dla rdzeni tensorowych, wykonujący mnożenia macierzowe z prawie dwukrotnym przyrostem wydajności w porównaniu z architekturą Ampere.
• TMA (Tensor Memory Accelerator): Sprzętowy moduł przyspieszający transfer danych między pamięcią globalną (HBM) a pamięcią współdzieloną (shared memory). TMA automatycznie wykonuje obliczenia adresowe, odciążając rdzenie obliczeniowe.
• Format FP8: Sprzętowa obsługa 8-bitowego formatu zmiennoprzecinkowego, który podwaja teoretyczną wydajność w porównaniu z FP16, lecz niesie ryzyko utraty dokładności z powodu ograniczonego zakresu dynamicznego.

Innowacje techniczne FlashAttention-3

Algorytm implementuje trzy kluczowe metody optymalizacji, opracowane specjalnie dla architektury Hopper:

1. Asynchroniczne wykonywanie i specjalizacja warpów

FlashAttention-3 wykorzystuje zasadę warp-specialization, w której różne grupy wątków (warps) na GPU specjalizują się w różnych zadaniach:

• Producer warps: Ładują dane z pamięci globalnej za pomocą TMA.
• Consumer warps: Wykonują mnożenia macierzowe na rdzeniach tensorowych.

Dzięki sprzętowej asynchroniczności Hopper operacje te nakładają się w czasie. Podczas gdy jedna grupa warpów wykonuje obliczenia, druga równolegle ładuje dane dla następnego bloku. To potokowe podejście (pipeline), zorganizowane na zasadzie „ping-ponga" (ping-pong scheduling), pozwala ukryć opóźnienia wynikające z wolnych operacji (np. Softmax) i maksymalnie obciążyć wszystkie moduły funkcjonalne GPU.

2. Minimalizacja operacji pamięciowych

Algorytm zachowuje ideologię tiling z poprzednich wersji, ale aktywnie wykorzystuje TMA do asynchronicznego ładowania kolejnych bloków danych równolegle z bieżącymi obliczeniami. Transfer danych z wolnej HBM do szybkiej SRAM odbywa się faktycznie „w cieniu" głównych obliczeń, dzięki czemu GPU rzadziej czeka bezczynnie na dane.

3. Niska precyzja (FP8) z redukcją błędu kwantyzacji

Przejście na FP8 podwaja szybkość, ale może prowadzić do znacznej utraty dokładności z powodu kwantyzacji. Aby temu zaradzić, twórcy wprowadzili metodę incoherent processing. Jej istota polega na następującym:

1. Przed obliczeniem uwagi wektory cech (zapytania Q i klucze K) są mnożone przez losową macierz ortogonalną (np. macierz Hadamarda).
2. To przekształcenie „rozmywa" wartości o anomalnie dużej wartości bezwzględnej (wartości odstające) po wszystkich współrzędnych, wyrównując ich rozkład.
3. Następnie wykonywana jest kwantyzacja do FP8, która przebiega teraz z mniejszym błędem.
4. Ponieważ przekształcenie jest ortogonalne, nie zniekształca końcowego wyniku uwagi (QKᵀ), gdyż efekt macierzy znosi się przy mnożeniu.

Technika ta pozwoliła zredukować błąd obliczania uwagi w FP8 około 2,6-krotnie w porównaniu ze standardowym zastosowaniem FP8 bez przekształceń.

Wydajność i znaczenie

Zastosowanie wymienionych technik pozwoliło FlashAttention-3 osiągnąć znaczącą przewagę nad poprzednimi wersjami na GPU H100:

• Przyspieszenie 1,5–2-krotne w porównaniu z FlashAttention-2.
• Wysoka utylizacja GPU: Osiąga ~75–85% teoretycznego maksimum wydajności H100.
• Przepustowość:
  • Do 740–840 TFLOPS dla połowicznej precyzji (FP16/BF16).
  • Do 1,2–1,3 PFLOPS (petaflopsów) przy użyciu 8-bitowej precyzji (FP8).

Wysoka efektywność FlashAttention-3 bezpośrednio wpływa na rozwój i zastosowanie LLM:

• Skrócenie czasu trenowania: Przyspieszenie uwagi o 75–100% znacząco skraca czas trenowania modeli, który może trwać tygodnie lub miesiące.
• Zwiększenie okna kontekstowego: Modele mogą efektywnie przetwarzać dłuższe sekwencje (setki tysięcy tokenów), co jest istotne przy analizie dużych dokumentów lub kodu.
• Racjonalne wykorzystanie zasobów: Pozwala osiągać tę samą wydajność przy mniejszej liczbie GPU lub uzyskiwać większą szybkość na tym samym sprzęcie, co obniża koszty wdrożenia modeli.

Dostępność i integracja

Autorzy opublikowali kod źródłowy FlashAttention-3 na licencji otwartego oprogramowania w serwisie GitHub. Oczekuje się jego integracji z wiodącymi frameworkami głębokiego uczenia, takimi jak PyTorch i biblioteki Hugging Face Transformers, co uczyni technologię dostępną dla szerokiego grona programistów i badaczy. Poprzednie wersje stały się już de facto standardem w branży i FlashAttention-3 prawdopodobnie podtrzyma tę tendencję.

Odnośniki

• Oficjalne repozytorium FlashAttention na GitHub
• Blog Together AI z ogłoszeniem FlashAttention-3

Literatura

• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Chen, Y. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Liu, Y. et al. (2024). FastAttention: Extending FlashAttention-2 to NPUs and Low-Resource GPUs. OpenReview: 76NYyOrnfk.
• 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. (2024). FlashMask: Efficient and Rich Mask Extension of FlashAttention. arXiv:2410.01359.
• Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.