FlashAttention-4: новый алгоритм внимания, обходящий аппаратные ограничения Blackwell

Введение

Разработчики компании Together AI официально представили FlashAttention-4 — новый алгоритм ядра внимания, оптимизированный специально под архитектуру Blackwell B200 и устраняющий узкие места этих передовых GPU. Разработчики сообщают о значительном приросте производительности: на B200 достигается 1605 TFLOPS (71% от пика), что в 1,3 раза быстрее cuDNN 9.13 и в 2,7 раза быстрее реализаций на Triton. Ключевые нововведения включают программную эмуляцию экспоненты для softmax (позволяет обойти медленные специализированные блоки), использование тензорной памяти TMEM для промежуточных данных (снижает нагрузку на общую память) и режим 2-CTA MMA (два блока потоков выполняют одно умножение с совместным доступом к TMEM).

Подробнее о FlashAttention-4

FlashAttention-4 строится на нескольких ключевых идеях совместного проектирования алгоритма и ядра, максимально использующих новые возможности Blackwell.

В прямом и обратном проходах внедрены программные конвейеры, которые полностью задействуют асинхронную многопоточность Blackwell. Укрупненные размеры тайлов позволяют перекрывать выполнение тензорных ядер, вычисление экспоненты softmax и операции с памятью. Это снижает простои и повышает утилизацию. Кроме того, вместо обращения к медленным SFU для вычисления exp, FlashAttention-4 использует программную эмуляцию — полиномиальную аппроксимацию на обычных FMA-блоках. Это устраняет бутылочное горлышко экспоненты. Дополнительно применяется условное масштабирование softmax, адаптирующееся к данным.

Еще одна функция — сохранение промежуточных результатов (например, значений softmax) во встроенной тензорной памяти (TMEM) Blackwell, а не в общей памяти. Это радикально уменьшает трафик. Кроме того, используется новый режим 2-CTA MMA, позволяющий двум блокам потоков совместно выполнять одно матричное умножение с разделением TMEM. Это вдвое сокращает объем атомарных операций и вдвое уменьшает потребность в общей памяти. Также реализована поддержка детерминированного выполнения для воспроизводимого обучения.

Для устранения дисбаланса нагрузки, вызванного причинно-следственными масками и переменной длиной последовательностей, в FlashAttention-4 встроен адаптивный планировщик тайлов, который динамически распределяет работу между потоковыми мультипроцессорами.

Эффективность FlashAttention-4 опирается на три ключевые инновации B200:

• Тензорная память (TMEM) — 256 КБ на SM, встроенная в тензорные ядра, для быстрого хранения промежуточных данных.
• Полностью асинхронные тензорные ядра 5-го поколения (tcgen05.mma), которые работают независимо и могут адресовать операнды прямо из TMEM. Размер плитки CTA UMMA увеличен до 128×256×16, что вдвое больше атома Hopper.
• 2-CTA MMA — режим, в котором одно матричное умножение выполняется парой CTA, совместно использующих TMEM. Это позволяет масштабировать плитку до 256×256×16 и сокращает избыточный трафик.

В результате на B200 с BF16 FlashAttention-4 достигает 1605 TFLOPS (71% от пика), что в 1,3 раза быстрее cuDNN 9.13 и в 2,7 раза быстрее Triton.

Выводы

FlashAttention-4 демонстрирует, что дальнейший прогресс в оптимизации внимания требует не просто улучшения отдельных операций, а переосмысления всего конвейера с учетом асимметричного роста аппаратных возможностей ИИ-ускорителей. Опираясь на новые функции Blackwell, такие как тензорную память, асинхронные тензорные ядра и 2-CTA MMA, разработчики смогли перенести узкие места из softmax и shared memory в тензорные ядра, добившись рекордной производительности.