Flash Attention

Wprowadzenie

Flash Attention to jeden z najważniejszych przełomów w optymalizacji Transformerów od czasu ich powstania. Zaproponowany w 2022 roku przez Tri Dao, całkowicie zmienił sposób, w jaki trenujemy i uruchamiamy duże modele językowe.

Głównym problemem standardowego mechanizmu Attention jest jego kwadratowa złożoność pamięciowa (O(n²)). Flash Attention rozwiązuje ten problem poprzez sprytne wykorzystanie hierarchii pamięci GPU (SRAM + HBM).

Jak działa Flash Attention?

Zamiast materializować całą macierz Attention (n × n), algorytm wykonuje obliczenia w sposób IO-aware – optymalizuje transfery danych między wolną pamięcią HBM a szybką pamięcią SRAM.

Dzięki temu unika zapisywania ogromnej macierzy attention scores w pamięci, co drastycznie zmniejsza zużycie VRAM i zwiększa prędkość.

Wersje Flash Attention

  • Flash Attention (v1) – 2022 – pierwsza rewolucyjna wersja
  • Flash Attention 2 – 2023 – jeszcze większa optymalizacja, lepsza paralelizacja
  • Flash Attention 3 – 2024/2025 – wsparcie dla Hopper GPU (RTX 40xx / H100), asynchroniczne obliczenia, FP8

Zalety Flash Attention

  • Do 3–4× szybszy trening i inferencja
  • Znacznie niższe zużycie pamięci VRAM (umożliwia dłuższe konteksty)
  • Możliwość trenowania modeli z kontekstem 64k–128k+ tokenów
  • Lepsze wykorzystanie GPU (wyższa occupancy)
  • Brak kompromisu w dokładności (bit-for-bit identyczne wyniki)

Zastosowania w praktyce (2026)

  • Wszystkie nowoczesne modele open-source (Llama 3, Mistral, Qwen2, Gemma 2)
  • Trening i fine-tuning długok kontekstowych modeli
  • Inferencja na konsumenckich kartach graficznych
  • Modele multimodalne i długie dokumenty
  • Agentyczne systemy wymagające długiej pamięci kontekstowej

Najlepsze praktyki

  • Zawsze włącz Flash Attention 2 lub 3 w bibliotekach (Hugging Face, vLLM, Axolotl, Unsloth)
  • Przy QLoRA + Flash Attention można trenować 70B modele na 2× RTX 4090
  • Używaj `torch.compile` razem z Flash Attention dla dodatkowych zysków
  • Przy bardzo długich kontekstach łącz z Ring Attention lub Infini-Attention

Powiązane pojęcia

Transformer • Attention Mechanism • Efficient Attention • SDPA (Scaled Dot Product Attention) • Long Context Training • vLLM • xFormers • Paged Attention