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

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

Diffusersのtransformerモデルの多くは、attention dispatcherを介してattentionをルーティングします。これにより、単一のAPIの背後で最適化されたバックエンドを切り替えることができます。dispatcherは登録された実装を管理し、それらに対して統一された呼び出しパスを公開します。ただし、多くのautoencoderなどの一部のモデルは、内部構造の都合上、dispatcherを使用しません。そのため、バックエンドを切り替えても、これらのモデルのattention層には影響を与えません。

利用可能なアテンションファミリーの概要については下表を、より完全なリストについては利用可能なバックエンドセクションを参照してください。最速のバックエンドは、モデル、GPU、dtype、入力形状によって異なります。

Hub バックエンド（*_hub という名前のもの）には Kernels ライブラリが必要であり、初回使用時にカーネルをダウンロードします。SageAttention の Hub バックエンドには、DIFFUSERS_TRUST_REMOTE_KERNELS=true も必要です（リモートカーネルの信頼を参照）。その他のバックエンドには、flash-attn や sageattention のような独自のパッケージが必要です。利用可能なバックエンド の表には、バックエンドを有効にした際に Diffusers がチェックする要件が記載されています。

モデルにバックエンドを設定する

set_attention_backend() メソッドは、モデルの注意層を走査し、各層に選択されたバックエンドを適用します。また、ディスパッチャーのプロセス全体で有効なバックエンドも同じ値に設定します。

reset_attention_backend() は、アテンションレイヤーのバックエンドのみをクリアします。プロセス全体のアクティブなバックエンドはクリアされません。終了時に以前のアクティブなバックエンドを復元する一時的な切り替えには、attention_backend コンテキストマネージャーを使用してください。

以下の例では、device_map="cuda" を使用して _flash_3_hub（Hub からの FlashAttention-3）を有効にしています。

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"
)
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]

非HubのFlashAttention-3バックエンド（_flash_3、_flash_varlen_3）は、FlashAttention-3をソースからビルドする必要があり、まもなく非推奨になります。代わりに_flash_3_hubまたは_flash_3_varlen_hubを使用してください。

バックエンドを一時的に試す

attention_backend() コンテキストマネージャは、ブロックの持続期間中、プロセス全体のアクティブなバックエンドを設定し、ブロックを抜けた際に以前のバックエンドを復元します。set_attention_backend() によって永続的なバックエンドが適用された状態を残さずに、1回の呼び出しでバックエンドを試す際に使用します。

import torch
from diffusers import QwenImagePipeline, attention_backend

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"
)
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 バックエンドでは、まずKernelsライブラリが必要です。その後、Diffusers が初回使用時に Hub からカーネルを取得します。

Hub attention backends は Kernels ライブラリを使用して計算カーネルをダウンロードし、ローカルで実行します。ほとんどの Hub attention 名（_flash_3_hub、flash_hub、およびその他の *_hub backends）は kernels-community 組織に解決されます。この組織は Kernels における信頼できるパブリッシャーであるため、Diffusers は DIFFUSERS_TRUST_REMOTE_KERNELS を設定せずにこれらの attention カーネルを読み込みます。

SageAttention Hub バックエンド（sage_hub および sage_blackwell_hub）は、代わりに SageAttention 組織から読み込みます。これらを使用するには、DIFFUSERS_TRUST_REMOTE_KERNELS=true を設定してください。

GGUF や Nunchaku Lite などの他のカーネルベースの機能は、GGUF と Nunchaku Lite が kernels-community 以外のパブリッシャーからカーネルを取得できます。これらのパスは、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 より優先されます。

チェック

attention dispatcherは、各ディスパッチされたアテンション呼び出しの前にデバッグチェックを実行できます。実行されるチェックは、アクティブなバックエンドに登録された制約によって決まります。

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

DIFFUSERS_ATTN_CHECKS 環境変数を使用してチェックを有効にします。チェックにはオーバーヘッドが発生するため、デフォルトでは無効になっています。

export DIFFUSERS_ATTN_CHECKS=yes

チェックが有効な場合、Diffusers はすべてのディスパッチされたアテンション呼び出しの前に、これらの制約を実行します。以下の低レベルの例では、dispatch_attention_fn を直接呼び出しています。パイプラインの推論では、そのインポートは不要です。set_attention_backend() または attention_backend() によってバックエンドを設定するだけで十分です。

import torch
from diffusers.models.attention_dispatch import attention_backend, dispatch_attention_fn

query = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")
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}")

利用可能なバックエンド

利用可能なアテンションバックエンドとそのバリアントの完全なリストについては、以下の表を参照してください。バックエンドを有効にすると、Diffusersはパッケージの可用性とバージョンの固定を確認します。