Как работает FlashAttention-3 технически? Чем отличается от FA2?

Краткий тезис

FlashAttention-3 (FA3) — это третья версия алгоритма быстрого внимания, оптимизированная для архитектуры Hopper (H100 и новее). Ключевые отличия от Attention 2|FlashAttention-2 (FA2): использование новых аппаратных возможностей — инструкций WGMMA (Warp Group Matrix Multiply-Accumulate) для асинхронного матричного умножения, TMA (Tensor Memory Accelerator) для асинхронного перемещения данных, увеличение размера блока с 64×64 до 128×128 и поддержка FP8. В результате FA3 показывает ускорение до 2× по сравнению с FA2 на H100 и до 4× по сравнению со стандартным attention|attention|dot-product attention|attention|dot-product attention|Scaled attention|dot-product attention|Dot-Product Attention (SDPA) при сохранении высокой точности.

1. Введение: эволюция FlashAttention

FlashAttention — это семейство алгоритмов, которые вычисляют внимание (attention) без материализации полной матрицы S (QK^T) в глобальной памяти. Вместо этого используется tiling (разбиение на блоки) и пересчёт softmax on-the-fly, что снижает потребление памяти с O(N^2) до O(N) и ускоряет вычисления за счёт лучшего использования кэша.

• FA1 (2022): базовая версия с tiling, поддержка Ampere (A100).
• FA2 (2023): улучшения в параллелизации по длине последовательности, оптимизации для A100, увеличенный размер блока (64×64).
• FA3 (2025): создана для архитектуры Hopper (H100, H200, B100). Использует новые аппаратные возможности, асинхронность и FP8.

2. Аппаратная база: архитектура Hopper (H100)

Architecture Hopper — это поколение GPU от NVIDIA (2023+), которое принесло ключевые нововведения:

3. WGMMA (Warp Group Matrix Multiply-Accumulate)

WGMMA — это инструкция в Hopper, которая выполняет матричное умножение-накопление сразу для группы варпов (warp group). В отличие от обычной инструкции HMMA (используемой в FA2), WGMMA работает асинхронно: инициирует GEMM и не блокирует warp до завершения.

• Синхронный GEMM (как в FA2): warp ждёт завершения умножения, пока все данные не будут готовы. Простой конвейера.
• Асинхронный GEMM (FA3): warp может запустить другой GEMM или обработать softmax, пока текущий считается на тензорных ядрах.

WGMMA работает с подгруппой из 4-х варпов (128 нитей) и оперирует с блоками размера 128×128, что соответствует размеру tile в FA3.

Пример псевдокода (упрощённо):

# Псевдокод для одного блока FA3
async_load(Q_block, TMA)   # асинхронная загрузка Q
async_load(K_block, TMA)   # асинхронная загрузка K
wgmma_result = wgmma(Q_block, K_block)  # асинхронное умножение
# Пока WGMMA выполняется, можно делать softmax или загружать V
wait_all()  # дождаться завершения

4. TMA (Tensor Memory Accelerator)

TMA — это отдельный аппаратный блок в Hopper, который выполняет асинхронное копирование данных из глобальной памяти (HBM) в разделяемую память (SRAM) и обратно.

• В FA2 копирование делалось программно через инструкцию cp.async. TMA автоматически разбивает данные на 128-байтовые сегменты, использует более широкую шину и поддерживает многомерные копии (2D-срезы тензоров).
• TMA может работать параллельно с тензорными ядрами, что позволяет перекрывать загрузку данных с вычислением GEMM.

Влияние на производительность уменьшение простоя конвейера (stall) при загрузке из HBM, которая является узким местом.

5. Изменение partitioning: блоки 128×128

В FA2 использовались блоки 64×64 для Q, K, V и выходных данных O. FA3 увеличивает размер блока до 128×128.

• Почему 128×128

  • Соответствует размеру, при котором WGMMA работает наиболее эффективно (128 нитей в warp group, каждая обрабатывает 4 элемента).
  • Увеличивает долю вычислений на каждый загруженный элемент, улучшая arithmetic intensity.
  • Снижает количество итераций по блокам для длинных последовательностей.

• Ограничение требует больше разделяемой памяти (128×128×2 байта на тензор при FP16 ≈ 32 КБ на один tile, а для Q,K,V,O — ~128 КБ). H100 предоставляет до 228 КБ, что достаточно.

6. Поддержка FP8

