diffusers/v0.43.0/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.

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.

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
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

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.

Développer

Nom du backend	Famille	Description
native	PyTorch natif	Backend par défaut utilisant scaled_dot_product_attention de PyTorch
flex	FlexAttention	Implémentation de FlexAttention pour PyTorch
_native_cudnn	PyTorch natif	Attention optimisée par CuDNN
_native_efficient	PyTorch natif	Attention efficace en mémoire
_native_flash	PyTorch natif	FlashAttention de PyTorch
_native_math	PyTorch natif	Attention basée sur les mathématiques (repli)
_native_npu	PyTorch natif	Attention optimisée pour NPU
_native_xla	PyTorch natif	Attention optimisée par XLA
flash	FlashAttention	FlashAttention-2
flash_hub	FlashAttention	FlashAttention-2 à partir de noyaux
flash_varlen	FlashAttention	FlashAttention de longueur variable
flash_varlen_hub	FlashAttention	FlashAttention de longueur variable à partir des noyaux
aiter_fa2_hub	Moteur de tenseurs IA pour ROCm	FlashAttention-2 pour AMD ROCm à partir des kernels
aiter_sage_hub	Moteur de tenseurs IA pour ROCm	Attention quantifiée (INT8 QK, FP8 PV) pour AMD RDNA4 (gfx1201) à partir de kernels
flash_4_hub	FlashAttention	FlashAttention-4
_flash_3	FlashAttention	FlashAttention-3
_flash_varlen_3	FlashAttention	FlashAttention-3 à longueur variable
_flash_3_hub	FlashAttention	FlashAttention-3 à partir de kernels
_flash_3_varlen_hub	FlashAttention	FlashAttention-3 de longueur variable à partir des noyaux
sage	SageAttention	Attention quantifiée (INT8 QK)
sage_hub	SageAttention	Attention quantifiée (INT8 QK) à partir de noyaux
sage_blackwell_hub	SageAttention	SageAttention3 attention FP4 pour les GPU Blackwell SM120 à partir de kernels
sage_varlen	SageAttention	SageAttention de longueur variable
_sage_qk_int8_pv_fp8_cuda	SageAttention	INT8 QK + FP8 PV (CUDA)
_sage_qk_int8_pv_fp8_cuda_sm90	SageAttention	INT8 QK + FP8 PV (SM90)
_sage_qk_int8_pv_fp16_cuda	SageAttention	INT8 QK + FP16 PV (CUDA)
_sage_qk_int8_pv_fp16_triton	SageAttention	INT8 QK + FP16 PV (Triton)
xformers	xFormers	Attention efficace en mémoire