# Backends de atenção

Todas as implementações de atenção executam o mesmo cálculo. Cada token é comparado com todos os outros tokens. A diferença está em *como* o cálculo é realizado. A atenção básica escala mal porque materializa a matriz de atenção completa na memória, criando gargalos que atrasam a inferência. As implementações otimizadas reorganizam a matemática para reduzir o tráfego de memória, resultando em inferências mais rápidas e acessíveis.

A [AttentionInterface](/docs/transformers/main/pt/internal/modeling_utils#transformers.AttentionInterface) fornece implementações de atenção otimizadas. Ela desacopla a implementação da atenção da implementação do modelo para simplificar a experimentação com diferentes funções. Adicione novos backends facilmente com essa interface consistente.

Backend de atençãoDescrição "flash_attention_3"melhora o FlashAttention-2 ao sobrepor operações e fundir as passadas de forward e backward de forma mais compacta "flash_attention_2"divide os cálculos em blocos menores e usa a memória rápida na placa "flex_attention"framework para especificar padrões de atenção personalizados (esparsos, locais por bloco, janela deslizante) sem escrever kernels de baixo nível manualmente "sdpa"implementação integrada do PyTorch para atenção de produto escalar "paged|flash_attention_3"Versão paginada do FlashAttention-3 "paged|flash_attention_2"Versão paginada do FlashAttention-2 "paged|sdpa"Versão paginada do SDPA "paged|eager"Versão paginada de eager

## Definir um backend de atenção

Use o argumento `attn_implementation` em [from_pretrained()](/docs/transformers/main/pt/main_classes/model#transformers.PreTrainedModel.from_pretrained) para instanciar um modelo com uma função de atenção específica.

```py
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="flash_attention_2"
)
```

Alterne entre backends de atenção em tempo de execução sem recarregar o modelo usando [set_attn_implementation()](/docs/transformers/main/pt/main_classes/model#transformers.PreTrainedModel.set_attn_implementation).

```py
model.set_attn_implementation("sdpa")
```

### Kernels

Baixe e carregue kernels de computação compilados diretamente do [Hub](https://huggingface.co/models?other=kernels) em tempo de execução com a biblioteca [Kernels](https://huggingface.co/docs/kernels/index). Isso evita problemas de empacotamento causados por versões incompatíveis de PyTorch ou CUDA.

Os kernels se registram automaticamente em [AttentionInterface](/docs/transformers/main/pt/internal/modeling_utils#transformers.AttentionInterface) ao serem detectados. Não é necessário instalar o pacote FlashAttention explicitamente. Solicitar FlashAttention pelo nome também recorre ao kernel do Hub, veja [FlashAttention fallback](./kernel_doc/loading_kernels#flashattention-fallback).

```py
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="kernels-community/flash-attn2"
)
```

### Gerenciador de contexto SDPA

O scaled dot product attention (SDPA) do PyTorch seleciona automaticamente a função de atenção mais rápida para backends CUDA. Para outros backends, ele usa por padrão a implementação em C++ do PyTorch.

Force o SDPA a usar uma implementação específica com o gerenciador de contexto [torch.nn.attention.sdpa_kernel](https://pytorch.org/docs/stable/generated/torch.nn.attention.sdpa_kernel.html).

```py
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="sdpa"
)

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    outputs = model.generate(**inputs)
```

## Atenção específica do backbone

Modelos multimodais usam backbones diferentes para cada modalidade. Otimize o desempenho atribuindo funções de atenção específicas a cada backbone. Alguns backbones de visão, por exemplo, têm melhor desempenho em fp32, o que o FlashAttention não suporta.

Mapeie os backbones de visão para diferentes funções de atenção com um dicionário, enquanto o backbone de texto continua a usar FlashAttention. As chaves na implementação de atenção devem corresponder aos nomes das subconfigurações.

```py
from transformers import AutoModelForImageTextToText

attention_implementation_per_backbone = {"vision_config": "sdpa", "text_config": "flash_attention_2"}

for key in attention_implementation_per_backbone:
    assert key in model.config.sub_configs, f"Invalid key in `attention_implementation`"

model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation=attention_implementation_per_backbone
)
```

Omita determinados backbones do dict para usar a função de atenção padrão (SDPA).

```py
model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation={"text_config": "flash_attention_2"}
)
```

Defina a mesma função de atenção para todos os backbones com uma única string.

```py
model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation="eager"
)
```

Defina a função de atenção globalmente com uma chave vazia.

```py
model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation={"": "eager"}
)
```

## Crie uma nova função de atenção

Personalize ou crie novas funções de atenção adicionando-as ao registro de atenção com [AttentionInterface.register()](/docs/transformers/main/pt/internal/modeling_utils#transformers.AttentionInterface.register). Os modelos utilizam essas funções por meio do argumento `attn_implementation`.

> [!WARNING]  
> Registre uma função de máscara de atenção correspondente ao registrar uma função de atenção personalizada. Se o nome da `attn_implementation` personalizada não estiver registrado em [AttentionMaskInterface](/docs/transformers/main/pt/internal/modeling_utils#transformers.AttentionMaskInterface), o Transformers ignora a criação da máscara e passa `attention_mask=None` para as camadas de atenção. Sua função de atenção deve lidar com as restrições causais, de preenchimento, de empacotamento ou de janela deslizante por conta própria, caso contrário, essas restrições podem ser descartadas silenciosamente.

Este exemplo personaliza a função de atenção para imprimir uma mensagem para cada camada. Ele mantém a máscara na implementação original, registrando `masking_utils.sdpa_mask` como a função de máscara de atenção.

```python
import torch
from transformers import AutoModelForCausalLM, AttentionInterface, AttentionMaskInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.masking_utils import sdpa_mask

def my_new_sdpa(*args, **kwargs):
    print("I just entered the attention computation")
    return sdpa_attention_forward(*args, **kwargs)

AttentionInterface.register("my_new_sdpa", my_new_sdpa)
AttentionMaskInterface.register("my_new_sdpa", sdpa_mask)  # must have the same name as the registered attention function

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", attn_implementation="my_new_sdpa")
model(torch.ones(1, 5, dtype=int))
```

Você também pode adicionar novos argumentos à função de atenção. Os modelos que suportam [AttentionInterface](/docs/transformers/main/pt/internal/modeling_utils#transformers.AttentionInterface) propagam kwargs para as camadas de atenção e para a função de atenção. Passe os argumentos como kwargs na função forward do modelo. Funções de atenção personalizadas devem seguir essa assinatura e formato de retorno.

```python
import torch
from transformers import AutoModelForCausalLM, AttentionInterface, AttentionMaskInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.masking_utils import sdpa_mask

def custom_attention(
    module: torch.nn.Module,  # required arg
    query: torch.Tensor,  # required arg
    key: torch.Tensor,  # required arg
    value: torch.Tensor,  # required arg
    attention_mask: Optional[torch.Tensor],  # required arg
    a_new_kwargs = None,  # You can now add as many kwargs as you need
    another_new_kwargs = None,  # You can now add as many kwargs as you need
    **kwargs,  # You need to accept **kwargs as models will pass other args
) -> tuple[torch.Tensor, Optional[torch.Tensor]]
    ...  # do your magic!
    return attn_output, attn_weights  # attn_weights are optional here

AttentionInterface.register("custom", custom_attention)
AttentionMaskInterface.register("custom", sdpa_mask)  # to leave the existing mask untouched

model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="custom")
model(torch.ones(1, 5, dtype=int), a_new_kwargs=..., another_new_kwargs=...)
```

Verifique o [código de modelagem](https://github.com/huggingface/transformers/tree/main/src/transformers/models) de um modelo para confirmar quais argumentos e kwargs ele envia para a função de atenção.

## AttentionMaskInterface

[AttentionMaskInterface](/docs/transformers/main/pt/internal/modeling_utils#transformers.AttentionMaskInterface) é o registro consultado pelas funções [`create_*_mask`](#build-an-attention-mask) para converter uma máscara no formato esperado pelo backend de atenção ativo. O FlexAttention requer uma [BlockMask](https://docs.pytorch.org/docs/stable/nn.attention.flex_attention.html#torch.nn.attention.flex_attention.BlockMask), o SDPA requer um tensor 4D e o FlashAttention requer a máscara de padding base 2D. Registre um backend personalizado ou substitua o formatador de um backend existente com [AttentionMaskInterface.register()](/docs/transformers/main/pt/internal/modeling_utils#transformers.AttentionInterface.register).

```python
import torch
from transformers import AttentionMaskInterface
from transformers.masking_utils import sdpa_mask

def my_new_sdpa_mask(*args, **kwargs):
    print("I just entered the attention mask computation")
    return sdpa_mask(*args, **kwargs)

AttentionMaskInterface.register("my_new_sdpa_mask", my_new_sdpa_mask)
```

Sem um formatador registrado para o `attn_implementation` ativo, a criação da máscara é ignorada e `attention_mask=None` é passado para as camadas de atenção.

Funções registradas devem corresponder a esta assinatura.

```python
def custom_attention_mask(
    batch_size: int,  # required arg
    q_length: int,  # required arg
    kv_length: int,  # required arg
    q_offset: int = 0,  # required arg
    kv_offset: int = 0,  # required arg
    mask_function: Callable = causal_mask_function,  # required arg
    attention_mask: Optional[torch.Tensor] = None,  # required arg
    **kwargs,  # a few additional args may be passed as kwargs, especially the model's config is always passed
) -> Optional[torch.Tensor]:
```

O argumento `mask_function` é um `Callable` que imita as funções [mask_mod](https://pytorch.org/blog/flexattention/) do PyTorch. Ele recebe 4 índices `(batch_idx, head_idx, q_idx, kv_idx)` e retorna um booleano indicando se aquela posição contribui para o cálculo de atenção. Essa é a mesma forma primitiva usada por `or_mask_function` e `and_mask_function` em [Construir uma máscara de atenção](#build-an-attention-mask).

> [!TIP]
> Use este [workaround](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/executorch.py) para o torch.export caso o `mask_function` falhe ao criar uma máscara.

## Construir uma máscara de atenção

Construa máscaras de atenção com as funções `create_*_mask` em [transformers.masking_utils](https://github.com/huggingface/transformers/blob/main/src/transformers/masking_utils.py#L894). Cada função lê o backend de atenção ativo da configuração do modelo, consulta o formatador de máscara do backend em [AttentionMaskInterface](/docs/transformers/main/pt/internal/modeling_utils#transformers.AttentionMaskInterface) e retorna o formato que o backend espera. Você não precisa inverter, expandir ou converter o tipo da máscara por conta própria.

Escolha a função que corresponde ao padrão de atenção.

| function | caso de uso |
|---|---|
| `create_causal_mask` | modelos somente de decodificador em que cada token presta atenção a si mesmo e aos tokens anteriores |
| `create_bidirectional_mask` | modelos de codificador, ou atenção cruzada de um decodificador para os estados do codificador |
| `create_sliding_window_causal_mask` | modelos de decodificador com um padrão de atenção de janela deslizante |
| `create_chunked_causal_mask` | modelos de decodificador que dividem a sequência em blocos de tamanho fixo |
| `create_bidirectional_sliding_window_mask` | modelos de codificador com um padrão de atenção de janela deslizante |

> [!WARNING]
> Os helpers legados de máscara chamáveis - `get_extended_attention_mask`, `create_extended_attention_mask_for_decoder`, `invert_attention_mask` - emitem um aviso de deprecação e serão removidos em uma versão futura. Use as funções `create_*_mask` em vez disso.

Chame `create_causal_mask` dentro de uma passagem de forward do decoder. Passe a config, os embeddings de entrada, a `attention_mask` 2D fornecida pelo usuário e o cache. A função usa os embeddings para ler o tamanho do lote, o comprimento da query, o dtype e o device, e usa o cache para calcular o comprimento da key.

```py
from transformers.masking_utils import create_causal_mask

attention_mask = create_causal_mask(
    config=self.config,
    inputs_embeds=inputs_embeds,
    attention_mask=attention_mask,
    past_key_values=past_key_values,
)
```

Chame `create_bidirectional_mask` para a autoatenção do codificador. Remova `past_key_values`, pois os codificadores não fazem cache.

```py
from transformers.masking_utils import create_bidirectional_mask

attention_mask = create_bidirectional_mask(
    config=self.config,
    inputs_embeds=embedding_output,
    attention_mask=attention_mask,
)
```

Para a atenção cruzada, passe os estados do codificador como `encoder_hidden_states` para que a máscara utilize o comprimento das chaves e valores do codificador em vez do comprimento das consultas do decodificador.

```py
encoder_attention_mask = create_bidirectional_mask(
    config=self.config,
    inputs_embeds=embedding_output,
    attention_mask=encoder_attention_mask,
    encoder_hidden_states=encoder_hidden_states,
)
```

Adicione restrições extras por cima da máscara base com os argumentos `or_mask_function` e `and_mask_function`. Use `or_mask_function` para permitir que posições adicionais participem da atenção e `and_mask_function` para restringir ainda mais o padrão base. Ambos seguem a assinatura de 4 índices de `mask_function` descrita em [AttentionMaskInterface](#attentionmaskinterface). Eles recebem `(batch_idx, head_idx, q_idx, kv_idx)` e retornam um booleano.

> [!WARNING]
> `or_mask_function` e `and_mask_function` podem expressar qualquer padrão de atenção, mas são mais lentos que os padrões integrados e não são compatíveis com ExecuTorch. O overhead é mais perceptível em modelos menores (~200M parâmetros), onde a criação da máscara ocupa uma proporção maior do tempo do forward pass. Utilize-os apenas quando as funções padrão `create_*_mask` não conseguem expressar o que você precisa.

Por exemplo, sobreponha uma função que retorna `True` em todos os pontos a uma máscara causal para transformá-la em uma totalmente bidirecional. A união com o padrão causal permite que cada token preste atenção a todos os outros tokens.

```py
mask_kwargs = {
    "config": self.config,
    "inputs_embeds": inputs_embeds,
    "attention_mask": attention_mask,
    "past_key_values": past_key_values,
    "position_ids": position_ids,
    "or_mask_function": lambda *args: torch.tensor(True, dtype=torch.bool),
}

attention_mask = create_causal_mask(**mask_kwargs)
```

Durante a geração, [generate()](/docs/transformers/main/pt/main_classes/text_generation#transformers.GenerationMixin.generate) constrói máscaras por meio de `create_masks_for_generate`, que direciona para o `create_*_mask` adequado com base na configuração do modelo. Sobrescreva-o em uma classe de modelo para incorporar uma estratégia de mascaramento personalizada para a geração.

## Passe uma máscara de atenção 4D personalizada

Use uma máscara 4D própria quando precisar de um padrão de atenção que as funções `create_*_mask` não conseguem expressar. Uma máscara 4D tem a forma `(batch_size, 1, query_length, kv_length)`, em que `1` faz o broadcast da mesma máscara para todas as cabeças de atenção. O Transformers detecta e usa a máscara como está, pulando `create_*_mask`.

Uma máscara 4D usa uma de duas convenções de valores.

| dtype | participar | mascarar |
|---|---|---|
| boolean | `True` | `False` |
| float | `0.0` | `-inf` |

A convenção float adiciona a máscara aos scores de atenção antes do softmax. Um score somado a `0.0` permanece inalterado, de modo que a posição contribui. Um score somado a `-inf` cai para zero após o softmax, de modo que a posição é excluída.

> [!IMPORTANT]
> A convenção aceita depende do backend de atenção. `sdpa` aceita uma máscara booleana ou de ponto flutuante. `eager` adiciona a máscara aos scores, portanto, aceita apenas uma máscara de ponto flutuante. `flash_attention_2` e `flex_attention` consomem seus próprios formatos (uma máscara de padding 2D e um [BlockMask](https://docs.pytorch.org/docs/stable/nn.attention.flex_attention.html#torch.nn.attention.flex_attention.BlockMask)) e não aceitam uma máscara 4D bruta.

Um erro comum é reutilizar a convenção `1`/`0` de uma máscara de padding 2D em uma máscara 4D em ponto flutuante. Como a máscara é somada às pontuações, `0.0` preserva uma posição e `1.0` apenas adiciona um pequeno viés.

O exemplo abaixo contrasta uma máscara incorreta com uma correta. Ambas partem do mesmo padrão causal `1`/`0`.

```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", attn_implementation="sdpa")
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")

input_ids = tokenizer("my favorite condiment on a", return_tensors="pt").input_ids
seq_len = input_ids.shape[1]

# 1 attends, 0 masks
causal = torch.tril(torch.ones(seq_len, seq_len))

# wrong: 1.0/0.0 floats are added to the scores, so 0.0 keeps a token and 1.0 barely changes it
wrong_mask = causal[None, None]

# correct: 0.0 attends, -inf masks
correct_mask = torch.where(causal.bool(), 0.0, float("-inf"))[None, None]
```

A máscara incorreta mantém todas as posições, pois `0.0` é o valor que mascara, de modo que nada é excluído.

```text
        wrong_mask                          correct_mask
   (1 attends, 0 masks)              (0 attends, -inf masks)

      k0 k1 k2 k3 k4                     k0   k1   k2   k3   k4
   q0  1  0  0  0  0                  q0  0  -inf -inf -inf -inf
   q1  1  1  0  0  0                  q1  0   0   -inf -inf -inf
   q2  1  1  1  0  0                  q2  0   0    0   -inf -inf
   q3  1  1  1  1  0                  q3  0   0    0    0   -inf
   q4  1  1  1  1  1                  q4  0   0    0    0    0
```

## Atenção bidirecional

Modelos apenas de decodificador usam atenção causal (unidirecional) por padrão, na qual cada token atende apenas a si mesmo e aos tokens anteriores. Defina `is_causal=False` para alternar para atenção bidirecional, na qual cada token atende a todos os outros tokens. Isso permite usar modelos apenas de decodificador como codificadores de texto, por exemplo, para gerar embeddings.

> [!NOTE]
> Isso só funciona para modelos causais (decodificadores). Não transforma modelos codificadores em modelos decodificadores.

Defina `is_causal=False` na configuração do modelo para tornar a atenção bidirecional o padrão em cada passagem de frente.

```py
from transformers import AutoModel, AutoConfig

config = AutoConfig.from_pretrained("meta-llama/Llama-3.2-1B")
config.is_causal = False

model = AutoModel.from_pretrained("meta-llama/Llama-3.2-1B", config=config)

# all forward passes now use bidirectional attention
outputs = model(**inputs)
```

Passe `is_causal` na chamada de forward em vez de na configuração do modelo para alternar entre atenção causal e bidirecional sem carregar o modelo duas vezes. O kwarg sobrescreve temporariamente a configuração e é restaurado após a chamada.

```py
from transformers import AutoModel

model = AutoModel.from_pretrained("meta-llama/Llama-3.2-1B")

# run with bidirectional attention
outputs = model(**inputs, is_causal=False)

# run with default causal attention
outputs = model(**inputs)
```