Flash Attention

Ускорение обучения и инференса трансформеров на длинных последовательностях.

Метаданные навыка

Справочник: полный SKILL.md

Flash Attention — быстрое внимание с экономией памяти

Быстрый старт

Flash Attention даёт 2–4x ускорение и 10–20x сокращение памяти для механизма внимания трансформеров за счёт IO-осведомлённого тайлинга и перевычисления активаций.

Нативный PyTorch (проще всего, PyTorch 2.2+):

import torch
import torch.nn.functional as F

q = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)  # [batch, heads, seq, dim]
k = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)

# Автоматически использует Flash Attention, если доступно
out = F.scaled_dot_product_attention(q, k, v)

Библиотека flash-attn (больше возможностей):

pip install flash-attn --no-build-isolation

from flash_attn import flash_attn_func

# q, k, v: [batch, seqlen, nheads, headdim]
out = flash_attn_func(q, k, v, dropout_p=0.0, causal=True)

Типовые сценарии

Сценарий 1: включение в существующей модели PyTorch

Скопируйте этот чеклист:

Flash Attention Integration:
- [ ] Step 1: Check PyTorch version (≥2.2)
- [ ] Step 2: Enable Flash Attention backend
- [ ] Step 3: Verify speedup with profiling
- [ ] Step 4: Test accuracy matches baseline

Шаг 1: проверить версию PyTorch

python -c "import torch; print(torch.__version__)"
# Должно быть ≥2.2.0

Если <2.2, обновите:

pip install --upgrade torch

Шаг 2: подключить бэкенд Flash Attention

Замените стандартный механизм внимания:

# До (стандартное внимание)
attn_weights = torch.softmax(q @ k.transpose(-2, -1) / math.sqrt(d_k), dim=-1)
out = attn_weights @ v

# После (Flash Attention)
import torch.nn.functional as F
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)

Принудительное использование бэкенда Flash Attention (torch.backends.cuda.sdp_kernel устарел, берите torch.nn.attention.sdpa_kernel вместе с SDPBackend):

from torch.nn.attention import SDPBackend, sdpa_kernel

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v)

Шаг 3: замерить ускорение с профилированием

import torch.utils.benchmark as benchmark

def test_attention(use_flash):
    q, k, v = [torch.randn(2, 8, 2048, 64, device='cuda', dtype=torch.float16) for _ in range(3)]

    if use_flash:
        from torch.nn.attention import SDPBackend, sdpa_kernel
        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
            return F.scaled_dot_product_attention(q, k, v)
    else:
        attn = (q @ k.transpose(-2, -1) / 8.0).softmax(dim=-1)
        return attn @ v

# Бенчмарк
t_flash = benchmark.Timer(stmt='test_attention(True)', globals=globals())
t_standard = benchmark.Timer(stmt='test_attention(False)', globals=globals())

print(f"Flash: {t_flash.timeit(100).mean:.3f}s")
print(f"Standard: {t_standard.timeit(100).mean:.3f}s")

Ожидаемый результат: 2–4x ускорение для последовательностей >512 токенов.

Шаг 4: проверить совпадение результатов с базовым вариантом

# Сравниваем результаты
q, k, v = [torch.randn(1, 8, 512, 64, device='cuda', dtype=torch.float16) for _ in range(3)]

# Flash Attention
out_flash = F.scaled_dot_product_attention(q, k, v)

# Стандартное внимание
attn_weights = torch.softmax(q @ k.transpose(-2, -1) / 8.0, dim=-1)
out_standard = attn_weights @ v

# Проверяем разницу
diff = (out_flash - out_standard).abs().max()
print(f"Max difference: {diff:.6f}")
# Должно быть <1e-3 для float16

Сценарий 2: библиотека flash-attn для расширенных функций

Для multi-query attention, sliding window или H100 FP8.

Скопируйте этот чеклист:

flash-attn Library Setup:
- [ ] Step 1: Install flash-attn library
- [ ] Step 2: Modify attention code
- [ ] Step 3: Enable advanced features
- [ ] Step 4: Benchmark performance

Шаг 1: установить библиотеку flash-attn

# GPU NVIDIA (CUDA 12.0+)
pip install flash-attn --no-build-isolation

# Проверяем установку
python -c "from flash_attn import flash_attn_func; print('Success')"

Шаг 2: изменить код механизма внимания

from flash_attn import flash_attn_func

# Вход: [batch_size, seq_len, num_heads, head_dim]
# Транспонируем из [batch, heads, seq, dim], если нужно
q = q.transpose(1, 2)  # [batch, seq, heads, dim]
k = k.transpose(1, 2)
v = v.transpose(1, 2)

out = flash_attn_func(
    q, k, v,
    dropout_p=0.1,
    causal=True,  # Для авторегрессионных моделей
    window_size=(-1, -1),  # Без скользящего окна
    softmax_scale=None  # Автомасштабирование
)

out = out.transpose(1, 2)  # Обратно в [batch, heads, seq, dim]

Шаг 3: включить расширенные функции

Multi-query attention (общие K/V для всех голов):

from flash_attn import flash_attn_func

# q: [batch, seq, num_q_heads, dim]
# k, v: [batch, seq, num_kv_heads, dim]  # Меньше голов KV
out = flash_attn_func(q, k, v)  # Автоматически обрабатывает MQA

Sliding window attention (локальное внимание):

# Внимание только к окну 256 токенов до/после
out = flash_attn_func(
    q, k, v,
    window_size=(256, 256),  # окно (слева, справа)
    causal=True
)

Шаг 4: замерить производительность

import torch
from flash_attn import flash_attn_func
import time

