FlashAttention-3 (CS)

FlashAttention-3 — je algoritmus pro optimalizaci mechanismu pozornosti (attention) v transformerových neuronových sítích, vyvinutý pro maximální využití hardwarových možností GPU architektury NVIDIA Hopper (H100). Algoritmus byl představen v roce 2024 skupinou výzkumníků ze společností Colfax Research, Meta, NVIDIA, Georgia Tech, Princetonské univerzity a Together AI. Práce byla přijata na konferenci NeurIPS 2024 a označena jako spotlight.

FlashAttention-3 je třetí iterací v rodině algoritmů navazujících na FlashAttention (2022) a FlashAttention-2 (2023). Jeho hlavním cílem je výrazně urychlit trénování a inferenci velkých jazykových modelů (LLM) při zachování přesnosti výpočtů.

Úvod a předhistorie

Problém mechanismu pozornosti

Klíčovou součástí transformerů je mechanismus vlastního pozornosti (self-attention), jehož výpočetní složitost a spotřeba paměti však rostou kvadraticky (O(n²)) s rostoucí délkou vstupní sekvence (n). To vytváří závažné „úzké hrdlo", protože moderní GPU jsou optimalizována pro rychlé maticové násobení, zatímco výpočet exponenciálních funkcí (například v Softmax) probíhá řádově pomaleji. Navíc při naivní implementaci musí být v paměti GPU uložen velký mezilehlý tensor pozornosti, což omezuje škálovatelnost modelů.

FlashAttention a FlashAttention-2

Pro řešení tohoto problému byl v roce 2022 navržen FlashAttention, který snížil počet přístupů k pomalé globální paměti (HBM) pomocí dvou technik:

• Bloková zpracování (tiling): Výpočty jsou rozděleny do bloků (dlaždic), které jsou zpracovávány v rychlé on-chip paměti (SRAM).
• Slučování operací: Všechny operace (maticové násobení, Softmax) jsou prováděny v jednom jádře GPU bez zápisu mezilehlých výsledků do globální paměti.

To umožnilo snížit paměťovou složitost z kvadratické na lineární a zrychlit výpočty 2–4krát.

V roce 2023 byla představena vylepšená verze — FlashAttention-2, která optimalizovala paralelizaci výpočtů. Na GPU architektury NVIDIA Ampere (A100) dosáhla ~70 % teoretického výkonnostního maxima. Na novější architektuře NVIDIA Hopper (H100) však byla její účinnost výrazně nižší — přibližně 35 %. Důvodem bylo, že algoritmus nevyužíval nové hardwarové možnosti Hopperu, což vedlo ke vzniku FlashAttention-3.

Nové hardwarové možnosti GPU Hopper (H100)

Architektura NVIDIA Hopper přinesla řadu nových funkcí, které FlashAttention-3 využívá k dosažení maximálního výkonu:

• WGMMA (Warpgroup Matrix Multiply-Accumulate): Nový typ instrukcí pro tensorová jádra provádějící maticové násobení s téměř dvojnásobným nárůstem výkonu oproti architektuře Ampere.
• TMA (Tensor Memory Accelerator): Hardwarový modul urychlující přenos dat mezi globální (HBM) a sdílenou (shared memory) pamětí. TMA automaticky provádí adresní výpočty a odlehčuje výpočetní jádra.
• Formát FP8: Hardwarová podpora 8bitového formátu dat s plovoucí desetinnou čárkou, který zdvojnásobuje teoretický výkon oproti FP16, avšak nese riziko ztráty přesnosti kvůli omezenému dynamickému rozsahu.

Technické inovace FlashAttention-3

Algoritmus implementuje tři klíčové optimalizační metody speciálně navržené pro architekturu Hopper:

1. Asynchronní provádění a specializace warpů

FlashAttention-3 využívá princip warp-specialization, při němž se různé skupiny vláken (warps) na GPU specializují na různé úkoly:

