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

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

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

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

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

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

利用可能なバックエンド

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