Что такое FlashAttention и как он ускоряет внимание

FlashAttention — это алгоритм точного внимания, который перестраивает порядок вычислений так, что блоки данных обрабатываются в быстрой памяти GPU. Ответ на вопрос «что такое FlashAttention и как он ускоряет внимание» короткий: он сокращает переносы данных между HBM и SRAM, поэтому GPU меньше ждёт память и больше считает. Математика внимания остаётся прежней.

Основной механизм простыми словами

Внимание сравнивает каждый токен с каждым другим. Modular Handbook описывает попарное сопоставление позиций в плотном внимании. При причинной маске будущие позиции закрыты; сама схема матрицы остаётся удобной для объяснения. Стандартная реализация умножает Q на Kᵀ и получает матрицу N×N, где N — длина последовательности. Для 4096 токенов эта матрица содержит около 16 миллионов элементов, и её приходится записывать в HBM, а затем читать обратно для следующей операции. Авторы FlashAttention в статье на arXiv отмечают, что само внимание имеет квадратичную по длине последовательности сложность по времени и памяти.

Ключевая цель этой оптимизации — сократить перемещения данных при сохранении вычисления плотного внимания. Modular Handbook прямо называет стандартное внимание «memory-bound»: алгоритм проводит большую часть времени в ожидании переносов памяти. Именно на этот слой и нацелен FlashAttention.

Механизм опирается на две техники, которые перечисляет Modular Handbook. Первая — тайлинг и пересчёт: вычисление разбивается на блоки, помещающиеся в SRAM, блоки Q, K, V загружаются из HBM в быструю память, внимание считается внутри SRAM, а выход обновляется инкрементально. Вторая — слияние ядер: вместо отдельных операций matmul → softmax → matmul всё сводится в одно ядро GPU, крупные промежуточные матрицы не материализуются целиком в HBM, а число отдельных запусков ядер сокращается.

Блог PyTorch описывает это так: тайлинг загружает блоки входов из HBM в SRAM, внимание считается относительно блока, выход обновляется в HBM. За счёт отказа от записи крупных промежуточных матриц внимание сокращает объём чтений и записей и даёт ускорение по настенным часам в 2–4 раза. Там же подчёркнуто, что результат получается корректным, без приближения.

Пределы механизма важно назвать сразу. Плотное внимание сохраняет O(N²d) операций с плавающей точкой. Экономия касается переносов памяти и объёма промежуточных матриц. Входные данные и KV-кэш остаются на месте. Формулировка «FlashAttention удаляет неважные токены» неверна для базового алгоритма: он не отбрасывает связи и не меняет смысл внимания.

Учебный вход: схема 8×8 и блоки 2×2

Разберём механизм на условном примере. Возьмём последовательность из 8 токенов. Матрица связей между ними — 8×8, то есть 64 элемента. Разобьём её на блоки 2×2: получится 16 блоков по 4 элемента. Это рисунок обхода по частям с накоплением результата и корректной нормализацией.

Число 16 здесь означает количество блоков в учебной сетке. Оно не равно числу реальных запусков ядра на GPU. Четыре элемента задают размер блока на рисунке. Расход всей памяти алгоритма и производительность требуют отдельной оценки.

Полный учебный текст схемы:

Учебная схема FlashAttention (условный вход, не измерение)
Последовательность: 8 токенов
Матрица связей: 8 x 8 = 64 элемента
Разбиение: блоки 2 x 2
Число блоков: 16
Элементов в блоке: 4

Полная сетка 8 x 8. Одинаковые буквы обозначают один блок 2 x 2:
A A | B B | C C | D D
A A | B B | C C | D D
----+-----+-----+----
E E | F F | G G | H H
E E | F F | G G | H H
----+-----+-----+----
I I | J J | K K | L L
I I | J J | K K | L L
----+-----+-----+----
M M | N N | O O | P P
M M | N N | O O | P P

Порядок обхода:
1. Загрузить блок Q, K, V в быструю память.
2. Посчитать внимание внутри блока.
3. Обновить выход с корректной нормализацией.
4. Перейти к следующему блоку.

Смысл: обход по частям с накоплением результата.
16 — число блоков сетки, не число запусков ядра.
4 — размер учебной ячейки, не объём памяти алгоритма.

Эталон проверки понимания — карточка утверждений. Ниже она приведена целиком как учебный вход.

Карточка утверждений (учебный вход)
1. «FlashAttention удаляет неважные токены» — неверно для базового алгоритма.
2. «Переорганизует точное внимание и сокращает перенос данных» — основной механизм.
3. «Любая GPU получит двукратное ускорение» — не обосновано.
4. «Сравнивать нужно одну модель, длину, точность и устройство» — корректно.

