FlashAttention-2

Материал из Systems analysis wiki
Перейти к навигации Перейти к поиску

FlashAttention-2 — это усовершенствованный алгоритм, предназначенный для вычисления механизма внимания (attention) в больших языковых моделях (LLM). Алгоритм был разработан Три Дао (Tri Dao) и исследователями из Стэнфордского университета и представлен в июле 2023 года[1]. Его ключевая цель — значительно ускорить обучение и инференс (вывод) моделей-трансформеров за счёт более эффективного использования аппаратных ресурсов GPU, при этом сохраняя полную идентичность вычислений со стандартным механизмом внимания, то есть без потери точности.

FlashAttention-2 является логическим продолжением алгоритма FlashAttention, представленного той же командой в 2022 году. Новая версия решает проблему неполной загрузки GPU, которая наблюдалась у предшественника, и достигает почти двукратного прироста скорости по сравнению с первой версией.

Предпосылки: проблема внимания в трансформерах

Стандартный механизм самовнимания (self-attention) является узким местом при работе с длинными последовательностями текста в трансформерах. Его вычислительная сложность и потребление памяти растут квадратично (O(N²)) в зависимости от длины последовательности (N), что накладывает серьёзные ограничения на максимальную длину контекста и масштабируемость LLM[1].

Для решения этой проблемы в 2022 году был представлен алгоритм FlashAttention[2]. Его ключевые идеи:

  • Учёт иерархии памяти GPU (IO-awareness): Алгоритм минимизирует дорогостоящие операции чтения/записи между медленной памятью GPU (HBM) и быстрой статической памятью (SRAM) на чипе.
  • Блочная обработка (tiling): Вычисления разбиваются на небольшие блоки (тайлы), которые обрабатываются в быстрой SRAM, что позволяет избежать материализации полной матрицы внимания в памяти.

Это позволило достичь линейного роста потребления памяти (O(N)) и ускорения в 2–4 раза по сравнению со стандартными реализациями[2]. FlashAttention получил широкое распространение и способствовал появлению моделей с значительно увеличенным контекстом, например, с 2–4 тыс. токенов (GPT-3) до 128 тыс. (GPT-4) и более[3]. Так, в модели Falcon-40B использование FlashAttention ускорило инференс в 3 раза, а общую производительность генерации — в 5 раз по сравнению с GPT-3[4].

Разработка и цели FlashAttention-2

Несмотря на успех, первая версия FlashAttention использовала вычислительные ресурсы GPU не полностью. На видеокартах NVIDIA A100 производительность достигала лишь 25–40% от теоретического максимума (FLOPs/s)[1]. Основной причиной была неоптимальная загрузка потоковых мультипроцессоров (Streaming Multiprocessors) и избыточные операции с общей памятью[5].

Целью FlashAttention-2 стало дальнейшее ускорение вычислений за счёт более эффективного распараллеливания работы и минимизации вспомогательных операций. Алгоритм был полностью переписан с использованием низкоуровневых примитивов библиотеки NVIDIA CUTLASS 3.x для достижения максимальной производительности[6].

Техническая архитектура и принципы работы

FlashAttention-2 вводит три ключевых улучшения для повышения параллелизма и эффективности[1]:

1. Минимизация не-матричных операций

Алгоритм сокращает количество вспомогательных операций с плавающей точкой, не являющихся матричным умножением (non-matmul FLOPs). Поскольку тензорные ядра GPU оптимизированы именно под матричные операции (GEMM) и выполняют их до 16 раз быстрее, это изменение позволяет большую часть времени использовать наиболее производительные блоки GPU.

2. Улучшенный параллелизм

В оригинальном FlashAttention работа над одной "головой" внимания не распараллеливалась, что приводило к простоям при длинных последовательностях и малых размерах батча. FlashAttention-2 вводит межблочное распараллеливание: теперь вычисления для одной головы внимания распределяются между разными потоковыми мультипроцессорами GPU, что значительно повышает их загрузку.

3. Оптимизированное разделение работы внутри блока

На уровне одного вычислительного блока работа была перераспределена между группами потоков (warps) для уменьшения обмена данными через общую память (shared memory). Это снижает количество избыточных операций чтения/записи, необходимых для нормализации Softmax.

