Backends de atención

El despachador de atención es una función experimental. Abra un problema si tiene algún comentario o encuentra algún problema.

Diffusers ofrece varios algoritmos de atención optimizados que son más eficientes en memoria y computación a través de su attention dispatcher. El dispatcher actúa como un enrutador para gestionar y cambiar entre diferentes implementaciones de atención y proporciona una interfaz unificada para interactuar con ellas.

Consulte la tabla siguiente para obtener un resumen de las familias de atención disponibles y consulte la sección Backends disponibles para obtener una lista más completa.

Esta guía te mostrará cómo configurar y utilizar los distintos backends de atención.

set_attention_backend

El método set_attention_backend() recorre todos los módulos del modelo y establece el backend de atención adecuado para usar. La configuración del backend de atención persiste hasta que se llama a reset_attention_backend().

El siguiente ejemplo muestra cómo habilitar la implementación _flash_3_hub para FlashAttention-3 desde la biblioteca kernels, lo que permite usar de inmediato kernels de cálculo optimizados del Hub sin necesidad de ninguna configuración.

FlashAttention-3 requiere como mínimo GPUs Ampere.

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]

Para restaurar el backend de atención predeterminado, llama a reset_attention_backend().

pipeline.transformer.reset_attention_backend()

context manager attention_backend

El gestor de contexto attention_backend establece temporalmente un backend de atención para un modelo dentro del contexto. Fuera del contexto, se utiliza la atención predeterminada (la atención de producto escalar de puntos nativa de PyTorch). Esto es útil si deseas usar backends diferentes para distintas partes de un pipeline o si quieres probar los distintos backends.

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]

La mayoría de los backends de atención admiten torch.compile sin interrupciones del grafo y pueden usarse para acelerar aún más la inferencia.

Confianza en kernels remotos

Hub backends and other kernel-backed features (such as GGUF and Nunchaku Lite) download compute kernels from the Hub with kernels and execute their code locally.

By default, kernels only loads a kernel when its publisher is a trusted kernel publisher on the Hub. Kernels published under the kernels-community organization are trusted, so Diffusers loads them without any additional configuration. The _flash_3_hub, flash_hub, sage_hub, and the other Hub attention backends all resolve to kernels-community repositories.

Kernels from any other publisher are not vetted. Loading one downloads and runs code that Diffusers cannot vouch for, so Diffusers keeps it disabled unless you explicitly opt in with the DIFFUSERS_TRUST_REMOTE_KERNELS environment variable. When set, Diffusers forwards trust_remote_code=True to kernels so it loads kernels from untrusted publishers too.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Solo habilita esto después de inspeccionar el repositorio del kernel, ya que otorga al código descargado la capacidad de ejecutarse en tu máquina. Sin ello, cargar un kernel de un editor no confiable genera un error. Diffusers realiza esta comprobación por sí mismo, por lo que también se aplica a kernels<0.14.0, que es anterior al argumento trust_remote_code. Establecer DIFFUSERS_DISABLE_REMOTE_CODE=true deshabilita el código remoto de forma global y tiene prioridad sobre DIFFUSERS_TRUST_REMOTE_KERNELS.

Comprobaciones

El despachador de atención incluye comprobaciones de depuración que detectan errores comunes antes de que causen problemas.

1. Las comprobaciones de dispositivo verifican que los tensores de consulta, clave y valor estén en el mismo dispositivo.
2. Las comprobaciones de tipo de datos confirman que los tensores tienen dtypes coincidentes y que se utiliza bfloat16 o float16.
3. Las comprobaciones de forma validan las dimensiones de los tensores y evitan mezclar máscaras de atención con banderas causales.

Active estas comprobaciones estableciendo la variable de entorno DIFFUSERS_ATTN_CHECKS. Las comprobaciones añaden sobrecarga a cada operación de atención, por lo que están desactivadas de forma predeterminada.

export DIFFUSERS_ATTN_CHECKS=yes

Las comprobaciones se ejecutan ahora antes de cada operación de atención.

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

También puedes configurar el registro directamente.

from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True

Backends disponibles

Consulte la tabla siguiente para obtener una lista completa de los backends de atención disponibles y sus variantes.