FP8 — 8-битный формат с плавающей точкой (e4m3 или e5m2). FA3 поддерживает FP8 для вычислений внимания.

• Преимущества

  • Увеличение пропускной способности тензорных ядер (в 2× больше операций в секунду по сравнению с FP16).
  • Снижение объёма передаваемых данных (2 байта на элемент → 1 байт).

• Технические детали

  • Используется смешанная точность: FP8 для Q, K, V и матрицы S, FP16/FP32 для softmax и накопления (для сохранения точности).
  • Накопление в WGMMA выполняется в FP16 или FP32, затем результат конвертируется обратно в FP8 при записи.

7. Сравнение FA2 и FA3

8. Технические детали реализации FA3 (шаг за шагом)

Предположим, мы вычисляем внимание для последовательности длиной N, с размерностью головы d (обычно 128 или 256). FA3 разбивает Q, K, V на блоки (tiles) размера B×d, где B = 128.

Алгоритм (упрощённо):

1. Инициализация выходного тензора O (N×d) нулями.
2. Внешний цикл по блокам K и V (по оси длины).
3. Внутренний цикл по блокам Q.
4. Для каждого блока Q (Q_block):
   • Асинхронно загрузить Q_block в разделяемую память через TMA.
   • Асинхронно загрузить K_block.
   • Запустить WGMMA для Q_block и K_block → получаем блок S_block (размер B×B) в регистрах.
   • Пока WGMMA выполняется, загрузить V_block асинхронно.
   • Дождаться завершения WGMMA.
   • Применить маску (если есть) и выполнить softmax с on-the-fly пересчётом (как в FA1/FA2).
   • Накопить результат: O_block += softmax(S_block) @ V_block (снова WGMMA).
   • Записать O_block обратно, используя TMA.

Ключевое различие с FA2: все копирования и умножения происходят асинхронно и перекрываются. В FA2 копирование и GEMM синхронизировались на каждом шаге.

9. Производительность и бенчмарки

На H100 (80GB SXM) FA3 демонстрирует:

Ускорение vs FA2 на H100 достигает 2× для FP16 и 4× для FP8, потому что FP8 удваивает throughput тензорных ядер.

10. Ограничения и требования

• Аппаратные требования только GPU архитектуры Hopper (H100, H200, B100). На Ampere (A100) FA3 не запустится из-за отсутствия WGMMA и TMA.
• Ограничение по длине последовательности из-за размера разделяемой памяти FA3 может не поддерживать блоки больше 128 в некоторых конфигурациях (если d велико).
• Точность: при использовании FP8 возможна незначительная потеря качества (несколько сотых perplexity), что может быть критично для некоторых задач.
• Совместимость FA3 пока не интегрирован во все фреймворки (например, Hugging Face может не иметь оптимизированного ядра).

11. Применение в контексте Agentic RAG

В Agentic RAG агенты часто обрабатывают длинные контексты (история диалога, retrieved documents). Высокая производительность внимания критична для минимизации задержки (latency) и увеличения пропускной способности.

• FA3 позволяет обрабатывать контексты длиной 128K токенов на H100 с задержкой менее 100 мс, что делает возможным реальное использование агентов с full-context retrieval.
• Поддержка FP8 снижает требования к памяти, что важно при работе с несколькими агентами параллельно.
• Асинхронность может быть использована для перекрытия retrieval и вычисления внимания (pipeline parallelism).

Пет-проект для закрепления

Задача реализовать симуляцию асинхронного выполнения FA3 на Python (без GPU) для понимания конвейерной архитектуры.
Инструменты Python, NumPy, библиотека asyncio или concurrent.futures.
Шаги:

1. Написать класс FlashAttention3Simulator, который принимает матрицы Q, K, V, размер блока B=128.
2. Разбить последовательность на блоки.
3. Для каждого блока создать корутины async_load (имитация TMA) и async_gemm (имитация WGMMA) с использованием asyncio.sleep для симуляции задержки.
4. Организовать конвейер: загрузить Q, K, одновременно запустить GEMM и загрузить V для следующего блока.
5. Измерить общее время выполнения и сравнить с синхронной версией (FA2).
   Ожидаемый результат программа покажет снижение времени выполнения за счёт перекрытия загрузки и вычислений; график Gantt.

Связь с другими вопросами

Навигация

• Предыдущий: 431
• Следующий: 433
• Индекс: 00. Индекс разборов