diffusers/main/fr/optimization/attention_backends.md

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.

famille d'attention	fonctionnalité principale
FlashAttention	minimise les lectures/écritures mémoire grâce au tuilage et à la recomputation
Moteur de tenseurs IA pour ROCm	Implémentation de FlashAttention optimisée pour les accélérateurs AMD ROCm
SageAttention	quantifie l'attention en int8
FlexAttention	PyTorch FlexAttention
PyTorch natif	implémentation intégrée de PyTorch utilisant scaled_dot_product_attention
xFormers	attention efficace en mémoire avec prise en charge de divers noyaux d'attention

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.

Nom du backend	Famille	Description	Prérequis
native	PyTorch natif	Backend par défaut utilisant scaled_dot_product_attention de PyTorch	Aucun
flex	FlexAttention	PyTorch FlexAttention	torch>=2.5.0
_native_cudnn	PyTorch natif	Attention optimisée par CuDNN	CUDA + CuDNN
_native_efficient	PyTorch natif	Attention efficace en mémoire	Aucun au-delà de PyTorch
_native_flash	PyTorch natif	FlashAttention de PyTorch	CUDA
_native_math	PyTorch natif	Attention basée sur les mathématiques (repli)	Aucun
_native_npu	PyTorch natif	Attention optimisée pour NPU	torch_npu
_native_xla	PyTorch natif	Attention optimisée par XLA	torch_xla>=2.2
flash	FlashAttention	FlashAttention-2	flash-attn>=2.6.3
flash_hub	FlashAttention	FlashAttention-2 des noyaux du Hub	kernels>=0.12
flash_varlen	FlashAttention	FlashAttention de longueur variable	flash-attn>=2.6.3
flash_varlen_hub	FlashAttention	FlashAttention de longueur variable à partir des noyaux du Hub	kernels>=0.12
aiter_fa2_hub	Moteur de tenseurs IA pour ROCm	FlashAttention-2 pour AMD ROCm à partir des noyaux du Hub (bfloat16)	kernels>=0.12, ROCm
flash_4_hub	FlashAttention	FlashAttention-4 des noyaux du Hub	kernels>=0.12.3
_flash_3	FlashAttention	FlashAttention-3 (local ; bientôt obsolète)	Compiler FA3 à partir du code source
_flash_varlen_3	FlashAttention	FlashAttention-3 de longueur variable (local ; bientôt obsolète)	Compiler FA3 à partir du code source
_flash_3_hub	FlashAttention	FlashAttention-3 des noyaux du Hub	kernels>=0.12
_flash_3_varlen_hub	FlashAttention	FlashAttention-3 de longueur variable à partir des noyaux du Hub	kernels>=0.12
sage	SageAttention	Attention quantifiée (INT8 QK)	sageattention>=2.1.1
sage_hub	SageAttention	Attention quantifiée (INT8 QK) à partir des noyaux du Hub	kernels>=0.12, DIFFUSERS_TRUST_REMOTE_KERNELS=true
sage_blackwell_hub	SageAttention	SageAttention3 attention FP4 pour les GPU Blackwell SM120 à partir des noyaux du Hub	kernels>=0.12, DIFFUSERS_TRUST_REMOTE_KERNELS=true
sage_varlen	SageAttention	SageAttention de longueur variable	sageattention>=2.1.1
_sage_qk_int8_pv_fp8_cuda	SageAttention	INT8 QK + FP8 PV (CUDA)	sageattention>=2.1.1
_sage_qk_int8_pv_fp8_cuda_sm90	SageAttention	INT8 QK + FP8 PV (SM90)	sageattention>=2.1.1 ; SM90
_sage_qk_int8_pv_fp16_cuda	SageAttention	INT8 QK + FP16 PV (CUDA)	sageattention>=2.1.1
_sage_qk_int8_pv_fp16_triton	SageAttention	INT8 QK + FP16 PV (Triton)	sageattention>=2.1.1
xformers	xFormers	Attention efficace en mémoire	xformers>=0.0.29