FlashAttention

Главное ограничение — обмен с памятью

Внимание для последовательности длины $N$ строит матрицу оценок $Q{K}^{\mathrm{\top }}$ размера $N\times N$, применяет к ней softmax и умножает результат на $V$. Число операций растёт как $N^2$, но вычислительная сложность — не единственная проблема. Квадратную матрицу нужно записать в память и прочитать обратно.

У GPU есть небольшая быстрая память рядом с вычислителями (SRAM) и большая, но более медленная память (HBM). Прямая реализация многократно переносит матрицу $N\times N$ между вычислителями и HBM: записывает после $Q{K}^{\mathrm{\top }}$, читает для softmax, снова записывает. При $N=8000$ это 64 миллиона чисел для каждой головы в каждом слое. Скорость начинает определяться обменом с памятью: внимание становится memory-bound (ограниченным пропускной способностью памяти).

Увеличьте длину с 2K до 8K — в четыре раза. Во сколько раз должны вырасти число оценок и память матрицы? Проверьте ответ, затем повторите рассуждение для следующего доступного размера.

Вычисление softmax по блокам

Обычный softmax вычитает максимум (ради устойчивости) и делит на сумму экспонент — для этого будто бы нужно увидеть всю строку. Но максимум и сумму можно накапливать потоково, блок за блоком. Достаточно поддерживать два скаляра и один вектор:

• $m$ — бегущий максимум уже просмотренных оценок;
• $\ell$ — бегущая сумма$\sum e^{s-m}$;
• $o$ — бегущий выход$\sum e^{s-m}v$ (ненормированный).

Если в новом блоке находится более высокий максимум, прежние накопления нужно привести к новому масштабу. Для этого достаточно одного множителя $\alpha =e^{m_{\text{old}}-m_{\text{new}}}$:

Итоговый результат равен $o/\ell$; полная строка весов при этом не хранится.

Блочная обработка

Разбиение на тайлы (tiling) позволяет одному запросу получать ключи блоками. При увеличении текущего максимума множитель $\alpha$ перенормирует накопленные ранее величины. Размер тайла и запрос можно выбрать до начала опыта; их смена начинает обработку заново.

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

Обратите внимание: в этой демонстрации в памяти в каждый момент находятся лишь текущий тайл, два скаляра и вектор-аккумулятор — но не матрица $N\times N$. Рабочие данные помещаются в быстрой SRAM, а обмен с HBM резко сокращается.

Путь данных

Теперь сопоставим два пути. Прямая реализация сохраняет квадратную матрицу в HBM, затем читает её для softmax и следующего умножения. FlashAttention переносит в SRAM только текущие блоки $Q_i,K_j,V_j$, вычисляет тайл $S=Q_i{K}_j^{\mathrm{\top }}$ и обновляет аккумулятор строки.

Пройдите обработку от первого до последнего тайла. Перед каждым шагом различайте уже выполненную работу и данные, которые ещё нужно хранить. Должны ли эти два числа расти одинаково?

В этом и состоит экономия: прямая реализация переносит через HBM всю квадратную матрицу $S$, а FlashAttention держит в SRAM лишь текущий тайл и записывает обратно только результат. Меньше обмена с медленной памятью — выше скорость и нет квадратного расхода памяти на матрицу внимания.

Как менялось следующее ограничение

Базовая математика внимания оставалась прежней. Каждое поколение устраняло одно ограничение и тем самым делало заметнее следующее: сначала обмен с HBM, затем загрузку GPU, наконец — паузы между softmax и матричными умножениями.

Сравните версии по порядку, начиная с FA-1. Для каждого перехода назовите устранённое ограничение и оставшуюся проблему. Проверьте также оборудование: можно ли напрямую сравнивать показанные ускорения между разными устройствами?

FA-1 (2022) — основа алгоритма

Первая версия объединяет всё, что мы разобрали в главах 1–4: тайлы, потоковый softmax и пересчёт весов на обратном проходе. Она убрала матрицу $N\times N$ из медленной памяти и резко сократила обмен с HBM. Но GPU всё ещё использовался не полностью. Работу делили только между парами «последовательность × голова», а внутри блока вычислители часто ждали друг друга. Отсюда — второе поколение.

FA-2 (2023) — больше параллелизма, меньше операций

Все три приёма направлены на то, чтобы GPU меньше простаивал. Матричные умножения выполняют Tensor Cores, а экспоненты, деления и поиск максимума выполняются на обычных CUDA-ядрах. Чем меньше вспомогательной работы и ожидания, тем полнее используется быстрая матричная часть GPU.

• Отложить деление. FA-1 нормировал выход — делил на бегущую сумму $\ell$ — на каждом блоке. В FA-2 деление выполняется один раз в самом конце. Стоимость одного деления невелика, но во внутреннем цикле эта операция повторяется множество раз. Её перенос сокращает долю вспомогательных вычислений.
• Больше параллельных блоков. FA-1 распределял работу по парам «последовательность × голова». Но при длинном контексте последовательностей мало, а сами они велики — независимых задач не хватает, чтобы загрузить все вычислительные блоки. FA-2 дополнительно делит работу по длине запроса: даже одна длинная последовательность даёт достаточно независимых задач для загрузки GPU.
• Каждому варпу — свои строки. Внутри блока работу делят варпы (группы потоков). FA-1 делил между ними ключи $K$: варпы вычисляли части одного результата, а затем объединяли их через разделяемую память и ждали друг друга. FA-2 делит запросы $Q$ — теперь у каждого варпа свои строки выхода целиком, поэтому им реже приходится объединять результаты и ждать друг друга.

FA-3 (2024) — совмещение этапов вычисления

Даже в FA-2 шаги идут по очереди: матричное умножение (Tensor Cores) → softmax (обычные ядра) → следующее матричное умножение. Пока обычные ядра считают экспоненты одного тайла, Tensor Cores ждут — и наоборот. На GPU архитектуры Hopper FA-3 выполняет эти этапы внахлёст: пока для тайла $N$ считается softmax, для тайла $N+1$ уже идёт матричное умножение. Асинхронные инструкции позволяют одновременно считать и загружать данные, а специализация варпов разделяет эти обязанности между группами потоков. Получается конвейер, в котором ожидание сведено к минимуму.

Второе нововведение — FP8, восьмибитный числовой формат. Такие матричные умножения быстрее, но менее точны, поэтому значения квантуют поблочно, контролируя численную ошибку. В этом режиме производительность на H100 достигает примерно 1,2 PFLOPS — уже с компромиссом между скоростью и численной точностью.

Эта идея получила дальнейшее развитие. Например, Flash-Decoding ускоряет генерацию с большим контекстом, распараллеливая работу по ключам. Основной принцип остаётся прежним: не хранить то, что можно пересчитать, и сокращать простои GPU.