diffusers/v0.43.0/pt/optimization/attention_backends.md

Backends de atenção

O despachante de atenção é um recurso experimental. Abra uma issue se tiver algum feedback ou encontrar algum problema.

O Diffusers fornece vários algoritmos de atenção otimizados, mais eficientes em termos de memória e computação, por meio do seu dispatcher de atenção. O dispatcher atua como um roteador para gerenciar e alternar entre diferentes implementações de atenção e fornece uma interface unificada para interagir com elas.

Consulte a tabela abaixo para uma visão geral das famílias de atenção disponíveis e a seção Backends disponíveis para uma lista mais completa.

família de atenção	principal recurso
FlashAttention	minimiza leituras/escritas de memória por meio de tiling e recomputação
AI Tensor Engine para ROCm	Implementação do FlashAttention otimizada para aceleradores AMD ROCm
SageAttention	quantiza a atenção para int8
PyTorch nativo	implementação embutida do PyTorch usando scaled_dot_product_attention
xFormers	atenção eficiente em memória com suporte para vários kernels de atenção

Este guia mostrará como configurar e usar os diferentes backends de atenção.

set_attention_backend

O método set_attention_backend() percorre todos os módulos do modelo e define o backend de atenção apropriado para uso. A configuração do backend de atenção permanece ativa até que reset_attention_backend() seja chamado.

O exemplo abaixo demonstra como habilitar a implementação _flash_3_hub para FlashAttention-3 a partir da biblioteca kernels, o que permite usar instantaneamente kernels de computação otimizados do Hub sem exigir nenhuma configuração.

FlashAttention-3 requer GPUs Ampere no mínimo.

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]

Para restaurar o backend de atenção padrão, chame reset_attention_backend().

pipeline.transformer.reset_attention_backend()

context manager attention_backend

O gerenciador de contexto attention_backend define temporariamente um backend de atenção para um modelo dentro do contexto. Fora do contexto, a atenção padrão (atenção de produto escalar com escala nativa do PyTorch) é utilizada. Isso é útil se você deseja usar backends diferentes para diferentes partes de um pipeline ou se deseja testar os diferentes 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]

A maioria dos backends de atenção suporta torch.compile sem quebras de grafo e pode ser usada para acelerar ainda mais a inferência.

Confiança em kernels remotos

Hub backends and other kernel-backed features (such as GGUF and Nunchaku Lite) download compute kernels from the Hub with kernels and execute their code locally.

By default, kernels only loads a kernel when its publisher is a trusted kernel publisher on the Hub. Kernels published under the kernels-community organization are trusted, so Diffusers loads them without any additional configuration. The _flash_3_hub, flash_hub, sage_hub, and the other Hub attention backends all resolve to kernels-community repositories.

Kernels de qualquer outro editor não são verificados. Carregar um deles baixa e executa código pelo qual o Diffusers não pode se responsabilizar, por isso o Diffusers o mantém desativado, a menos que você opte explicitamente por ativá-lo com a variável de ambiente DIFFUSERS_TRUST_REMOTE_KERNELS. Quando definida, o Diffusers encaminha trust_remote_code=True para kernels, permitindo que ele carregue kernels de editores não confiáveis também.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Ative isso apenas após inspecionar o repositório do kernel, pois isso concede ao código baixado a capacidade de ser executado na sua máquina. Sem isso, carregar um kernel de um editor não confiável gera um erro. O Diffusers executa essa verificação por conta própria, portanto, ela também se aplica a kernels<0.14.0, que antecede o argumento trust_remote_code. Definir DIFFUSERS_DISABLE_REMOTE_CODE=true desativa o código remoto globalmente e tem precedência sobre DIFFUSERS_TRUST_REMOTE_KERNELS.

Verificações

O despachador de atenção inclui verificações de depuração que capturam erros comuns antes que causem problemas.

1. As verificações de dispositivo confirmam que os tensores de consulta, chave e valor estão no mesmo dispositivo.
2. As verificações de tipo de dados confirmam que os tensores possuem dtypes correspondentes e usam bfloat16 ou float16.
3. As verificações de forma validam as dimensões dos tensores e impedem a mistura de máscaras de atenção com bandeiras causais.

Ative essas verificações definindo a variável de ambiente DIFFUSERS_ATTN_CHECKS. As verificações adicionam sobrecarga a cada operação de atenção, por isso estão desativadas por padrão.

export DIFFUSERS_ATTN_CHECKS=yes

As verificações são executadas agora antes de cada operação de atenção.

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

Você também pode configurar o registro diretamente.

from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True

Backends disponíveis

Consulte a tabela abaixo para obter a lista completa dos backends de atenção disponíveis e suas variantes.

Expandir

Nome do Backend	Família	Descrição
native	PyTorch nativo	Backend padrão usando scaled_dot_product_attention do PyTorch
flex	FlexAttention	Implementação do PyTorch FlexAttention
_native_cudnn	PyTorch nativo	Atenção otimizada pelo CuDNN
_native_efficient	PyTorch nativo	Atenção eficiente em memória
_native_flash	PyTorch nativo	FlashAttention do PyTorch
_native_math	PyTorch nativo	Atenção baseada em matemática (fallback)
_native_npu	PyTorch nativo	Atenção otimizada para NPU
_native_xla	PyTorch nativo	Atenção otimizada com XLA
flash	FlashAttention	FlashAttention-2
flash_hub	FlashAttention	FlashAttention-2 de kernels
flash_varlen	FlashAttention	FlashAttention de comprimento variável
flash_varlen_hub	FlashAttention	FlashAttention de comprimento variável a partir de kernels
aiter_fa2_hub	Motor de Tensor de IA para ROCm	FlashAttention-2 para AMD ROCm a partir de kernels
aiter_sage_hub	Motor de Tensor de IA para ROCm	Atenção quantizada (INT8 QK, FP8 PV) para AMD RDNA4 (gfx1201) a partir de kernels
flash_4_hub	FlashAttention	FlashAttention-4
_flash_3	FlashAttention	FlashAttention-3
_flash_varlen_3	FlashAttention	FlashAttention-3 de comprimento variável
_flash_3_hub	FlashAttention	FlashAttention-3 de kernels
_flash_3_varlen_hub	FlashAttention	FlashAttention-3 de comprimento variável a partir de kernels
sage	SageAttention	Atenção quantizada (INT8 QK)
sage_hub	SageAttention	Atenção quantizada (INT8 QK) a partir de kernels
sage_blackwell_hub	SageAttention	Atenção SageAttention3 FP4 para GPUs Blackwell SM120 a partir de kernels
sage_varlen	SageAttention	SageAttention de comprimento variável
_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	Atenção eficiente em memória