diffusers/main/ja/optimization/attention_backends.md

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

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

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

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

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

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はパッケージの可用性とバージョンの固定を確認します。

バックエンド名	ファミリー	説明	前提条件
native	PyTorch ネイティブ	PyTorchのを使用するデフォルトのバックエンド	None
flex	FlexAttention	PyTorch FlexAttention	torch>=2.5.0
_native_cudnn	PyTorch ネイティブ	CuDNN最適化アテンション	CUDA + CuDNN
_native_efficient	PyTorch ネイティブ	メモリ効率の高いアテンション	PyTorch以外にはありません
_native_flash	PyTorch ネイティブ	PyTorchのFlashAttention	CUDA
_native_math	PyTorch ネイティブ	数学ベースの注意機構（フォールバック）	None
_native_npu	PyTorch ネイティブ	NPU最適化されたアテンション	torch_npu
_native_xla	PyTorch ネイティブ	XLA最適化されたアテンション	torch_xla>=2.2
flash	FlashAttention	FlashAttention-2	flash-attn>=2.6.3
flash_hub	FlashAttention	HubカーネルからのFlashAttention-2	kernels>=0.12
flash_varlen	FlashAttention	可変長のFlashAttention	flash-attn>=2.6.3
flash_varlen_hub	FlashAttention	Hubカーネルからの可変長FlashAttention	kernels>=0.12
aiter_fa2_hub	ROCm 向け AI Tensor Engine	HubカーネルからのAMD ROCm向けFlashAttention-2(bfloat16)	kernels>=0.12、ROCm
flash_4_hub	FlashAttention	HubカーネルからのFlashAttention-4	kernels>=0.12.3
_flash_3	FlashAttention	FlashAttention-3（ローカル；まもなく非推奨）	ソースからFA3をビルドする
_flash_varlen_3	FlashAttention	可変長FlashAttention-3（ローカル；まもなく非推奨）	ソースからFA3をビルドする
_flash_3_hub	FlashAttention	HubカーネルからのFlashAttention-3	kernels>=0.12
_flash_3_varlen_hub	FlashAttention	Hubカーネルからの可変長FlashAttention-3	kernels>=0.12
sage	SageAttention	量子化されたアテンション（INT8 QK）	sageattention>=2.1.1
sage_hub	SageAttention	Hubカーネルによる量子化アテンション（INT8 QK）	kernels>=0.12, DIFFUSERS_TRUST_REMOTE_KERNELS=true
sage_blackwell_hub	SageAttention	HubカーネルによるSM120 Blackwell GPU向けのSageAttention3 FP4アテンション	kernels>=0.12, DIFFUSERS_TRUST_REMOTE_KERNELS=true
sage_varlen	SageAttention	可変長のSageAttention	sageattention>=2.1.1
_sage_qk_int8_pv_fp8_cuda	SageAttention	INT8 QK + FP8 PV (CUDA)	sageattention>=2.1.1
_sage_qk_int8_pv_fp8_cuda_sm90	SageAttention	INT8 QK + FP8 PV (SM90)	sageattention>=2.1.1; SM90
_sage_qk_int8_pv_fp16_cuda	SageAttention	INT8 QK + FP16 PV (CUDA)	sageattention>=2.1.1
_sage_qk_int8_pv_fp16_triton	SageAttention	INT8 QK + FP16 PV (Triton)	sageattention>=2.1.1
xformers	xFormers	メモリ効率の高いアテンション	xformers>=0.0.29