diffusers/main/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.

A maioria dos modelos transformer do Diffusers roteia a atenção por meio de um dispatcher de atenção, permitindo alternar entre backends otimizados por trás de uma única API. O dispatcher gerencia as implementações registradas e expõe um caminho de chamada unificado para elas. Alguns modelos, como muitos autoencoders, não utilizam o dispatcher devido à sua estrutura interna, e a troca do backend não afeta suas camadas de atenção.

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. O backend mais rápido depende do modelo, da GPU, do dtype e da forma da entrada.

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

Os backends do Hub (os nomes *_hub) precisam da biblioteca Kernels e baixam o kernel no primeiro uso. Os backends do Hub do SageAttention também precisam de DIFFUSERS_TRUST_REMOTE_KERNELS=true (veja Confiança em kernels remotos). Os outros backends precisam do próprio pacote, como flash-attn ou sageattention. A tabela Backends disponíveis lista os requisitos que o Diffusers verifica ao você habilitar um backend.

Defina um backend no modelo

O método set_attention_backend() percorre as camadas de atenção do modelo e aplica o backend escolhido em cada uma delas. Ele também define o backend ativo do dispatcher em todo o processo para o mesmo valor.

reset_attention_backend() limpa o backend apenas nas camadas de atenção. Ele não limpa o backend ativo em todo o processo. Para uma troca temporária que restaura o backend ativo anterior ao sair, use o gerenciador de contexto attention_backend.

O exemplo abaixo habilita _flash_3_hub (FlashAttention-3 do Hub) com 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]

Os backends de FlashAttention-3 não-Hub (_flash_3, _flash_varlen_3) exigem a compilação do FlashAttention-3 a partir do código-fonte e serão descontinuados em breve. Use _flash_3_hub ou _flash_3_varlen_hub no lugar deles.

Tente um backend temporariamente

O gerenciador de contexto attention_backend() define o backend ativo em todo o processo durante a duração do bloco e restaura o backend anterior quando o bloco é encerrado. Use-o para testar um backend em uma única chamada sem deixar um backend permanente aplicado a partir de 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]

A maioria dos backends de atenção funciona com torch.compile. Se isso acelera seu pipeline depende do modelo e do backend. Consulte Precisão e compilação.

Confiança em kernels remotos

Os backends do Hub precisam primeiro da biblioteca Kernels, e em seguida o Diffusers busca o kernel do Hub no primeiro uso.

Os backends de atenção do Hub baixam kernels de computação com a biblioteca Kernels e os executam localmente. A maioria dos nomes de atenção do Hub (_flash_3_hub, flash_hub e os outros backends *_hub) resolve para a organização kernels-community. Essa organização é um publicador confiável no Kernels, portanto o Diffusers carrega esses kernels de atenção sem definir DIFFUSERS_TRUST_REMOTE_KERNELS.

Os backends do SageAttention Hub (sage_hub e sage_blackwell_hub) são carregados a partir da organização SageAttention. Defina DIFFUSERS_TRUST_REMOTE_KERNELS=true para usá-los.

Outros recursos baseados em kernels, como GGUF e Nunchaku Lite, podem obter kernels de publicadores fora do kernels-community. Esses caminhos permanecem bloqueados, a menos que você opte por habilitá-los com DIFFUSERS_TRUST_REMOTE_KERNELS. Quando definido, o Diffusers encaminha trust_remote_code=True para o Kernels, permitindo que publicadores não confiáveis também sejam carregados.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Ative isso apenas após inspecionar o repositório do kernel. Sem ele, 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 pode executar verificações de depuração antes de cada chamada de atenção despachada. Quais verificações são executadas depende das restrições registradas para o backend ativo.

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, quando registradas, confirmam a correspondência dos dtypes e frequentemente exigem 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 as verificações com a variável de ambiente DIFFUSERS_ATTN_CHECKS. As verificações adicionam sobrecarga, portanto, estão desativadas por padrão.

export DIFFUSERS_ATTN_CHECKS=yes

Com as verificações ativadas, o Diffusers executa essas restrições antes de cada chamada de atenção despachada. O exemplo de baixo nível abaixo chama dispatch_attention_fn diretamente. A inferência de pipeline não precisa dessa importação. Ela só precisa do backend definido 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 disponíveis

Consulte a tabela abaixo para obter a lista completa dos backends de atenção disponíveis e suas variantes. O Diffusers verifica a disponibilidade do pacote e as versões fixadas quando você habilita um backend.

Nome do Backend	Família	Descrição	Pré-requisito
native	PyTorch nativo	Backend padrão usando scaled_dot_product_attention do PyTorch	None
flex	FlexAttention	PyTorch FlexAttention	torch>=2.5.0
_native_cudnn	PyTorch nativo	Atenção otimizada pelo CuDNN	CUDA + CuDNN
_native_efficient	PyTorch nativo	Atenção eficiente em memória	Nenhuma além do PyTorch
_native_flash	PyTorch nativo	FlashAttention do PyTorch	CUDA
_native_math	PyTorch nativo	Atenção baseada em matemática (fallback)	None
_native_npu	PyTorch nativo	Atenção otimizada para NPU	torch_npu
_native_xla	PyTorch nativo	Atenção otimizada com XLA	torch_xla>=2.2
flash	FlashAttention	FlashAttention-2	flash-attn>=2.6.3
flash_hub	FlashAttention	FlashAttention-2 dos kernels do Hub	kernels>=0.12
flash_varlen	FlashAttention	FlashAttention de comprimento variável	flash-attn>=2.6.3
flash_varlen_hub	FlashAttention	FlashAttention de comprimento variável dos kernels do Hub	kernels>=0.12
aiter_fa2_hub	Motor de Tensor de IA para ROCm	FlashAttention-2 para AMD ROCm a partir dos kernels do Hub (bfloat16)	kernels>=0.12, ROCm
flash_4_hub	FlashAttention	FlashAttention-4 dos kernels do Hub	kernels>=0.12.3
_flash_3	FlashAttention	FlashAttention-3 (local; em breve será descontinuado)	Compilar o FA3 a partir do código-fonte
_flash_varlen_3	FlashAttention	FlashAttention-3 de comprimento variável (local; será descontinuado em breve)	Compilar o FA3 a partir do código-fonte
_flash_3_hub	FlashAttention	FlashAttention-3 dos kernels do Hub	kernels>=0.12
_flash_3_varlen_hub	FlashAttention	FlashAttention-3 de comprimento variável dos kernels do Hub	kernels>=0.12
sage	SageAttention	Atenção quantizada (INT8 QK)	sageattention>=2.1.1
sage_hub	SageAttention	Atenção quantizada (INT8 QK) dos kernels do Hub	kernels>=0.12, DIFFUSERS_TRUST_REMOTE_KERNELS=true
sage_blackwell_hub	SageAttention	Atenção SageAttention3 FP4 para GPUs Blackwell SM120 a partir de kernels do Hub	kernels>=0.12, DIFFUSERS_TRUST_REMOTE_KERNELS=true
sage_varlen	SageAttention	SageAttention de comprimento variável	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	Atenção eficiente em memória	xformers>=0.0.29