q, k, v = [torch.randn(4, 4096, 32, 64, device='cuda', dtype=torch.float16) for _ in range(3)]

# Прогрев
for _ in range(10):
    _ = flash_attn_func(q, k, v)

# Бенчмарк
torch.cuda.synchronize()
start = time.time()
for _ in range(100):
    out = flash_attn_func(q, k, v)
    torch.cuda.synchronize()
end = time.time()

print(f"Time per iteration: {(end-start)/100*1000:.2f}ms")
print(f"Memory allocated: {torch.cuda.max_memory_allocated()/1e9:.2f}GB")

Сценарий 3: оптимизация H100 FP8 (FlashAttention-3)

Для максимальной производительности на GPU семейства Hopper (H100).

Важно: pip-пакет flash-attn (2.8.x) содержит только FlashAttention-2 — ни FA3, ни FP8-ядер для H100 в нём нет, и flash_attn_func FP8 сам по себе не включает. FlashAttention-3 — отдельная бета-сборка из исходников каталога hopper/ в репозитории, доступная через модуль flash_attn_interface. FA3 поддерживает FP16/BF16 в прямом и обратном проходе и FP8 только в прямом.

FP8 Setup:
- [ ] Step 1: Verify Hopper (H100) GPU available
- [ ] Step 2: Build & install FlashAttention-3 from source (hopper/)
- [ ] Step 3: Use the FA3 interface (FP8 forward)

Шаг 1: проверить наличие GPU H100

nvidia-smi --query-gpu=name --format=csv
# Должно показать «H100» или «H800»

Шаг 2: собрать и установить FlashAttention-3 из исходников

FA3 не входит в pip install flash-attn. Собирайте его из подкаталога hopper/:

git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install
# (компиляция тяжёлая, нужны CUDA toolchain и GPU семейства Hopper)

Шаг 3: использовать интерфейс FA3 (прямой проход на FP8)

FA3 предоставляет собственный модуль flash_attn_interface (это не FA2-модуль flash_attn). FP8 работает только в прямом проходе и ждёт входные данные типа float8_e4m3fn:

import torch
from flash_attn_interface import flash_attn_func  # FA3 (сборка hopper), не `flash_attn`

# q, k, v: [batch, seqlen, nheads, headdim]
q = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)
k = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)

# Прямой проход FP8 (инференс / только forward): приводим к float8_e4m3fn
q_fp8 = q.to(torch.float8_e4m3fn)
k_fp8 = k.to(torch.float8_e4m3fn)
v_fp8 = v.to(torch.float8_e4m3fn)

out = flash_attn_func(q_fp8, k_fp8, v_fp8, causal=True)
# Интерфейс FA3 также поддерживает FP16/BF16 в прямом и обратном проходе.

Когда использовать, а когда — альтернативы

Flash Attention подходит, если:

• Обучаете трансформеры с последовательностями >512 токенов
• Запускаете инференс с длинным контекстом (>2K токенов)
• Не хватает VRAM (OOM при стандартном attention)
• Нужно ускорение 2–4x без потери точности
• Используете PyTorch 2.2+ или готовы установить flash-attn

Лучше взять альтернативу, когда:

• Стандартное внимание: последовательности <256 токенов — накладные расходы не окупаются
• xFormers: нужны дополнительные варианты механизма внимания, а не просто скорость
• Memory-efficient attention: инференс на CPU (Flash Attention работает только на GPU)

Частые проблемы

Проблема: ImportError: cannot import flash_attn

Установите с флагом --no-build-isolation:

pip install flash-attn --no-build-isolation

Или предварительно установите CUDA toolkit:

conda install cuda -c nvidia
pip install flash-attn --no-build-isolation

Проблема: скорость ниже ожидаемой (нет ускорения)

Выигрыш от Flash Attention растёт с длиной последовательности:

• <512 токенов: минимальное ускорение (10–20%)
• 512–2K токенов: ускорение 2–3x

• 2K токенов: ускорение 3–4x

Убедитесь, что длина последовательности достаточна.

Проблема: RuntimeError: CUDA error

Проверьте, поддерживает ли GPU Flash Attention:

import torch
print(torch.cuda.get_device_capability())
# Должно быть ≥(7, 5) для Turing и новее

Требования по архитектурам:

• Ampere (A100, A10): ✅ Полная поддержка
• Turing (T4): ✅ Поддерживается
• Volta (V100): ❌ Не поддерживается

Проблема: деградация точности

Убедитесь, что тип данных — float16 или bfloat16, а не float32:

q = q.to(torch.float16)  # Или torch.bfloat16

Flash Attention работает с float16/bfloat16. Float32 не поддерживается.

Дополнительные темы

Интеграция с HuggingFace Transformers: инструкции по включению Flash Attention в моделях BERT, GPT, Llama — в references/transformers-integration.md.

Бенчмарки производительности: детальное сравнение скорости и памяти на разных GPU и длинах последовательностей — в references/benchmarks.md.

Требования к оборудованию

• GPU: NVIDIA Ampere+ (A100, A10, A30) или AMD MI200+
• VRAM: столько же, сколько для стандартного внимания (Flash Attention не увеличивает потребление памяти)
• CUDA: 12.0+ (минимум 11.8)
• PyTorch: 2.2+ для нативной поддержки

Не поддерживается: V100 (Volta), инференс на CPU

Ресурсы

• Статья: “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” (NeurIPS 2022)
• Статья: “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning” (ICLR 2024)
• Блог: https://tridao.me/blog/2024/flash3/
• GitHub: https://github.com/Dao-AILab/flash-attention
• Документация PyTorch: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html