Attention backends

The attention dispatcher is an experimental feature. Please open an issue if you have any feedback or encounter any problems.

Most Diffusers transformer models route attention through an attention dispatcher so you can switch optimized backends behind one API. The dispatcher manages registered implementations and exposes a unified call path for them. Some models, such as many autoencoders, don’t use the dispatcher because of their internals, and switching the backend has no effect on their attention layers.

Refer to the table below for an overview of the available attention families and to the Available backends section for a more complete list. The fastest backend depends on the model, GPU, dtype, and input shape.

Hub backends (the *_hub names) need the Kernels library and download the kernel on first use. The SageAttention Hub backends also need DIFFUSERS_TRUST_REMOTE_KERNELS=true (see Trusting remote kernels). Other backends need their own package, such as flash-attn or sageattention. The Available backends table lists the requirements Diffusers checks when you enable a backend.

Set a backend on the model

The set_attention_backend() method walks the model’s attention layers and applies the chosen backend on each one. It also sets the dispatcher’s process-wide active backend to the same value.

reset_attention_backend() clears the backend on attention layers only. It does not clear the process-wide active backend. For a temporary switch that restores the previous active backend on exit, use the attention_backend context manager.

The example below enables _flash_3_hub (FlashAttention-3 from the Hub) with device_map="cuda".

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]

The non-Hub FlashAttention-3 backends (_flash_3, _flash_varlen_3) require building FlashAttention-3 from source and will be deprecated soon. Use _flash_3_hub or _flash_3_varlen_hub instead.

Try a backend temporarily

The attention_backend() context manager sets the process-wide active backend for the duration of the block and restores the previous backend when the block exits. Use it to try a backend for one call without leaving a permanent backend applied from set_attention_backend().

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]

Most attention backends work with torch.compile. Whether that speeds up your pipeline depends on the model and backend. See Precision and compilation.

Trusting remote kernels

Hub backends need the Kernels library first, and then Diffusers fetches the Hub kernel on first use.

Hub attention backends download compute kernels with the Kernels library and run them locally. Most Hub attention names (_flash_3_hub, flash_hub, and the other *_hub backends) resolve to the kernels-community organization. That organization is a trusted publisher in Kernels, so Diffusers loads those attention kernels without setting DIFFUSERS_TRUST_REMOTE_KERNELS.

The SageAttention Hub backends (sage_hub and sage_blackwell_hub) load from the SageAttention organization instead. Set DIFFUSERS_TRUST_REMOTE_KERNELS=true to use them.

Other kernel-backed features such as GGUF and Nunchaku Lite can pull kernels from publishers outside kernels-community. Those paths stay blocked unless you opt in with DIFFUSERS_TRUST_REMOTE_KERNELS. When set, Diffusers forwards trust_remote_code=True to Kernels so untrusted publishers can load too.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Only enable this after inspecting the kernel repository. Without it, loading a kernel from an untrusted publisher raises an error. Diffusers performs this check itself, so it also applies to kernels<0.14.0, which predates the trust_remote_code argument. Setting DIFFUSERS_DISABLE_REMOTE_CODE=true disables remote code globally and takes precedence over DIFFUSERS_TRUST_REMOTE_KERNELS.

Checks

The attention dispatcher can run debugging checks before each dispatched attention call. Which checks run depends on the constraints registered for the active backend.

1. Device checks verify that query, key, and value tensors live on the same device.
2. Data type checks, where registered, confirm matching dtypes and often require bfloat16 or float16.
3. Shape checks validate tensor dimensions and prevent mixing attention masks with causal flags.

Enable checks with the DIFFUSERS_ATTN_CHECKS environment variable. Checks add overhead, so they are disabled by default.

export DIFFUSERS_ATTN_CHECKS=yes

With checks on, Diffusers runs those constraints before every dispatched attention call. The low-level example below calls dispatch_attention_fn directly. Pipeline inference does not need that import. It only needs the backend set via set_attention_backend() or 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}")

Available backends

Refer to the table below for a complete list of available attention backends and their variants. Diffusers checks package availability and version pins when you enable a backend.