FlashAttention v1

Хотя пиковая вычислительная мощность современных GPU измеряется триллионами FLOPS, на практике производительность на многих задачах глубокого обучения ограничивается скоростью доступа к памяти.

На GPU есть иерархия памяти: регистры и shared memory (SRAM) очень быстрые, но их объём крайне мал (менее 0.1% от объёма VRAM), глобальная память (VRAM) — объёмная, но значительно медленнее. Эффективность алгоритма на GPU определяется тем, насколько успешно он минимизирует обращения к VRAM, удерживая промежуточные данные в SRAM.

Во многих стандартных реализациях нейросетевых операций вычисления оказываются memory-bound, классический пример — attention в трансформерах.

Проблема базового attention

Scaled Dot-Product Attention вычисляется по формуле 1 (см. картинку). На практике это означает явное вычисление и хранение матрицы QKᵀ размером n x n.

При росте длины последовательности это приводит к квадратичному росту памяти и обращений к медленной VRAM-памяти. Attention остаётся memory-bound даже при использовании библиотеки cuBLAS. Масштабирование по длине контекста упирается не в вычисления, а в пропускную способность памяти.

Идея FlashAttention

FlashAttention ускоряет вычисление attention без изменения самой формулы расчёта. Ключевая идея — не хранить матрицу QKᵀ целиком в глобальной памяти.

Это достигается за счёт двух техник: тайлинга и онлайн-Softmax.

Тайлинг (Tiling)

Алгоритм разбивает все входные матрицы Q,K,V на блоки, которые поочередно загружаются в быструю SRAM. 

Каждый блок запросов Q по очереди сопоставляется с блоками K и V. Для каждого такого сочетания прямо в SRAM вычисляется фрагмент внимания, который сразу агрегируется в итоговый результат. Таким образом, не нужно ждать обхода всей последовательности или сохранять промежуточную матрицу QKᵀ. Финальный attention заполняется постепенно теми данными, которые в данный момент лежат в быстрой памяти.

Подробнее про математику tiling-механизма можно прочитать по ссылке.

Онлайн-Softmax

Тайлинг в flashattention был бы не так эффективен без онлайн-Softmax. В стандартной реализации для вычисления весов внимания используется формула 2 (см. картинку).

По формуле знаменатель представляет собой сумму экспонент по строке длиной N. В классической реализации attention это создает проблему: даже используя тайлинг вычисления, необходимы результаты всех блоков Q,K,V для расчёта Softmax. Это вынуждает снова возвращаться к проблеме сохранения матриц целиком в VRAM.

Онлайн-Softmax решает проблему, позволяя обновлять знаменатель и нормализовывать результат по мере подгрузки новых блоков в SRAM, вычисляя результат инкрементально. При обработке каждого нового блока в SRAM алгоритм «на лету» обновляет:

• текущий максимум строки;
• нормализующую сумму экспонент (знаменатель формулы);
• промежуточный результат O — если в новом блоке нашёлся максимум больше предыдущего, алгоритм пересчитывает уже накопленный результат.

Это позволяет получить результат, полностью эквивалентный классическому Softmax, при этом матрица QKᵀ вообще не покидает пределов SRAM.

Итог

FlashAttention показал, что существенное ускорение трансформеров возможно за счёт переосмысления вычислений под архитектуру GPU. Алгоритм снижает потребление памяти с квадратичной до линейной, уменьшает число обращений к медленной памяти VRAM и переводит вычисление attention из memory-bound в compute-bound проблему, когда скорость вычислений ограничена скоростью работы GPU.

Про ускорение LLM, дообучение и их деплой рассказываем на курсе LLM. Запишитесь в лист ожидания до 8 февраля, чтобы первыми занять место и получить скидки до 20% 🔥

Автор: Антон Наумов

🪔 DeepSchool