FlashAttention-3

FlashAttention-3 — это алгоритм для оптимизации механизма внимания (attention) в трансформерных нейросетях, разработанный для максимального использования аппаратных возможностей GPU архитектуры NVIDIA Hopper (H100). Алгоритм был представлен в 2024 году группой исследователей из компаний Colfax Research, Meta, NVIDIA, Georgia Tech, Принстонский университет и Together AI. Работа была принята на конференцию NeurIPS 2024 и отмечена как spotlight.

FlashAttention-3 является третьей итерацией в семействе алгоритмов, следующих за FlashAttention (2022) и FlashAttention-2 (2023). Его главная цель — значительно ускорить обучение и инференс больших языковых моделей (LLM), сохранив при этом точность вычислений.

Введение и предыстория

Проблема механизма внимания

Ключевым компонентом трансформеров является механизм самовнимания (self-attention), однако его вычислительная сложность и потребление памяти растут квадратично (O(n²)) с увеличением длины входной последовательности (n). Это создаёт серьёзное "узкое место", так как современные GPU оптимизированы для быстрых матричных умножений, но вычисление экспоненциальных функций (например, в Softmax) происходит на порядки медленнее. Кроме того, при наивной реализации в памяти GPU должен храниться большой промежуточный тензор внимания, что ограничивает масштабируемость моделей.

FlashAttention и FlashAttention-2

Для решения этой проблемы в 2022 году был предложен FlashAttention, который сократил объём обращений к медленной глобальной памяти (HBM) за счёт двух техник:

• Блочная обработка (tiling): Вычисления разбиваются на блоки (тайлы), которые обрабатываются в быстрой on-chip памяти (SRAM).
• Слияние операций: Все операции (матричное умножение, Softmax) выполняются в одном ядре GPU без записи промежуточных результатов в глобальную память.

Это позволило снизить сложность по памяти с квадратичной до линейной и ускорило вычисления в 2–4 раза.

В 2023 году была представлена улучшенная версия — FlashAttention-2, которая оптимизировала распараллеливание вычислений. На GPU архитектуры NVIDIA Ampere (A100) она достигла ~70% от пиковой теоретической производительности. Однако на более новой архитектуре NVIDIA Hopper (H100) её эффективность оказалась значительно ниже — около 35%. Это было связано с тем, что алгоритм не использовал новые аппаратные возможности Hopper, что и послужило толчком к созданию FlashAttention-3.

Новые аппаратные возможности GPU Hopper (H100)

Архитектура NVIDIA Hopper предоставила ряд новых функций, которые FlashAttention-3 использует для достижения максимальной производительности:

• WGMMA (Warpgroup Matrix Multiply-Accumulate): Новый тип инструкций для тензорных ядер, выполняющий матричные умножения с почти двукратным приростом производительности по сравнению с архитектурой Ampere.
• TMA (Tensor Memory Accelerator): Аппаратный модуль, который ускоряет передачу данных между глобальной (HBM) и общей (shared memory) памятью. TMA автоматически выполняет адресные расчёты, разгружая вычислительные ядра.
• Формат FP8: Аппаратная поддержка 8-битного формата данных с плавающей точкой, который удваивает теоретическую производительность по сравнению с FP16, но несёт риск потери точности из-за ограниченного динамического диапазона.

Технические инновации FlashAttention-3

Алгоритм реализует три ключевых метода оптимизации, специально разработанных для архитектуры Hopper:

1. Асинхронное выполнение и специализация варпов

FlashAttention-3 использует принцип warp-specialization, при котором разные группы потоков (warps) на GPU специализируются на разных задачах:

• Producer warps: Загружают данные из глобальной памяти с помощью TMA.
• Consumer warps: Выполняют матричные умножения на тензорных ядрах.

Благодаря аппаратной асинхронности Hopper, эти операции перекрываются по времени. Пока одна группа варпов выполняет вычисления, другая параллельно загружает данные для следующего блока. Этот конвейерный подход (pipeline), организованный по принципу «пинг-понга» (ping-pong scheduling), позволяет скрыть задержки от медленных операций (например, Softmax) и максимально полно загрузить все функциональные модули GPU.

2. Минимизация операций с памятью

Алгоритм сохраняет идеологию tiling из предыдущих версий, но активно использует TMA для асинхронной загрузки следующих блоков данных параллельно с текущими вычислениями. Передача данных из медленной HBM в быструю SRAM фактически выполняется «в тени» основных вычислений, благодаря чему GPU меньше простаивает в ожидании данных.

3. Низкая точность (FP8) с уменьшением ошибки квантования

Переход на FP8 удваивает скорость, но может привести к существенной потере точности из-за квантования. Для борьбы с этим разработчики внедрили метод incoherent processing. Его суть заключается в следующем:

1. Перед вычислением внимания векторы признаков (запросы Q и ключи K) умножаются на случайную ортогональную матрицу (например, матрицу Адамара).
2. Это преобразование «размазывает» значения с аномально большим модулем (выбросы) по всем координатам, выравнивая их распределение.
3. После этого выполняется квантование в FP8, которое теперь проходит с меньшей ошибкой.
4. Поскольку преобразование ортогонально, оно не искажает итоговый результат внимания (QKᵀ), так как эффект от матрицы нивелируется при перемножении.

Эта техника позволила сократить ошибку вычисления внимания в FP8 примерно в 2,6 раза по сравнению со стандартным применением FP8 без преобразований.

Производительность и значение

Применение перечисленных техник позволило FlashAttention-3 достичь значительного превосходства над предыдущими версиями на GPU H100:

• Ускорение в 1.5–2 раза по сравнению с FlashAttention-2.
• Высокая утилизация GPU: Достигает ~75–85% от теоретического максимума производительности H100.
• Пропускная способность:
  • До 740–840 TFLOPS для половинной точности (FP16/BF16).
  • До 1.2–1.3 PFLOPS (петафлопс) при использовании 8-битной точности (FP8).

Высокая эффективность FlashAttention-3 напрямую влияет на разработку и применение LLM:

• Сокращение времени обучения: Ускорение внимания на 75–100% значительно сокращает время обучения моделей, которое может занимать недели или месяцы.
• Увеличение контекстного окна: Модели могут эффективно обрабатывать более длинные последовательности (сотни тысяч токенов), что важно для анализа больших документов или кода.
• Рациональное использование ресурсов: Позволяет достигать той же производительности на меньшем количестве GPU или получать большую скорость на том же оборудовании, что снижает стоимость развёртывания моделей.

Доступность и интеграция

Авторы опубликовали исходный код FlashAttention-3 под открытой лицензией в GitHub. Ожидается его интеграция в ведущие фреймворки глубокого обучения, такие как PyTorch и библиотеки Hugging Face Transformers, что сделает технологию доступной для широкого круга разработчиков и исследователей. Предыдущие версии уже стали стандартом де-факто в индустрии, и FlashAttention-3, вероятно, продолжит эту тенденцию.

Ссылки

• Официальный репозиторий FlashAttention на GitHub
• Блог Together AI с анонсом FlashAttention-3

Литература

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