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.

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.