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.

La plupart des modèles transformer de Diffusers acheminent l’attention via un dispatcheur d’attention afin de pouvoir basculer entre des backends optimisés derrière une seule API. Le dispatcheur gère les implémentations enregistrées et expose un chemin d’appel unifié pour celles-ci. Certains modèles, comme de nombreux autoencodeurs, n’utilisent pas le dispatcheur en raison de leur architecture interne, et le changement de backend n’a aucun effet sur leurs couches d’attention.

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. Le backend le plus rapide dépend du modèle, du GPU, du dtype et de la forme des entrées.

Les backends du Hub (les noms *_hub) nécessitent la bibliothèque Kernels et téléchargent le noyau lors de la première utilisation. Les backends du Hub SageAttention nécessitent également DIFFUSERS_TRUST_REMOTE_KERNELS=true (voir Faire confiance aux noyaux distants). Les autres backends nécessitent leur propre package, tel que flash-attn ou sageattention. Le tableau Backends disponibles répertorie les exigences que Diffusers vérifie lorsque vous activez un backend.

Définir un backend sur le modèle

La méthode set_attention_backend() parcourt les couches d’attention du modèle et applique le backend choisi à chacune d’entre elles. Elle définit également le backend actif du dispatcher pour l’ensemble du processus à la même valeur.

reset_attention_backend() efface le backend uniquement sur les couches d’attention. Il n’efface pas le backend actif à l’échelle du processus. Pour un changement temporaire qui restaure le backend actif précédent à la sortie, utilisez le gestionnaire de contexte attention_backend.

L'exemple ci-dessous active _flash_3_hub (FlashAttention-3 depuis le Hub) avec 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]

Les backends FlashAttention-3 non-Hub (_flash_3, _flash_varlen_3) nécessitent la compilation de FlashAttention-3 depuis le code source et seront bientôt dépréciés. Utilisez plutôt _flash_3_hub ou _flash_3_varlen_hub.

Essayer un backend temporairement

Le gestionnaire de contexte attention_backend() définit le backend actif pour l’ensemble du processus pendant la durée du bloc et restaure le backend précédent à la sortie du bloc. Utilisez-le pour tester un backend lors d’un seul appel sans laisser un backend permanent appliqué par 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 plupart des backends d'attention fonctionnent avec torch.compile. L'accélération de votre pipeline dépend du modèle et du backend. Consultez Précision et compilation.

Faire confiance aux noyaux distants

Les backends du Hub ont d'abord besoin de la bibliothèque Kernels, puis Diffusers récupère le noyau du Hub lors de la première utilisation.

Les backends d'attention du Hub téléchargent des noyaux de calcul via la bibliothèque Kernels et les exécutent localement. La plupart des noms d'attention du Hub (_flash_3_hub, flash_hub et les autres backends *_hub) résolvent vers l'organisation kernels-community. Cette organisation est un éditeur de confiance dans Kernels, Diffusers charge donc ces noyaux d'attention sans définir DIFFUSERS_TRUST_REMOTE_KERNELS.

Les backends du Hub SageAttention (sage_hub et sage_blackwell_hub) sont chargés depuis l’organisation SageAttention. Définissez DIFFUSERS_TRUST_REMOTE_KERNELS=true pour les utiliser.

D’autres fonctionnalités basées sur des noyaux, telles que GGUF et Nunchaku Lite, peuvent récupérer des noyaux auprès d’éditeurs extérieurs à kernels-community. Ces chemins restent bloqués à moins que vous n’optiez pour leur activation via DIFFUSERS_TRUST_REMOTE_KERNELS. Une fois défini, Diffusers transmet trust_remote_code=True à Kernels afin que les éditeurs non approuvés puissent également être chargés.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Activez cette option uniquement après avoir inspecté le dépôt du noyau. Sans elle, 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 peut exécuter des vérifications de débogage avant chaque appel d'attention dispatché. Les vérifications exécutées dépendent des contraintes enregistrées pour le backend actif.

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, lorsqu’elles sont enregistrées, confirment la correspondance des dtypes et exigent souvent bfloat16 ou 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 les vérifications avec la variable d’environnement DIFFUSERS_ATTN_CHECKS. Les vérifications ajoutent une surcharge, elles sont donc désactivées par défaut.

export DIFFUSERS_ATTN_CHECKS=yes

Avec les vérifications activées, Diffusers exécute ces contraintes avant chaque appel d’attention dispatché. L’exemple de bas niveau ci-dessous appelle directement dispatch_attention_fn. L’inférence de pipeline n’a pas besoin de cet import. Il suffit de définir le backend via set_attention_backend() ou 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

Reportez-vous au tableau ci-dessous pour la liste complète des backends d'attention disponibles et de leurs variantes. Diffusers vérifie la disponibilité des packages et les versions épinglées lorsque vous activez un backend.