diffusers/v0.43.0/ja/optimization/attention_backends.md

アテンションのバックエンド

attention dispatcher は実験的な機能です。フィードバックや問題が発生した場合は、issue を作成してください。

Diffusers は、attention dispatcher を通じて、メモリおよび計算効率のより高い複数の最適化されたアテンションアルゴリズムを提供しています。この dispatcher は、異なるアテンション実装の管理と切り替えを行うルーターとして機能し、それらとの対話のための統一されたインターフェースを提供します。

利用可能なアテンションファミリーの概要については下表を、より完全なリストについては利用可能なバックエンドセクションを参照してください。

アテンションファミリー	主な機能
FlashAttention	タイル化と再計算により、メモリの読み書きを最小限に抑えます
ROCm 向け AI Tensor Engine	AMD ROCm アクセラレータ向けに最適化されたFlashAttentionの実装
SageAttention	attentionをint8に量子化します
PyTorch ネイティブ	scaled_dot_product_attention を使用した PyTorch 組み込み実装
xFormers	各種アテンションカーネルに対応するメモリ効率の高いアテンション

このガイドでは、異なるアテンションバックエンドを設定し、使用する方法について説明します。

set_attention_backend

set_attention_backend() メソッドは、モデル内のすべてのモジュールを反復処理し、使用する適切なアテンションバックエンドを設定します。アテンションバックエンドの設定は、reset_attention_backend() が呼び出されるまで維持されます。

以下の例では、kernels ライブラリから FlashAttention-3 の _flash_3_hub 実装を有効にする方法を示しています。これにより、セットアップなしで Hub から最適化された計算カーネルをすぐに利用できます。

FlashAttention-3には、最低限Ampere GPUが必要です。

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)
pipeline.transformer.set_attention_backend("_flash_3_hub")

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""
pipeline(prompt).images[0]

デフォルトのattentionバックエンドを復元するには、reset_attention_backend()を呼び出します。

pipeline.transformer.reset_attention_backend()

attention_backend コンテキストマネージャー

attention_backend コンテキストマネージャは、コンテキスト内でモデルの注意機構バックエンドを一時的に設定します。コンテキストの外では、デフォルトの注意機構（PyTorch ネイティブのスケーリングドットプロダクトアテンション）が使用されます。パイプラインの異なる部分で異なるバックエンドを使用したい場合や、異なるバックエンドをテストしたい場合に便利です。

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)
prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

with attention_backend("_flash_3_hub"):
    image = pipeline(prompt).images[0]

ほとんどのアテンションバックエンドは、グラフの中断なしでtorch.compileをサポートしており、推論のさらなる高速化に利用できます。

リモートカーネルの信頼

Hub バックエンドおよびその他のカーネルベースの機能（GGUF や Nunchaku Lite など）は、kernels を使用して Hub から計算カーネルをダウンロードし、そのコードをローカルで実行します。

デフォルトでは、kernels は、その公開者が Hub 上の信頼できるカーネル公開者である場合のみカーネルを読み込みます。kernels-community 組織の下で公開されたカーネルは信頼できるため、Diffusers は追加の設定なしでそれらを読み込みます。_flash_3_hub、flash_hub、sage_hub およびその他の Hub アテンションバックエンドは、すべて kernels-community リポジトリに解決されます。

他のパブリッシャーによるカーネルは検証されていません。1つを読み込むと、Diffusers が保証できないコードがダウンロードされて実行されるため、Diffusers は DIFFUSERS_TRUST_REMOTE_KERNELS 環境変数で明示的にオプトインしない限り、これを無効にしています。設定すると、Diffusers は kernels に trust_remote_code=True を渡し、信頼できないパブリッシャーのカーネルも読み込むようになります。

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

カーネルのリポジトリを確認した後にのみこれを有効にしてください。有効にすると、ダウンロードされたコードがあなたのマシン上で実行されるようになります。これを有効にしない場合、信頼できないパブリシャーからのカーネルを読み込むとエラーが発生します。Diffusers はこのチェックを自身で実行するため、trust_remote_code 引数が存在する以前の kernels<0.14.0 にも適用されます。DIFFUSERS_DISABLE_REMOTE_CODE=true を設定すると、リモートコードがグローバルに無効化され、DIFFUSERS_TRUST_REMOTE_KERNELS より優先されます。

