# Como o cache funciona

Imagine que você está tendo uma conversa com alguém e, em vez de lembrar do que a pessoa disse anteriormente, ela precisa recomeçar do zero a cada vez que você responde. Isso seria lento e ineficiente, certo?

Você pode estender essa analogia para modelos transformer. A geração em modelos autoregressivos pode ser lenta porque faz uma previsão um token de cada vez. Cada nova previsão depende de todo o contexto anterior.

Para prever o 1000º token, o modelo requer informações dos 999 tokens anteriores. Essas informações são representadas como multiplicações de matrizes entre as representações dos tokens.

Para prever o 1001º token, você precisa das mesmas informações dos 999 tokens anteriores, além de qualquer informação do 1000º token. Isso representa muitas multiplicações de matriz que o modelo precisa calcular repetidamente para cada token!

Um cache de chave-valor (KV) elimina essa ineficiência armazenando pares de kv derivados das camadas de atenção dos tokens processados anteriormente. Os pares de kv armazenados são recuperados do cache e reutilizados para os tokens subsequentes, evitando a necessidade de recálculo.

> [!WARNING]
> O uso de cache deve ser restrito à **inferência**. Ele pode causar erros inesperados se estiver habilitado durante o treinamento.

Para entender melhor como e por que o cache funciona, vamos examinar de perto a estrutura das matrizes de atenção.

## Matrizes de atenção

O **atendimento de produto escalar** é calculado conforme mostrado abaixo para um lote de tamanho `b`, número de cabeças de atenção `h`, comprimento da sequência até o momento `T` e dimensão por cabeça de atenção `d_head`.

$$ \text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V $$

As matrizes de consulta (`Q`), chave (`K`) e valor (`V`) são projeções dos embeddings de entrada, com forma `(b, h, T, d_head)`.

Para a atenção causal, a máscara impede que o modelo preste atenção a tokens futuros. Uma vez que um token é processado, sua representação nunca muda em relação aos tokens futuros, o que significa que $ K_{\text{past}} $ e $ V_{\text{past}} $ podem ser armazenados em cache e reutilizados para calcular a representação do último token.

$$ \text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{em cache}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{em cache}}, v_{t}]) $$

No momento da inferência, você só precisa da query do último token para calcular a representação $ x_t $ que prevê o próximo token $ t+1 $. A cada etapa, os novos vetores de chave e valor são **armazenados** no cache e **adicionados** às chaves e valores anteriores.

$$ K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t) $$

A atenção é calculada de forma independente em cada camada do modelo, e o cache é feito por camada.

Consulte a tabela abaixo para comparar como o cache melhora a eficiência.

| sem cache | com cache |
|---|---|
| para cada etapa, recalcular todos os `K` e `V` anteriores  | para cada etapa, calcule apenas o `K` e o `V` atuais |
| o custo de atenção por passo é **quadrático** em relação ao comprimento da sequência | o custo de atenção por passo é **linear** em relação ao comprimento da sequência (a memória cresce linearmente, mas o custo de computação por token permanece baixo) |

## Classe de cache

Uma interface básica de cache de KV recebe os tensores de chave e valor para o token atual e retorna os tensores `K` e `V` atualizados. Isso é gerenciado internamente pelo método `forward` de um modelo.

```py
new_K, new_V = cache.update(k_t, v_t, layer_idx)
attn_output = attn_layer_idx_fn(q_t, new_K, new_V)
```

Ao usar a classe [Cache](/docs/transformers/main/pt/internal/generation_utils#transformers.Cache) do Transformers, o módulo de autoatenção executa várias etapas críticas para integrar informações passadas e presentes.

1. O módulo de atenção concatena os pares kv atuais com os pares kv anteriores armazenados no cache. Isso cria pesos de atenção com a forma `(new_tokens_length, past_kv_length + new_tokens_length)`. Os pares kv atuais e anteriores são essencialmente combinados para calcular as pontuações de atenção, garantindo que o modelo esteja ciente do contexto anterior e da entrada atual.

2. Quando o método `forward` é chamado iterativamente, é crucial que a forma da attention mask corresponda ao comprimento combinado dos pares kv passados e atuais. A attention mask deve ter a forma `(batch_size, past_kv_length + new_tokens_length)`. Isso geralmente é tratado internamente em [generate()](/docs/transformers/main/pt/main_classes/text_generation#transformers.GenerationMixin.generate), mas, se você quiser implementar seu próprio loop de geração com [Cache](/docs/transformers/main/pt/internal/generation_utils#transformers.Cache), tenha isso em mente! A attention mask deve conter os valores dos tokens passados e atuais.

## Implementação de armazenamento em cache

Os caches são estruturados como uma lista de camadas, em que cada camada contém um cache de chave e um cache de valor. Os caches de chave e de valor são tensores com a forma `[batch_size, num_heads, seq_len, head_dim]`.

As camadas podem ser de diferentes tipos (por exemplo, `DynamicLayer`, `StaticLayer`, `StaticSlidingWindowLayer`), o que altera principalmente o modo como o comprimento da sequência é tratado e como o cache é atualizado.

A mais simples é uma `DynamicLayer` que cresce à medida que mais tokens são processados. A dimensão do comprimento da sequência (`seq_len`) aumenta a cada novo token:

```py
cache.layers[idx].keys = torch.cat([cache.layers[idx].keys, key_states], dim=-2)
cache.layers[idx].values = torch.cat([cache.layers[idx].values, value_states], dim=-2)
```

Outros tipos de camada, como `StaticLayer` e `StaticSlidingWindowLayer`, possuem um comprimento de sequência fixo que é definido quando o cache é criado. Isso os torna compatíveis com `torch.compile`. No caso de `StaticSlidingWindowLayer`, os tokens existentes são deslocados para fora do cache quando um novo token é adicionado.

O exemplo abaixo demonstra como criar um loop de geração com [DynamicCache](/docs/transformers/main/pt/internal/generation_utils#transformers.DynamicCache). Conforme discutido, a máscara de atenção é uma concatenação dos valores de tokens passados e atuais.

```py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
from accelerate import Accelerator

device = Accelerator().device

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map=device)
tokenizer = AutoTokenizer.from_pretrained(model_id)

past_key_values = DynamicCache(config=model.config)
messages = [{"role": "user", "content": "Hello, what's your name."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)

generated_ids = inputs.input_ids
max_new_tokens = 10

for _ in range(max_new_tokens):
    outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
    # Greedily sample one next token
    next_token_ids = outputs.logits[:, -1:].argmax(-1)
    generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
    # Prepare inputs for the next generation step by leaving unprocessed tokens, in our case we have only one new token
    # and expanding attn mask for the new token, as explained above
    attention_mask = inputs["attention_mask"]
    attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
    inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}

print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
"[INST] Hello, what's your name. [/INST]  Hello! My name is LLaMA,"
```