Ускорьте внимание с FlashAttention-3: новые возможности и производительность

Внимание, как ключевой слой универсальной архитектуры Transformer, представляет собой узкое место для больших языковых моделей и приложений с длинным контекстом. FlashAttention и FlashAttention-2 стали пионерами в подходе к ускорению внимания на GPU, минимизируя чтение и запись памяти, и сейчас используются большинством библиотек для ускорения обучения и вывода Transformer. Это способствовало значительному увеличению длины контекста LLM за последние два года, с 2-4K (GPT-3, OPT) до 128K (GPT-4) и даже 1M (Llama 3). Однако, несмотря на успех, FlashAttention еще не использовал новые возможности современного оборудования. FlashAttention-2 достигал только 35% использования теоретической максимальной производительности FLOPs на GPU H100.

В этом блоге мы описываем три основных техники, позволяющих ускорить внимание на GPU Hopper: использование асинхронности Tensor Cores и TMA для (1) наложения общей вычислительной нагрузки и перемещения данных через специальизацию ворков и (2) интерливинг блочного умножения матриц и операций softmax, и (3) некоординированная обработка, использующая поддержку оборудования для низкой точности FP8. Мы рады представить FlashAttention-3, который включает эти техники. Он в 1.5-2.0 раза быстрее FlashAttention-2 с FP16, достигая до 740 TFLOPS, то есть 75% использования теоретической максимальной производительности FLOPs H100. С FP8 FlashAttention-3 достигает почти 1.2 PFLOPS с ошибкой в 2.6 раза меньше, чем базовое внимание FP8.

Улучшения от FlashAttention-3 приведут к более эффективному использованию GPU: новая техника использует до 75% максимальных возможностей GPU H100, в отличие от 35% ранее. Это обеспечивает значительно более быстрое (1.5-2x) выполнение по сравнению с предыдущими версиями для обучения и работы с большими языковыми моделями. Также FlashAttention-3 может работать с числами низкой точности (FP8), сохраняя при этом точность. Это позволяет еще быстрее обрабатывать данные и потенциально снижать потребление памяти, что может привести к экономии затрат и повышению эффективности для клиентов, работающих с крупномасштабными операциями ИИ.

Кроме того, ускорение механизма внимания позволяет моделям ИИ работать с гораздо более длинными текстами более эффективно. Это может открыть возможности для приложений, способных понимать и генерировать более длинный и сложный контент без замедления. FlashAttention-3 доступен на GitHub.

FlashAttention — это алгоритм, который переупорядочивает вычисление внимания и использует тайлинг и повторное вычисление для значительного ускорения и снижения использования памяти с квадратичного до линейного в зависимости от длины последовательности. Мы используем тайлинг для загрузки блоков входных данных из HBM (память GPU) в SRAM (быстрый кэш), выполняем внимание относительно этого блока и обновляем вывод в HBM. Не записывая большие промежуточные матрицы внимания в HBM, мы уменьшаем количество операций чтения/записи памяти, что приводит к ускорению времени выполнения в 2-4 раза.

FlashAttention-3 использует все новые функции Hopper, применяя мощные абстракции из библиотеки CUTLASS от NVIDIA. Переписывая FlashAttention для использования этих новых функций, мы уже можем значительно ускорить его (например, с 350 TFLOPS в FlashAttention-2 FP16 до около 540-570 TFLOPS). Однако асинхронная природа новых инструкций на Hopper (WGMMA и TMA) открывает дополнительные алгоритмические возможности для наложения операций и, таким образом, извлечения еще большей производительности.