Backends d'attention

Le dispatcheur d'attention est une fonctionnalité expérimentale. Veuillez ouvrir un ticket si vous avez des retours ou si vous rencontrez des problèmes.

Diffusers fournit plusieurs algorithmes d'attention optimisés, plus efficaces en termes de mémoire et de calcul, via son attention dispatcher. Le dispatcher agit comme un routeur pour gérer et basculer entre différentes implémentations d'attention et offre une interface unifiée pour interagir avec elles.

Reportez-vous au tableau ci-dessous pour un aperçu des familles d'attention disponibles et à la section Backends disponibles pour une liste plus complète.

Ce guide vous montre comment configurer et utiliser les différents backends d'attention.

set_attention_backend

La méthode set_attention_backend() parcourt tous les modules du modèle et définit le backend d'attention approprié à utiliser. Le réglage du backend d'attention est conservé jusqu'à l'appel de reset_attention_backend().

L’exemple ci-dessous montre comment activer l’implémentation _flash_3_hub de FlashAttention-3 à partir de la bibliothèque kernels, ce qui permet d’utiliser instantanément des noyaux de calcul optimisés depuis le Hub, sans aucune configuration préalable.

FlashAttention-3 nécessite au minimum des GPU 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]

Pour rétablir le backend d'attention par défaut, appelez reset_attention_backend().

pipeline.transformer.reset_attention_backend()

gestionnaire de contexte attention_backend

Le gestionnaire de contexte attention_backend définit temporairement un backend d'attention pour un modèle au sein du contexte. En dehors du contexte, l'attention par défaut (l'attention produit scalaire normalisée native de PyTorch) est utilisée. Cela est utile si vous souhaitez utiliser des backends différents pour différentes parties d'un pipeline ou si vous souhaitez tester les différents 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 plupart des backends d'attention prennent en charge torch.compile sans rupture de graphe et peuvent être utilisés pour accélérer davantage l'inférence.

Faire confiance aux noyaux distants

Les backends du Hub et les autres fonctionnalités basées sur des noyaux (tels que GGUF et Nunchaku Lite) téléchargent les noyaux de calcul depuis le Hub à l'aide de kernels et exécutent leur code localement.

Par défaut, kernels ne charge un noyau que si son éditeur est un éditeur de noyaux de confiance sur le Hub. Les noyaux publiés sous l’organisation kernels-community sont de confiance, Diffusers les charge donc sans configuration supplémentaire. Les backends d’attention du Hub _flash_3_hub, flash_hub, sage_hub et les autres se résolvent tous vers des dépôts kernels-community.

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

Activez cette option uniquement après avoir inspecté le dépôt du noyau, car elle accorde au code téléchargé la capacité de s’exécuter sur votre machine. Sans cela, le chargement d’un noyau provenant d’un éditeur non fiable provoque une erreur. Diffusers effectue lui-même ce contrôle, ce qui s’applique donc également à kernels<0.14.0, qui précède l’argument trust_remote_code. La définition de DIFFUSERS_DISABLE_REMOTE_CODE=true désactive le code distant de manière globale et prend le pas sur DIFFUSERS_TRUST_REMOTE_KERNELS.

Vérifications

Le dispatcheur d'attention inclut des vérifications de débogage qui détectent les erreurs courantes avant qu'elles ne causent des problèmes.

1. Les vérifications de périphérique confirment que les tenseurs de requête, de clé et de valeur sont sur le même périphérique.
2. Les vérifications de type de données confirment que les tenseurs ont des dtypes correspondants et utilisent soit bfloat16, soit float16.
3. Les vérifications de forme valident les dimensions des tenseurs et empêchent le mélange des masques d'attention avec les drapeaux causaux.

Activez ces vérifications en définissant la variable d’environnement DIFFUSERS_ATTN_CHECKS. Les vérifications ajoutent une surcharge à chaque opération d’attention, elles sont donc désactivées par défaut.

export DIFFUSERS_ATTN_CHECKS=yes

Les vérifications sont désormais exécutées avant chaque opération d'attention.

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

Vous pouvez également configurer le registre directement.

from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True

Backends disponibles

Reportez-vous au tableau ci-dessous pour obtenir la liste complète des backends d'attention disponibles et de leurs variantes.