FlashAttention (BG)

FlashAttention — това е революционен алгоритъм за изчисляване на механизма на внимание (attention), разработен за значително ускоряване на обучението и инференса на големи езикови модели (LLM) при запазване на пълна точност на изчисленията. Алгоритъмът е представен за първи път през 2022 година от екип изследователи от Станфордския университет под ръководството на Три Дао (Tri Dao).

Ключовата идея на FlashAttention се състои в реорганизацията на изчисленията съобразно йерархията на паметта на GPU, което позволява да се минимизира броят на обръщенията към бавната памет и да се отстрани главното тясно място на стандартния механизъм на внимание.

Проблематика на стандартното внимание

Стандартният механизъм на самовнимание в трансформерите се изчислява по формулата: $extAttention(Q,K,V)=extsoftmax\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ където Q, K, V са матриците на заявките, ключовете и стойностите.

Основният проблем на този подход е квадратичната сложност по време и памет (O(N²)) спрямо дължината на последователността N. При наивна реализация е необходимо да се изчислява и съхранява в паметта на GPU пълната матрица на вниманието S с размер N×N, което води до два критични проблема:

1. Голямо потребление на памет: Съхраняването на матрицата N×N става невъзможно при работа с дълги контексти.
2. Операции за вход-изход (IO): Главното тясно място не е броят на аритметичните операции, а постоянните обръщения към бавната памет на GPU.

Йерархия на паметта на GPU

За разбирането на проблема е важно да се разграничат два вида памет в GPU (на примера на NVIDIA A100):

• SRAM (статична памет): Бърза вътрешнокристална памет с малък обем (~20 МБ) с огромна пропускателна способност (до 19 ТБ/с).
• HBM (памет с висока пропускателна способност): Бавна памет с голям обем (40–80 ГБ) с много по-малка пропускателна способност (около 1.5 ТБ/с).

Тази асиметрия прави стандартния алгоритъм за внимание ограничен от пропускателната способност на паметта (memory-bound), тъй като той постоянно чете и записва големи матрици от бавната HBM, което е главният източник на закъснения.

Ключови иновации на FlashAttention

FlashAttention е IO-осъзнат (IO-aware) алгоритъм, който решава проблема чрез минимизиране на обръщенията към HBM. Това се постига чрез три основни техники.

Тайлинг и блокова обработка

Вместо да обработва цялата матрица наведнъж, FlashAttention разделя входните матрици Q, K, V на малки блокове (тайлове), които се побират в бързата SRAM. Алгоритъмът последователно зарежда тези блокове, извършва за тях всички изчисления на вниманието и актуализира крайния резултат, без да съхранява пълната матрица на вниманието в бавната HBM.

Онлайн изчисляване на Softmax

Ключов технически пробив стана „онлайн" изчисляването на Softmax. Стандартният Softmax изисква познаване на всички елементи на входния вектор за нормализация. FlashAttention използва модифициран алгоритъм, който позволява изчисляването на Softmax на части. Той поддържа две междинни стойности (текущия максимум и сумата от експоненти), които се актуализират при обработката на нови блокове, което позволява получаването на точен резултат без достъп до цялата матрица наведнъж.

Сливане на операциите в едно CUDA ядро

Всички операции на вниманието (матрично умножение QKᵀ, маскиране, Softmax, умножение по V) са обединени в единно слято CUDA ядро (fused kernel). Това кардинално намалява броя на операциите за четене/запис в HBM: вместо многократни преминавания по цялата матрица, алгоритъмът зарежда блока в SRAM веднъж, извършва всички изчисления и записва само крайния резултат.

Теоретична и практическа ефективност

Сложност и оптималност

FlashAttention намалява потреблението на памет от O(N²) до O(N), което осигурява линейно мащабиране. Доказано е, че IO-сложността на алгоритъма е теоретично оптимална за изчисляване на внимание в двустепенна йерархия на паметта, тоест по-бързото изпълнение на точното внимание е невъзможно без промяна на хардуерната част.

Емпирични резултати

Първата версия на FlashAttention демонстрира значителни подобрения:

• Ускорение:
  • BERT-large (дължина 512): 15% ускорение на обучението.
  • GPT-2 (дължина 1K): 3-кратно ускорение.
  • Задачи Long-Range Arena (1K-4K): 2.4-кратно ускорение.
• Икономия на памет: До 20-кратна икономия на памет в сравнение с точните базови реализации.
• Подобряване на качеството на моделите: Благодарение на възможността за работа с по-дълги контексти, FlashAttention не само не губи, но и подобрява качеството на моделите. Например, перплексността на GPT-2 се подобри с 0.7 точки, а точността при задачи за класификация на дълги документи нарасна с 6.4 точки.

Еволюция и по-нататъшно развитие

Успехът на FlashAttention постави началото на цяла серия от хардуерно-ориентирани алгоритми.

FlashAttention-2 (2023)

Втората версия беше насочена към по-пълно използване на ресурсите на GPU. В оригиналния FlashAttention ефективността на NVIDIA A100 беше едва 25–40% от максимума. FlashAttention-2 въведе подобрения в паралелизацията на изчисленията, което позволи:

• Да се постигне двукратно ускорение в сравнение с първата версия.
• Да се увеличи утилизацията на GPU до 50–73% от теоретичния максимум.
• Да се разшири поддръжката до глави на внимание с размер 256, както и за архитектури Multi-Query Attention (MQA).

FlashAttention-3 (2024)

Третата версия беше оптимизирана специално за архитектурата на GPU NVIDIA Hopper (H100). Тя използва нови хардуерни възможности, като асинхронност на Tensor Cores и поддръжка на FP8, което позволи:

• Да се постигне допълнително 1.5–2-кратно ускорение в сравнение с FlashAttention-2.
• Да се достигне производителност до 740 TFLOPS при FP16 и близо до 1.2 PFLOPS при FP8.

Специализирани решения

Идеите на FlashAttention са развити в други проекти:

• FlashInfer (2025): Настраиваем двигател за внимание, оптимизиран специално за задачи за инференс на LLM. Той се фокусира върху ефективната работа с KV-кеша в режим на поточна генерация.
• FlashMLA (2024): Реализация на внимание със свиване на контекстния кеш (latent attention), позволяваща икономия на памет при много дълги последователности с минимална загуба на информация.

Влияние върху индустрията и екосистемата

FlashAttention се превърна в фундаментален пробив и бързо стана индустриален стандарт за ефективно обучение и инференс на LLM. Той е интегриран в ключови библиотеки като PyTorch и Hugging Face и се използва в повечето големи езикови модели (LLaMA, MPT, Falcon, Claude и др.).

Именно FlashAttention и неговите следващи версии изиграха решаваща роля в увеличаването на контекстните прозорци на езиковите модели: от 2–4 хиляди токени (GPT-3) до 128 хиляди токени (GPT-4) и дори до милиони токени в експериментални модели. Алгоритъмът отстрани една от главните пречки по пътя на мащабирането на трансформерите, откривайки нови възможности за AI приложения — от анализ на дълги документи до мултимодално разбиране.

Препратки

• Официално хранилище на FlashAttention в GitHub

Литература

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