FlashAttention (CS)

FlashAttention — je revoluční algoritmus pro výpočet mechanismu attention, vyvinutý za účelem výrazného zrychlení trénování a inference velkých jazykových modelů (LLM) při zachování plné přesnosti výpočtů. Algoritmus byl poprvé představen v roce 2022 týmem výzkumníků ze Stanfordovy univerzity pod vedením Tri Dao.

Klíčová myšlenka FlashAttention spočívá v reorganizaci výpočtů s ohledem na hierarchii paměti GPU, což umožňuje minimalizovat počet přístupů k pomalé paměti a odstranit hlavní úzké místo standardního mechanismu attention.

Problematika standardního attention

Standardní mechanismus self-attention v transformerech se vypočítává podle vzorce: $extAttention(Q,K,V)=extsoftmax\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ kde Q, K, V jsou matice dotazů, klíčů a hodnot.

Hlavní problém tohoto přístupu je kvadratická složitost z hlediska času i paměti (O(N²)) vůči délce sekvence N. Při naivní implementaci je nutné vypočítat a uložit v paměti GPU celou matici attention S o velikosti N×N, což vede ke dvěma kritickým problémům:

1. Vysoká spotřeba paměti: Uložení matice N×N se stává nemožným při práci s dlouhými kontexty.
2. Operace vstupu a výstupu (IO): Hlavním úzkým místem není počet aritmetických operací, ale neustálé přístupy k pomalé paměti GPU.

Hierarchie paměti GPU

Pro pochopení problému je důležité rozlišovat dva typy paměti v GPU (na příkladu NVIDIA A100):

• SRAM (statická paměť): Rychlá paměť malého objemu přímo na čipu (~20 MB) s obrovskou propustností (až 19 TB/s).
• HBM (vysokopropustná paměť): Pomalá paměť velkého objemu (40–80 GB) s výrazně nižší propustností (přibližně 1,5 TB/s).

Tato asymetrie činí standardní algoritmus attention omezeným propustností paměti (memory-bound), protože neustále čte a zapisuje velké matice z pomalé HBM, což je hlavním zdrojem zpoždění.

Klíčové inovace FlashAttention

FlashAttention je IO-vnímavý (IO-aware) algoritmus, který řeší tento problém minimalizací přístupů k HBM. Toho je dosaženo pomocí tří základních technik.

Tiling a bloková zpracování

Místo zpracování celé matice najednou rozděluje FlashAttention vstupní matice Q, K, V na malé bloky (dlaždice), které se vejdou do rychlé SRAM. Algoritmus postupně načítá tyto bloky, provádí pro ně veškeré výpočty attention a aktualizuje výsledek, aniž by ukládal celou matici attention do pomalé HBM.

Online výpočet Softmax

Klíčovým technickým průlomem byl „online" výpočet Softmax. Standardní Softmax vyžaduje znalost všech prvků vstupního vektoru pro normalizaci. FlashAttention používá modifikovaný algoritmus, který umožňuje výpočet Softmax po částech. Udržuje dvě průběžné hodnoty (aktuální maximum a součet exponenciál), které jsou aktualizovány při zpracování nových bloků, což umožňuje získat přesný výsledek bez přístupu k celé matici najednou.

Sloučení operací do jednoho CUDA jádra

Všechny operace attention (maticové násobení QKᵀ, maskování, Softmax, násobení maticí V) jsou sloučeny do jediného fused CUDA jádra (fused kernel). To radikálně snižuje počet operací čtení/zápisu do HBM: místo opakovaných průchodů přes celou matici algoritmus načte blok do SRAM jednou, provede všechny výpočty a zapíše pouze konečný výsledek.

Teoretická a praktická účinnost

Složitost a optimalita

FlashAttention snižuje spotřebu paměti z O(N²) na O(N), což zajišťuje lineární škálování. Bylo prokázáno, že IO-složitost algoritmu je teoreticky optimální pro výpočet attention ve dvouúrovňové hierarchii paměti, to znamená, že přesnější attention nelze provést rychleji bez změny hardwaru.

Empirické výsledky

První verze FlashAttention prokázala výrazná zlepšení:

• Zrychlení:
  • BERT-large (délka 512): 15% zrychlení trénování.
  • GPT-2 (délka 1K): 3násobné zrychlení.
  • Úlohy Long-Range Arena (1K–4K): 2,4násobné zrychlení.
• Úspora paměti: Až 20násobná úspora paměti ve srovnání s přesnými základními implementacemi.
• Zlepšení kvality modelů: Díky možnosti pracovat s delšími kontexty FlashAttention nejenže neztrácí, ale i zlepšuje kvalitu modelů. Například perplexita GPT-2 se zlepšila o 0,7 bodu a přesnost v úlohách klasifikace dlouhých dokumentů vzrostla o 6,4 bodu.

Vývoj a další rozšíření

Úspěch FlashAttention odstartoval celou řadu hardwarově orientovaných algoritmů.

FlashAttention-2 (2023)

Druhá verze se zaměřila na plnější využití zdrojů GPU. V originálním FlashAttention dosahovala efektivita na NVIDIA A100 pouze 25–40 % maxima. FlashAttention-2 přinesla zlepšení v paralelizaci výpočtů, díky čemuž bylo dosaženo:

• Dvojnásobného zrychlení oproti první verzi.
• Zvýšení využití GPU na 50–73 % teoretického maxima.
• Rozšíření podpory na hlavy attention o velikosti 256 a také pro architektury Multi-Query Attention (MQA).

FlashAttention-3 (2024)

Třetí verze byla optimalizována speciálně pro architekturu GPU NVIDIA Hopper (H100). Využívá nové hardwarové možnosti, jako je asynchronnost Tensor Cores a podpora FP8, což umožnilo:

• Dosáhnout dalšího 1,5–2násobného zrychlení oproti FlashAttention-2.
• Dosáhnout výkonu až 740 TFLOPS na FP16 a blízko 1,2 PFLOPS na FP8.

Specializovaná řešení

Myšlenky FlashAttention byly rozvinuty v dalších projektech:

• FlashInfer (2025): Konfigurovatelný engine pro attention, optimalizovaný speciálně pro úlohy inference LLM. Zaměřuje se na efektivní práci s KV-cache v režimu proudové generace.
• FlashMLA (2024): Implementace attention se kompresí kontextové cache (latent attention), umožňující šetřit paměť na velmi dlouhých sekvencích s minimální ztrátou informací.

Vliv na průmysl a ekosystém

FlashAttention se stal základním průlomem a rychle se proměnil ve standard průmyslu pro efektivní trénování a inference LLM. Byl integrován do klíčových knihoven, jako jsou PyTorch a Hugging Face, a je využíván ve většině velkých jazykových modelů (LLaMA, MPT, Falcon, Claude a další).

Právě FlashAttention a jeho následující verze sehrály rozhodující roli při zvětšování kontextových oken jazykových modelů: z 2–4 tisíc tokenů (GPT-3) na 128 tisíc tokenů (GPT-4) a dokonce na miliony tokenů v experimentálních modelech. Algoritmus odstranil jednu z hlavních překážek škálování transformerů a otevřel nové možnosti pro AI aplikace, od analýzy dlouhých dokumentů až po multimodální porozumění.

Odkazy

• Oficiální repozitář FlashAttention na GitHubu

Literatura

• 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.