Flash Attention

Flash Attention — IO-эффективный алгоритм вычисления механизма внимания в трансформерах, сокращающий потребление памяти GPU с O(n²) до O(n) за счёт тайлирования и перевычисления. Разработан в Стэнфордском университете, впервые опубликован в 2022 году; де-факто стандарт обучения и инференса LLM.

Flash Attention — IO-эффективная реализация scaled dot-product attention, стандартного блока трансформерной архитектуры. Ключевое нововведение состоит не в изменении математики внимания, а в оптимизации схемы обращений к памяти GPU: алгоритм избегает записи промежуточной матрицы внимания размером n×n в медленную высокопропускную память (HBM) и вместо этого выполняет вычисления в быстром SRAM-кеше (shared memory).

Стандартный attention вычисляет A = softmax(QKᵀ/√d)V, где для последовательности длиной n требуется O(n²) памяти HBM для хранения A. Flash Attention разбивает Q, K, V на тайлы, умещающиеся в SRAM, и итеративно накапливает результат с применением численно устойчивой online-версии softmax. Для обратного прохода матрица A не хранится явно — используются исходные входные данные и recomputation, что принципиально снижает пиковое потребление памяти. Flash Attention 2 (2023) оптимизировал параллелизм по head-измерению и достиг порядка 70% теоретического максимума MFU на A100. Flash Attention 3 (2024) использует асинхронное выполнение WGMMA-инструкций и FP8-арифметику на архитектуре H100/H800.

Flash Attention стал стандартным компонентом практически всех современных фреймворков: PyTorch (встроен в F.scaled_dot_product_attention начиная с версии 2.0), HuggingFace Transformers, vLLM, llama.cpp. Алгоритм позволяет работать со значительно более длинными последовательностями при том же GPU-бюджете: переход на Flash Attention 2 позволяет увеличить длину последовательности в 4–8 раз без роста VRAM, что прямо ускорило развитие моделей с длинным контекстом.

К 2026 году Flash Attention де-факто обязателен при обучении и инференсе LLM. Три Дао (Tri Dao), первый автор оригинальных работ, продолжает публиковать улучшенные версии алгоритма под новые GPU-архитектуры, включая Blackwell от NVIDIA. Параллельно развиваются смежные подходы: FlexAttention (PyTorch), Ring Attention для распределённого контекста, Paged Attention (vLLM) для эффективного управления KV-кешем при инференсе.