optimizing-attention-flash

Flash Attention - 高速でメモリ効率に優れたアテンション

クイックスタート

Flash Attention は、IO を考慮したタイル化と再計算により、Transformer のアテンションを 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 ライブラリを使用する

マルチクエリアテンション、スライディングウィンドウ、または 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：高度な機能を有効にする

マルチクエリアテンション（ヘッド間で 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 トークンを超えるシーケンスで Transformer を学習する場合
• 長いコンテキスト（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 を参照してください。

高度な機能：ロータリー埋め込み、ALiBi、ページングされた KV キャッシュ、カスタムアテンションマスクについては、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

エージェントの作業手順

1. 環境を確認します。 GPU の種類、メモリ、ドライバー／CUDA のバージョン（nvidia-smi）、フレームワークのバージョン、データセットの場所とサイズを確認します。
2. 入力を明確にします。 データまたはユーザーに確認して、形式、識別子、パラメーターを確定します。結果に影響する値を推測せず、質問してください。
3. まず小規模に実行します。 スモークテストとして、小さなモデルまたはサブセットを使い、少ないステップ数で実行し、損失が減少することとチェックポイントが保存されることを確認します。
4. 上記の手順と参考資料を使って、タスク全体を実行します。
5. 結果を検証します。 保留データのメトリクスを追跡し、ベースラインと比較して、シード、設定、ハードウェアを記録します。
6. 報告します。 実行した内容（バージョン、コマンド、パラメーター）、確認した内容、未確定の点を伝えます。

発生した場合	対応
CUDA のメモリ不足	バッチサイズを小さくする、勾配累積／チェックポイントまたは混合精度を有効にする、あるいはモデルをシャーディングします。
損失が NaN になる、または発散する	学習率を下げ、データに無効な値がないか確認し、勾配クリッピングを有効にします。
これらの手順に記載された関数、フラグ、エンドポイントがインストール済みのバージョンにない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認して適応し、ユーザーに伝えます。API を捏造してはいけません。
必須の入力、識別子、パラメーターが曖昧	ユーザーに質問するか、実行前に仮定を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計値を決して捏造しないでください。実行または検証できない場合は、その旨を明確に伝えてください。
• ログなしにトレーニング結果を主張してはいけません。大規模なジョブを開始する前に、計算コストを見積もってください。
• ここに記載されているバージョン固有の詳細は古くなっている可能性があるため、インストール済みバージョンの公式ドキュメントで確認してください。
• 費用が発生する操作、共有 GPU やクラウドクォータを消費する操作、個人データや患者データを扱う操作、または元に戻せない操作を行う前には、許可を得てください。

関連スキル

• rwkv-architecture：RWKV（Receptance Weighted Key Value）アーキテクチャについて説明します。RWKV は O(n) の推論が可能で KV キャッシュを必要としない RNN／Transformer ハイブリッドで、RWKV…
• llama-cpp：NVIDIA のハードウェアなしで、CPU、Apple Silicon、コンシューマー向け GPU 上で LLM 推論を実行します。
• tensorrt-llm：NVIDIA TensorRT を使って LLM 推論を最適化し、最大スループットと最小レイテンシを実現します。