• Producer warps: Načítají data z globální paměti pomocí TMA.
• Consumer warps: Provádějí maticové násobení na tensorových jádrech.

Díky hardwarové asynchronnosti Hopperu se tyto operace časově překrývají. Zatímco jedna skupina warpů provádí výpočty, druhá paralelně načítá data pro následující blok. Tento zřetězený přístup (pipeline) organizovaný podle principu „ping-pong" (ping-pong scheduling) umožňuje skrýt latence pomalých operací (například Softmax) a maximálně vytížit všechny funkční moduly GPU.

2. Minimalizace operací s pamětí

Algoritmus zachovává filozofii tiling z předchozích verzí, ale aktivně využívá TMA pro asynchronní načítání dalších bloků dat paralelně s aktuálními výpočty. Přenos dat z pomalé HBM do rychlé SRAM probíhá fakticky „ve stínu" hlavních výpočtů, díky čemuž GPU méně čeká na data.

3. Nízká přesnost (FP8) se snížením chyby kvantizace

Přechod na FP8 zdvojnásobuje rychlost, ale může vést k výrazné ztrátě přesnosti kvůli kvantizaci. K boji proti tomuto jevu vývojáři zavedli metodu incoherent processing. Její podstata spočívá v následujícím:

1. Před výpočtem pozornosti jsou vektory příznaků (dotazy Q a klíče K) vynásobeny náhodnou ortogonální maticí (například Hadamardovou maticí).
2. Tato transformace „rozprostře" hodnoty s anomálně velkým modulem (odlehlé hodnoty) po všech souřadnicích a vyrovná jejich rozložení.
3. Poté je provedena kvantizace do FP8, která nyní probíhá s menší chybou.
4. Protože transformace je ortogonální, nezkresuje výsledek pozornosti (QKᵀ), neboť efekt matice se při násobení vyruší.

Tato technika umožnila snížit chybu výpočtu pozornosti v FP8 přibližně 2,6krát ve srovnání se standardním použitím FP8 bez transformací.

Výkonnost a význam

Uplatnění uvedených technik umožnilo FlashAttention-3 dosáhnout výrazné převahy nad předchozími verzemi na GPU H100:

• Zrychlení 1,5–2krát oproti FlashAttention-2.
• Vysoké využití GPU: Dosahuje ~75–85 % teoretického výkonnostního maxima H100.
• Propustnost:
  • Až 740–840 TFLOPS pro poloviční přesnost (FP16/BF16).
  • Až 1,2–1,3 PFLOPS (petaflopů) při použití 8bitové přesnosti (FP8).

Vysoká účinnost FlashAttention-3 přímo ovlivňuje vývoj a nasazení LLM:

• Zkrácení doby trénování: Zrychlení pozornosti o 75–100 % výrazně zkracuje dobu trénování modelů, která může trvat týdny nebo měsíce.
• Zvětšení kontextového okna: Modely mohou efektivně zpracovávat delší sekvence (stovky tisíc tokenů), což je důležité pro analýzu rozsáhlých dokumentů nebo kódu.
• Racionální využití zdrojů: Umožňuje dosáhnout stejného výkonu na menším počtu GPU nebo získat vyšší rychlost na stejném hardwaru, čímž se snižují náklady na nasazení modelů.

Dostupnost a integrace

Autoři zveřejnili zdrojový kód FlashAttention-3 pod otevřenou licencí na GitHubu. Očekává se jeho integrace do předních frameworků hlubokého učení, jako jsou PyTorch a knihovny Hugging Face Transformers, což zpřístupní technologii širokému okruhu vývojářů a výzkumníků. Předchozí verze se již staly de facto standardem v odvětví a FlashAttention-3 pravděpodobně v tomto trendu pokračuje.

Odkazy

• Oficiální repozitář FlashAttention na GitHubu
• Blog Together AI s oznámením 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.