Пункт 4 задаёт правило сравнения. Ускорение зависит от модели, длины последовательности, точности вычислений и конкретного устройства. Замер на одной конфигурации переносится на другую только после проверки.

Структурная корректность и фактическая правильность

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

FlashAttention относится к точному вниманию. Блог PyTorch указывает, что результат получается корректным без приближения. Это утверждение о структуре алгоритма: он не заменяет softmax приближением и не отбрасывает часть связей.

Численная погрешность — отдельная тема. Блог PyTorch описывает переход к FP8: он удваивает пропускную способность тензорных ядер, приводя значения 989 TFLOPS для FP16 и 1978 TFLOPS для FP8, и обменивает это на точность, поскольку для представления чисел с плавающей точкой используется меньше битов. В том же материале сказано, что FlashAttention-3 с FP8 даёт ошибку в 2,6 раза меньше, чем базовое FP8-внимание. Это сравнение двух реализаций FP8 в условиях авторов. Численная погрешность сохраняется.

Практическая проверка сопоставляет выход с эталонной реализацией при явно заданном допуске на численное расхождение. Точное внимание означает сохранение математической операции. Истинность текста, который выдаёт языковая модель, из этого свойства не следует.

Версии, оборудование и первоисточники

Версии FlashAttention различаются тем, какие аппаратные возможности они используют. Блог PyTorch сообщает, что FlashAttention-2 достигает до 70% теоретического максимума FLOPS на GPU Ampere (A100) и не использует новые возможности Hopper. FlashAttention-3 задействует три особенности Hopper: WGMMA для тензорных ядер, TMA для ускорения переноса данных между глобальной и разделяемой памятью, а также FP8 для низкой точности. По данным того же блога, FlashAttention-3 работает в 1,5–2,0 раза быстрее FlashAttention-2 с FP16, достигает 740 TFLOPS и 75% использования теоретического максимума H100, а с FP8 доходит почти до 1,2 PFLOPS.

Для реализации FlashAttention-2 с CUDA репозиторий проекта перечисляет GPU Ampere, Ada или Hopper, например A100, RTX 3090, RTX 4090, H100. Для GPU Turing, таких как T4 и RTX 2080, существует отдельный репозиторий flash-attention-turing с подмножеством функций. Там же указано, что поддерживаются размерности голов до 256, а для обратного прохода при размерности выше 192 перечислены аппаратные ограничения. В том же README уточнено исключение: начиная с flash-attn 2.5.5 размерность 256 работает на потребительских GPU при отсутствии dropout.

Сборка описана в том же репозитории. Без ninja компиляция может занять около 2 часов, поскольку не используются несколько ядер CPU; с ninja на 64-ядерной машине с CUDA toolkit сборка занимает 3–5 минут. При объёме RAM меньше 96 ГБ и большом числе ядер ninja может запустить слишком много параллельных задач и исчерпать память, поэтому число задач ограничивают переменной окружения MAX_JOBS. Репозиторий сообщает о нескольких положительных отчётах о работе на Windows начиная с версии 2.3.2, но указывает, что сборка для Windows требует дополнительного тестирования.

Для AMD в репозитории проекта описаны два бэкенда ROCm: composable_kernel как бэкенд по умолчанию и Triton. Реализация на Triton поддерживает GPU CDNA (MI200, MI300) и RDNA с типами fp16, bf16 и fp32, прямой и обратный проходы с причинной маской, переменные длины последовательностей, MQA/GQA, dropout, ротационные эмбеддинги, ALiBi и paged attention. Скользящее окно внимания на момент описания находится в работе.

Как подключают в Transformers

Документация Hugging Face описывает слой AttentionInterface, который отделяет реализацию внимания от реализации модели. Все реализации внимания выполняют одно и то же вычисление, различается способ выполнения. Базовая реализация материализует полную матрицу внимания в памяти и создаёт узкие места, замедляющие вывод; оптимизированные реализации перестраивают вычисления, чтобы сократить трафик памяти.

Там же сказано, что библиотека Kernels загружает скомпилированные ядра прямо из Hub во время выполнения и это снимает проблемы упаковки при несовпадении версий PyTorch или CUDA. Ядра автоматически регистрируются в AttentionInterface при обнаружении, и устанавливать пакет FlashAttention явно не требуется. Запрос FlashAttention по имени также переходит на ядро из Hub.

Отдельный момент из той же документации: PyTorch SDPA сам выбирает самую быструю функцию внимания для CUDA-бэкендов. Для других бэкендов по умолчанию используется реализация на C++ из PyTorch.

Учебный конфиг для запуска модели с выбранной реализацией внимания:

