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.

La mayoría de los modelos transformer de Diffusers enrutan la atención a través de un attention dispatcher para que puedas cambiar entre backends optimizados detrás de una única API. El dispatcher gestiona las implementaciones registradas y expone una ruta de llamada unificada para ellas. Algunos modelos, como muchos autoencoders, no utilizan el dispatcher debido a su estructura interna, y cambiar el backend no tiene efecto en sus capas de atención.

Consulte la tabla siguiente para obtener un resumen de las familias de atención disponibles y consulte la sección Backends disponibles para una lista más completa. El backend más rápido depende del modelo, la GPU, el dtype y la forma de la entrada.

Los backends de Hub (los nombres *_hub) requieren la biblioteca Kernels y descargan el kernel en el primer uso. Los backends de Hub de SageAttention también requieren DIFFUSERS_TRUST_REMOTE_KERNELS=true (ver Confianza en kernels remotos). Otros backends requieren su propio paquete, como flash-attn o sageattention. La tabla Backends disponibles enumera los requisitos que Diffusers comprueba al habilitar un backend.

Establecer un backend en el modelo

El método set_attention_backend() recorre las capas de atención del modelo y aplica el backend elegido en cada una. También establece el backend activo a nivel de proceso del despachador con el mismo valor.

reset_attention_backend() solo limpia el backend en las capas de atención. No limpia el backend activo a nivel de proceso. Para un cambio temporal que restaure el backend activo anterior al salir, use el gestor de contexto attention_backend.

El ejemplo siguiente habilita _flash_3_hub (FlashAttention-3 desde el Hub) con 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]

Los backends de FlashAttention-3 no-Hub (_flash_3, _flash_varlen_3) requieren compilar FlashAttention-3 desde el código fuente y se deprecarán pronto. Utilice _flash_3_hub o _flash_3_varlen_hub en su lugar.

Prueba un backend temporalmente

El gestor de contexto attention_backend() establece el backend activo a nivel de proceso durante la duración del bloque y restaura el backend anterior al salir del bloque. Úselo para probar un backend en una sola llamada sin dejar un backend permanente aplicado desde 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]

La mayoría de los backends de atención funcionan con torch.compile. Si eso acelera tu pipeline depende del modelo y del backend. Consulta Precisión y compilación.

Confianza en kernels remotos

Los backends de Hub requieren primero la biblioteca Kernels y, a continuación, Diffusers descarga el kernel de Hub en su primer uso.

Los backends de atención de Hub descargan kernels de cómputo con la biblioteca Kernels y los ejecutan localmente. La mayoría de los nombres de atención de Hub (_flash_3_hub, flash_hub y los demás backends *_hub) se resuelven a la organización kernels-community. Dicha organización es un editor de confianza en Kernels, por lo que Diffusers carga esos kernels de atención sin establecer DIFFUSERS_TRUST_REMOTE_KERNELS.

Los backends de SageAttention Hub (sage_hub y sage_blackwell_hub) se cargan desde la organización SageAttention. Establezca DIFFUSERS_TRUST_REMOTE_KERNELS=true para usarlos.

Otras funciones respaldadas por kernels, como GGUF y Nunchaku Lite, pueden obtener kernels de editores externos a kernels-community. Esas rutas permanecen bloqueadas a menos que se activen con DIFFUSERS_TRUST_REMOTE_KERNELS. Cuando se establece, Diffusers reenvía trust_remote_code=True a Kernels para que también puedan cargarse editores no confiables.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Solo active esto después de inspeccionar el repositorio del kernel. 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 puede ejecutar comprobaciones de depuración antes de cada llamada de atención despachada. Las comprobaciones que se ejecutan dependen de las restricciones registradas para el backend activo.

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, cuando están registradas, confirman la coincidencia de dtypes y a menudo requieren 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.

Habilite las comprobaciones con la variable de entorno DIFFUSERS_ATTN_CHECKS. Las comprobaciones añaden sobrecarga, por lo que están desactivadas de forma predeterminada.

export DIFFUSERS_ATTN_CHECKS=yes

Con las comprobaciones activadas, Diffusers ejecuta esas restricciones antes de cada llamada de atención despachada. El ejemplo de bajo nivel a continuación llama directamente a dispatch_attention_fn. La inferencia de la pipeline no necesita esa importación. Solo necesita que el backend esté configurado mediante set_attention_backend() o 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}")

Backends disponibles

Consulte la tabla siguiente para obtener una lista completa de los backends de atención disponibles y sus variantes. Diffusers comprueba la disponibilidad del paquete y las versiones fijadas cuando se habilita un backend.