Flash Attention

Flash Attention (FlashAttention) — это алгоритм и семейство ядер GPU, которые вычисляют точное внимание в моделях трансформеров, не записывая полную матрицу внимания в основную память GPU. При стандартном самовнимании по N токенам строится матрица оценок N × N, которая записывается в память с высокой пропускной способностью (HBM), а затем считывается обратно для вычисления softmax и произведения на значения, поэтому большая часть времени выполнения внимания уходит на перемещение данных. Три Дао и его коллеги из Стэнфорда представили этот метод в статье 2022 года «FlashAttention: быстрое и эффективное по памяти точное внимание с учётом ввода-вывода». Вычисления разбиваются на блоки, помещающиеся в быструю встроенную память SRAM, поэтому потребление памяти растёт линейно с длиной последовательности, а выход совпадает с результатом стандартного внимания с точностью до округления с плавающей точкой. Ядра работают внутри PyTorch и в движках инференса больших языковых моделей (LLM), например vLLM.

Принцип работы Flash Attention#

У GPU есть небольшой объём очень быстрой памяти рядом с вычислительными блоками и большой, но более медленной памяти вне чипа. В статье о FlashAttention описывается NVIDIA A100 с 40–80 ГБ HBM, пропускной способностью 1.5–2.0 ТБ/с и 192 КБ SRAM в каждом из 108 потоковых мультипроцессоров, для которой оценочная пропускная способность составляет 19 ТБ/с. Внимание выполняет мало операций на каждый прочитанный байт, поэтому его скорость зависит от частоты перемещения данных между двумя уровнями памяти. Метод учитывает число таких перемещений — именно это означает термин «с учётом ввода-вывода».

FlashAttention загружает блок запросов, а затем последовательно пропускает блоки ключей и значений через SRAM. Для каждого блока алгоритм вычисляет оценки, обновляет текущие максимум и сумму для каждой строки запросов, пересчитывает масштаб частичного результата при увеличении максимума и добавляет вклад нового блока. Такой онлайн-softmax даёт точный результат softmax, не обрабатывая строку целиком за один раз; в HBM возвращаются только итоговый результат и два числа на строку. При обратном проходе оценки блоков пересчитываются по Q, K и V вместо хранения матрицы N × N. Это требует дополнительных операций, но всё равно выполняется быстрее, поскольку считывается гораздо меньше данных из памяти. Та же схема на PyTorch, где все запросы обрабатываются по одному блоку ключей и значений за раз:

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

Самый большой блок оценок имеет размер 4,096 × 256 — это одна шестнадцатая матрицы 4,096 × 4,096. Результат отличается от стандартного внимания не более чем на 4.3e-07 — величину ошибки округления FP32. В настоящем ядре запросы также разбиваются на блоки, по одному на каждый блок потоков GPU, а весь цикл объединяется в одно ядро, чтобы блоки не покидали SRAM.

FlashAttention-2, -3 и -4#

Каждая версия рассчитана на более новое поколение GPU. FlashAttention-2 (2023) изменила способ распределения вычислительной работы между потоками и работала примерно вдвое быстрее исходной версии, достигая 50–73% пиковой пропускной способности A100. FlashAttention-3 (2024) использует асинхронные Tensor Cores H100 и поддерживает FP8, обеспечивая 740 TFLOPs/s в FP16 и почти 1.2 PFLOPs/s в FP8. FlashAttention-4 (2026) предназначена для GPU Blackwell и, согласно опубликованным данным, обеспечивает до 1 613 TFLOPs/s в BF16 на B200, что составляет 71% от пиковой производительности. Ядра входят в библиотеку с открытым исходным кодом flash-attention, которая работает с входными данными FP16 и BF16, а функция PyTorch scaled_dot_product_attention автоматически выбирает ядро FlashAttention-2, если GPU и входные данные его поддерживают, и в противном случае переключается на реализацию с эффективным использованием памяти или обычную реализацию. Модель SAM 2 от Meta, которая работает в пакете Ultralytics, вычисляет внимание в кодировщике изображений Hiera и модуле внимания к памяти с помощью этой функции, поэтому PyTorch направляет эти слои в FlashAttention, когда это позволяют оборудование и точность входных данных.

Flash Attention и смежные понятия#

• Стандартное внимание: Та же математическая операция, но полные матрицы оценок и вероятностей записываются в HBM.
• Разреженное внимание: Вычисляет оценки только для выбранных пар запросов и ключей. В статье о FlashAttention также представлена блочно-разреженная версия, пропускающая пустые блоки.
• Линейное внимание: Заменяет softmax картой признаков, чтобы вычислительная стоимость росла линейно с длиной последовательности; при этом меняются выходы модели.
• PagedAttention: Менеджер памяти, который при обслуживании запросов хранит KV-кэш в страницах фиксированного размера. Ядра внимания, включая FlashAttention, считывают ключи и значения из этих страниц.
• Кольцевое внимание: Разделяет очень длинную последовательность между устройствами. Каждое устройство вычисляет блочное внимание для своего фрагмента, а блоки ключей и значений передаются по кольцу.

Влияние на обучение и длинные контексты#

Благодаря линейному росту объёма памяти с длиной последовательности FlashAttention сделал практичным обучение на длинных контекстных окнах. В оригинальной статье сообщалось об ускорении полного цикла обучения BERT-large на 15% при длине последовательности 512 относительно рекорда обучения MLPerf 1.1, об ускорении GPT-2 в 3 раза при длине последовательности 1,024 и в 2.4 раза на тесте Long Range Arena. Такая же экономия ресурсов полезна и в компьютерном зрении, где Vision Transformer (ViT) или диффузионный трансформер преобразует изображения и видео высокого разрешения в длинные последовательности токенов. Более длинные последовательности также повысили точность: в статье описаны первые трансформеры, превзошедшие случайный выбор на Path-X — задаче, в которой нужно определить, соединены ли путём две точки на изображении размером 128 × 128, преобразованном в последовательность из 16,384 пикселей.