Производительность и эффективность

Благодаря архитектурным улучшениям, FlashAttention-2 демонстрирует значительный прирост производительности:

  • Двукратное ускорение: Алгоритм работает примерно в 2 раза быстрее по сравнению с первой версией FlashAttention[1].
  • Высокая утилизация GPU: На GPU NVIDIA A100 достигается 50–73% от теоретической максимальной пропускной способности (TFLOPs), что близко к эффективности оптимизированных операций матричного умножения (GEMM)[1].
  • Рекордная скорость вычислений:
    • На GPU A100 достигается скорость до 225 TFLOP/s в сквозном цикле обучения модели типа GPT, что соответствует 72% утилизации вычислительных блоков. Для сравнения, стандартное внимание в тех же условиях нагружало GPU менее чем на 100 TFLOP/s[7].
    • На GPU H100 производительность достигает 335 TFLOP/s[7].

Такой прирост производительности позволяет, например, обучать модель с контекстным окном 16k токенов за то же время, что ранее требовалось для окна 8k токенов[5]. Важно, что алгоритм остаётся точным и детерминированным, поэтому его применение не влияет на качество предсказаний модели[8].

Применение и интеграция в экосистему

FlashAttention-2 быстро стал стандартным инструментом в экосистеме LLM. Он интегрирован во многие популярные фреймворки и библиотеки:

  • PyTorch: Нативная поддержка.
  • Hugging Face Transformers: Поддержка включается параметром `attn_implementation="flash_attention_2"` при загрузке модели[9]. Совместим с десятками архитектур (GPT, Llama, Falcon, BERT и др.)[10].
  • TensorRT-LLM, xFormers и Triton: Алгоритм реализован для этих платформ, что обеспечивает широкое применение[7].

Интеграция позволяет легко сочетать FlashAttention-2 с другими методами оптимизации, такими как квантование (GPTQ, QLoRA) и эффективное дообучение (PEFT)[9].

Сравнение с последующими версиями

FlashAttention-3

Исследования в области оптимизации внимания продолжаются. В июле 2024 года Три Дао представил FlashAttention-3, нацеленный на использование возможностей архитектуры GPU NVIDIA Hopper (H100/H200). Ключевые нововведения[3]:

  • Поддержка FP8: Использует 8-битные вычисления с плавающей точкой для дальнейшего ускорения.
  • Асинхронные операции: Более эффективно использует асинхронные возможности GPU.

FlashAttention-3 обеспечивает ускорение в 1.5–2 раза по сравнению с FlashAttention-2 на GPU H100, достигая производительности до 740 TFLOP/s (75% от теоретического максимума)[11].

Литература

  • 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.
  • 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.
  • 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. OpenReview: rog0J435OO.
  • Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.

Примечания

  1. ↑ 1,0 1,1 1,2 1,3 1,4 1,5 Дао, Три. «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning». arXiv:2307.08691 [cs.LG], 17 июля 2023 г. [1]
  2. ↑ 2,0 2,1 «Optimizing LLMs for Speed and Memory». Hugging Face Documentation. [2]
  3. ↑ 3,0 3,1 Дао, Три. «FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision». Tri Dao's Blog. [3]
  4. ↑ «FlashAttention vs FlashAttention-2 - an Analysis». E2E Networks Blog. [4]
  5. ↑ 5,0 5,1 «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning». OpenReview. [5]
  6. ↑ «FlashAttention-2». Hazy Research, Stanford University. [6]
  7. ↑ 7,0 7,1 7,2 Дао, Три. «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning» (PDF). arXiv:2307.08691. [7]
  8. ↑ Рашка, Себастьян. «Llama 2 and FlashAttention 2». Ahead of AI Magazine. [8]
  9. ↑ 9,0 9,1 Белькада, Юнес. «Faster and more memory efficient models with Flash Attention 2!». LinkedIn. [9]
  10. ↑ «GPU inference». Hugging Face Documentation. [10]
  11. ↑ Дао, Три, и др. «FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision». arXiv:2407.08608 [cs.LG], 11 июля 2024 г. [11]


Источник — https://systems-analysis.ru/wiki/index.php?title=FlashAttention-2&oldid=2537