import os
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = os.environ["MODEL_ID"]
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    attn_implementation="flash_attention_2",
    torch_dtype="auto",
    device_map="auto",
)
inputs = tokenizer("Проверка внимания", return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=16)
print(tokenizer.decode(out[0], skip_special_tokens=True))

Здесь MODEL_ID берётся из переменной окружения, чтобы не хранить идентификатор в коде. Значение attn_implementation подставляется под доступный бэкенд. Этот фрагмент — учебный вход, показывающий форму вызова; запуск на конкретном оборудовании требует проверки совместимости версии, типа данных и устройства.

Документация Hugging Face предупреждает о ловушке при регистрации собственной функции внимания. Если имя attn_implementation не зарегистрировано в AttentionMaskInterface, Transformers пропускает создание маски и передаёт в слои внимания attention_mask=None. Тогда причинные, паддинговые, упаковочные и скользящие ограничения должен обрабатывать сам код функции, иначе они могут быть молча потеряны.

Случай, где подход не решает задачу

Есть конфигурации, в которых FlashAttention не даёт выигрыша. Первый случай — сравнение с методом, который меняет само вычисление внимания. В экспериментах исходной статьи на arXiv при длине выше 1K некоторые приближённые методы, например Linformer, оказывались быстрее. Такое сравнение требует отдельно оценивать качество: FlashAttention сохраняет точное плотное внимание, а приближённый метод решает иную вычислительную задачу. Результат конкретного эксперимента не означает преимущества любой реализации при любой длине.

Второй случай — несовместимая точность выбранного бэкенда. В примере Hugging Face некоторые визуальные части модели лучше работают в fp32, поэтому для них предлагается другой backend внимания. Поддержку точности нужно проверять для конкретной реализации: официальный README отдельно указывает fp32 для AMD Triton backend. Для мультимодальных моделей там же предложено назначать разные функции внимания разным бэкендам через словарь, чтобы текстовый бэкенд продолжал использовать FlashAttention.

Третий случай — оборудование вне списка поддержки выбранной реализации. Для Turing репозиторий предлагает отдельную ветку с подмножеством функций, а для Windows сообщает лишь о положительных отчётах, требующих дальнейшего тестирования. Попытка собрать основную версию на неподходящем устройстве приведёт к ошибке компиляции или отсутствию нужного ядра.

Четвёртый случай — задача, где узкое место лежит вне внимания. FlashAttention оптимизирует уровень ядра. Если время уходит на другие операции модели или на сетевой обмен, экономия переносов памяти внутри внимания даст ограниченный эффект.

Общее правило: ускорение проверяется на той конфигурации, где модель будет работать. Сравнение проводится при одной модели, одной длине последовательности, одной точности и одном устройстве.

Частые вопросы

Что такое FlashAttention простыми словами?
Это алгоритм точного внимания, который обрабатывает данные блоками в быстрой памяти GPU и сокращает переносы между HBM и SRAM. Математика внимания сохраняется, меняется порядок вычислений и движения данных. За счёт этого GPU меньше простаивает в ожидании памяти.

Чем FlashAttention-2 отличается от первой версии?
Версия 2 достигает до 70% теоретического максимума FLOPS на GPU Ampere (A100), как сообщает блог PyTorch. Она не задействует новые возможности Hopper, для которых позже сделали версию 3. Требования к оборудованию и поддерживаемые размерности голов описаны в репозитории проекта.

Что известно про FlashAttention-3 и FlashAttention-4?
Версия 3 использует WGMMA, TMA и FP8 на Hopper: блог PyTorch приводит ускорение1,5–2,0 раза относительно версии2 с FP16 в своих условиях. Сохранённый официальный README также содержит FlashAttention-4: реализация написана на CuTeDSL и оптимизирована для Hopper и Blackwell, например H100 и B200. Требования и способ подключения нужно брать из раздела соответствующей версии; описанные выше результаты версии3 не становятся результатами версии4.

Ускоряет ли FlashAttention вывод на любой видеокарте?
Нет. Для реализации FlashAttention-2 с CUDA репозиторий перечисляет GPU Ampere, Ada и Hopper, а для Turing предлагает отдельную ветку с подмножеством функций. Ускорение зависит от модели, длины последовательности, точности и устройства. Требования других версий и бэкендов проверяют в их собственных разделах.

Убирает ли FlashAttention квадратичную сложность внимания?
Математическая сложность плотного внимания остаётся квадратичной по длине последовательности: авторы статьи на arXiv описывают её как O(N²d) операций. Экономия достигается на переносах памяти и промежуточных матрицах. Отдельно авторы показали, что блочно-разреженный вариант масштабируется до 64k токенов и работает в 2–4 раза быстрее базового FlashAttention.