チェック

アテンションディスパッチャーには、問題を引き起こす前に一般的なエラーを検出するデバッグチェックが含まれています。

1. Device checksは、クエリ、キー、バリューのテンソルが同じデバイス上にあることを確認します。
2. データ型チェックにより、テンソルが一致するdtypeを持ち、bfloat16またはfloat16のいずれかを使用していることを確認します。
3. Shape checksはテンソルの次元を検証し、attention maskとcausal flagの混在を防ぎます。

DIFFUSERS_ATTN_CHECKS 環境変数を設定することで、これらのチェックを有効にできます。チェックはすべてのアテンション操作にオーバーヘッドを伴うため、デフォルトでは無効になっています。

export DIFFUSERS_ATTN_CHECKS=yes

これらのチェックは、すべてのアテンション操作の前に実行されます。

import torch

query = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")  # or "mps", "xpu", "cpu"
key = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")
value = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")

try:
    with attention_backend("flash"):
        output = dispatch_attention_fn(query, key, value)
        print("✓ Flash Attention works with checks enabled")
except Exception as e:
    print(f"✗ Flash Attention failed: {e}")

レジストリを直接設定することもできます。

from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True

利用可能なバックエンド

利用可能なアテンションバックエンドとそのバリアントの完全なリストについては、以下の表を参照してください。

展開

バックエンド名	ファミリー	説明
native	PyTorch ネイティブ	PyTorchのを使用するデフォルトのバックエンド
flex	FlexAttention	PyTorch FlexAttention 実装
_native_cudnn	PyTorch ネイティブ	CuDNN最適化アテンション
_native_efficient	PyTorch ネイティブ	メモリ効率の高いアテンション
_native_flash	PyTorch ネイティブ	PyTorchのFlashAttention
_native_math	PyTorch ネイティブ	数学ベースの注意機構（フォールバック）
_native_npu	PyTorch ネイティブ	NPU最適化されたアテンション
_native_xla	PyTorch ネイティブ	XLA最適化されたアテンション
flash	FlashAttention	FlashAttention-2
flash_hub	FlashAttention	カーネルからのFlashAttention-2
flash_varlen	FlashAttention	可変長のFlashAttention
flash_varlen_hub	FlashAttention	カーネルからの可変長FlashAttention
aiter_fa2_hub	ROCm 向け AI Tensor Engine	kernels による AMD ROCm 向け FlashAttention-2
aiter_sage_hub	ROCm 向け AI Tensor Engine	kernels による AMD RDNA4 (gfx1201) 向けの量子化アテンション (INT8 QK、FP8 PV)
flash_4_hub	FlashAttention	FlashAttention-4
_flash_3	FlashAttention	FlashAttention-3
_flash_varlen_3	FlashAttention	可変長のFlashAttention-3
_flash_3_hub	FlashAttention	kernels による FlashAttention-3
_flash_3_varlen_hub	FlashAttention	カーネルからの可変長FlashAttention-3
sage	SageAttention	量子化されたアテンション（INT8 QK）
sage_hub	SageAttention	カーネルによる量子化アテンション（INT8 QK）
sage_blackwell_hub	SageAttention	kernels による SM120 Blackwell GPU 向けの SageAttention3 FP4 アテンション
sage_varlen	SageAttention	可変長のSageAttention
_sage_qk_int8_pv_fp8_cuda	SageAttention	INT8 QK + FP8 PV (CUDA)
_sage_qk_int8_pv_fp8_cuda_sm90	SageAttention	INT8 QK + FP8 PV (SM90)
_sage_qk_int8_pv_fp16_cuda	SageAttention	INT8 QK + FP16 PV (CUDA)
_sage_qk_int8_pv_fp16_triton	SageAttention	INT8 QK + FP16 PV (Triton)
xformers	xFormers	メモリ効率の高いアテンション