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

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 (a atenção de produto escalar 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.