optimizing-attention-flash

Flash Attention - 高速でメモリ効率の良いアテンション

クイックスタート

Flash Attention は、IO を意識したタイリングと再計算によって、トランスフォーマーのアテンションを 2～4 倍高速化し、メモリ使用量を 10～20 分の 1 に削減します。

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)

# Automatically uses Flash Attention if available
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__)"
# Should be ≥2.2.0

2.2 未満の場合は、アップグレードしてください:

pip install --upgrade torch

ステップ 2: Flash Attention バックエンドを有効化

標準のアテンションを置き換えます:

# Before (standard attention)
attn_weights = torch.softmax(q @ k.transpose(-2, -1) / math.sqrt(d_k), dim=-1)
out = attn_weights @ v

# After (Flash Attention)
import torch.nn.functional as F
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)

Flash Attention バックエンドを強制的に使用します:

with torch.backends.cuda.sdp_kernel(
    enable_flash=True,
    enable_math=False,
    enable_mem_efficient=False
):
    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:
        with torch.backends.cuda.sdp_kernel(enable_flash=True):
            return F.scaled_dot_product_attention(q, k, v)
    else:
        attn = (q @ k.transpose(-2, -1) / 8.0).softmax(dim=-1)
        return attn @ v

# Benchmark
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")

想定される結果: 512 トークンを超えるシーケンスで 2～4 倍の高速化。

ステップ 4: 精度がベースラインと一致することをテスト

# Compare outputs
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)

# Standard attention
attn_weights = torch.softmax(q @ k.transpose(-2, -1) / 8.0, dim=-1)
out_standard = attn_weights @ v

# Check difference
diff = (out_flash - out_standard).abs().max()
print(f"Max difference: {diff:.6f}")
# Should be <1e-3 for float16

ワークフロー 2: 高度な機能のために flash-attn ライブラリを使用

Multi-query アテンション、スライディングウィンドウ、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 ライブラリをインストール

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

# Verify installation
python -c "from flash_attn import flash_attn_func; print('Success')"

ステップ 2: アテンションのコードを変更

from flash_attn import flash_attn_func

# Input: [batch_size, seq_len, num_heads, head_dim]
# Transpose from [batch, heads, seq, dim] if needed
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,  # For autoregressive models
    window_size=(-1, -1),  # No sliding window
    softmax_scale=None  # Auto-scale
)

out = out.transpose(1, 2)  # Back to [batch, heads, seq, dim]

ステップ 3: 高度な機能を有効化

Multi-query アテンション（ヘッド間で K/V を共有）:

from flash_attn import flash_attn_func

# q: [batch, seq, num_q_heads, dim]
# k, v: [batch, seq, num_kv_heads, dim]  # Fewer KV heads
out = flash_attn_func(q, k, v)  # Automatically handles MQA

スライディングウィンドウアテンション（ローカルアテンション）:

# Only attend to window of 256 tokens before/after
out = flash_attn_func(
    q, k, v,
    window_size=(256, 256),  # (left, right) window
    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)]

# Warmup
for _ in range(10):
    _ = flash_attn_func(q, k, v)

# Benchmark
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）

H100 GPU で最大限のパフォーマンスを得るためのワークフローです。

FP8 Setup:
- [ ] Step 1: Verify H100 GPU available
- [ ] Step 2: Install flash-attn with FP8 support
- [ ] Step 3: Convert inputs to FP8
- [ ] Step 4: Run with FP8 attention

ステップ 1: H100 GPU を確認

nvidia-smi --query-gpu=name --format=csv
# Should show "H100" or "H800"

ステップ 2: FP8 サポート付きで flash-attn をインストール

pip install flash-attn --no-build-isolation
# FP8 support included for H100

ステップ 3: 入力を FP8 に変換

import torch

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)

# Convert to float8_e4m3 (FP8)
q_fp8 = q.to(torch.float8_e4m3fn)
k_fp8 = k.to(torch.float8_e4m3fn)
v_fp8 = v.to(torch.float8_e4m3fn)

ステップ 4: FP8 アテンションで実行

from flash_attn import flash_attn_func

# FlashAttention-3 automatically uses FP8 kernels on H100
out = flash_attn_func(q_fp8, k_fp8, v_fp8)
# Result: ~1.2 PFLOPS, 1.5-2x faster than FP16

使うべき場合と代替手段

Flash Attention を使うべき場合:

• 512 トークンを超えるシーケンスでトランスフォーマーを学習する
• 長いコンテキスト（2K トークン超）で推論を実行する
• GPU メモリが不足している（標準のアテンションでは OOM になる）
• 精度を損なわずに 2～4 倍の高速化が必要
• PyTorch 2.2 以降を使用している、または flash-attn をインストールできる

代わりに代替手段を使うべき場合:

• 標準のアテンション: 256 トークン未満のシーケンス（オーバーヘッドに見合わない）
• xFormers: 速度だけでなく、より多くのアテンションのバリエーションが必要
• メモリ効率の良いアテンション: 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～3 倍の高速化
• 2K トークン超: 3～4 倍の高速化

シーケンス長が十分かどうかを確認してください。

問題: RuntimeError: CUDA error

GPU が Flash Attention をサポートしているかを確認します:

import torch
print(torch.cuda.get_device_capability())
# Should be ≥(7, 5) for Turing+

Flash Attention の要件:

• Ampere（A100、A10）: ✅ 完全サポート
• Turing（T4）: ✅ サポート
• Volta（V100）: ❌ 非サポート

問題: 精度の低下

dtype が float16 または bfloat16 であることを確認します（float32 ではない）:

q = q.to(torch.float16)  # Or torch.bfloat16

Flash Attention は高速化のために float16/bfloat16 を使用します。float32 はサポートされていません。

応用トピック

HuggingFace Transformers との統合: BERT、GPT、Llama モデルで Flash Attention を有効にする方法については、references/transformers-integration.md を参照してください。

パフォーマンスベンチマーク: GPU とシーケンス長ごとの詳細な速度とメモリの比較については、references/benchmarks.md を参照してください。

アルゴリズムの詳細: タイリング戦略、再計算、IO 計算量の分析については、references/algorithm.md を参照してください。

高度な機能: Rotary embeddings、ALiBi、paged KV cache、カスタムアテンションマスクについては、references/advanced-features.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