Attention backends

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

Diffusers provides several optimized attention algorithms that are more memory and computationally efficient through it’s attention dispatcher. The dispatcher acts as a router for managing and switching between different attention implementations and provides a unified interface for interacting with them.

Refer to the table below for an overview of the available attention families and to the Available backends section for a more complete list.

This guide will show you how to set and use the different attention backends.

set_attention_backend

The set_attention_backend() method iterates through all the modules in the model and sets the appropriate attention backend to use. The attention backend setting persists until reset_attention_backend() is called.

The example below demonstrates how to enable the _flash_3_hub implementation for FlashAttention-3 from the kernels library, which allows you to instantly use optimized compute kernels from the Hub without requiring any setup.

FlashAttention-3 requires Ampere GPUs at a minimum.

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", torch_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]

To restore the default attention backend, call reset_attention_backend().

pipeline.transformer.reset_attention_backend()

attention_backend context manager

The attention_backend context manager temporarily sets an attention backend for a model within the context. Outside the context, the default attention (PyTorch’s native scaled dot product attention) is used. This is useful if you want to use different backends for different parts of a pipeline or if you want to test the different backends.

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", torch_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 support torch.compile without graph breaks and can be used to further speed up inference.

Checks

The attention dispatcher includes debugging checks that catch common errors before they cause problems.

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

Enable these checks by setting the DIFFUSERS_ATTN_CHECKS environment variable. Checks add overhead to every attention operation, so they’re disabled by default.

export DIFFUSERS_ATTN_CHECKS=yes

The checks are run now before every attention operation.

import torch

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}")

You can also configure the registry directly.

from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True

Available backends

Refer to the table below for a complete list of available attention backends and their variants.