# Как работает кэширование

Представьте, что вы ведёте беседу, и собеседник каждый раз, когда вы отвечаете, начинает с чистого листа, не помня, что вы говорили ранее. Разве это не было бы медленно и неэффективно?

Эту аналогию можно распространить на трансформерные модели. Генерация в авто-регрессионных моделях может быть медленной, поскольку предсказание выполняется по одному токену за раз. Каждое новое предсказание зависит от всего предыдущего контекста.

Для предсказания 1000-го токена модели требуется информация из предыдущих 999 токенов. Эта информация представлена в виде матричных умножений по представлениям токенов.

Чтобы предсказать 1001-й токен, вам нужна та же самая информация из предыдущих 999 токенов, а также любая информация из 1000-го токена. Это огромное количество матричных умножений, которые модели приходится вычислять снова и снова для каждого токена!

Кэш ключей и значений (KV-кэш) устраняет эту неэффективность, сохраняя пары kv, полученные из слоев внимания для ранее обработанных токенов. Сохраненные пары kv извлекаются из кэша и переиспользуются для последующих токенов, что исключает необходимость повторного вычисления.

> [!WARNING]
> Кэширование следует использовать только для **inference**. Его включение во время обучения может привести к непредвиденным ошибкам.

Чтобы лучше понять, как и почему работает кэширование, давайте подробнее рассмотрим структуру матриц внимания.

## Матрицы внимания

Механизм **масштабированного внимания на основе скалярного произведения** вычисляется следующим образом для батча размера `b`, количества головок внимания `h`, текущей длины последовательности `T` и размерности каждой головки внимания `d_head`.

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

Матрицы запроса (`Q`), ключа (`K`) и значения (`V`) представляют собой проекции входных эмбеддингов с формой `(b, h, T, d_head)`.

Для каузального внимания маска не позволяет модели обращать внимание на будущие токены. Как только токен обработан, его представление больше не изменяется относительно будущих токенов, что означает, что $ K_{\text{past}} $ и $ V_{\text{past}} $ можно кэшировать и повторно использовать для вычисления представления последнего токена.

$$ \text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{кэшировано}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{кэшировано}}, v_{t}]) $$

На этапе инференса для вычисления представления $ x_t $, предсказывающего следующий токен $ t+1 $, достаточно запроса последнего токена. На каждом шаге новые векторы ключей и значений **сохраняются** в кэше и **добавляются** к прошлым ключам и значениям.

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

Внимание вычисляется независимо в каждом слое модели, а кэширование выполняется для каждого слоя отдельно.

См. таблицу ниже, чтобы сравнить, как кэширование повышает эффективность.

| без кэширования | с кэшированием |
|---|---|
| на каждом шаге пересчитывать все предыдущие `K` и `V`  | на каждом шаге вычисляются только текущие `K` и `V` |
| стоимость внимания на каждом шаге **квадратично** зависит от длины последовательности | стоимость внимания на каждом шаге **линейна** относительно длины последовательности (память растёт линейно, но вычислительная стоимость на токен остаётся низкой) |

## Класс кэша

Базовый интерфейс кэша KV принимает тензоры ключей и значений для текущего токена и возвращает обновлённые тензоры `K` и `V`. Это управляется внутренним методом `forward` модели.

```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)
```

При использовании класса [Cache](/docs/transformers/main/ru/internal/generation_utils#transformers.Cache) в Transformers модуль self-attention выполняет несколько ключевых шагов для интеграции прошлой и текущей информации.

1. Модуль внимания объединяет текущие пары kv с прошлыми парами kv, хранящимися в кэше. Это создаёт веса внимания с формой `(new_tokens_length, past_kv_length + new_tokens_length)`. Текущие и прошлые пары kv фактически объединяются для вычисления оценок внимания, что обеспечивает осведомлённость модели о предыдущем контексте и текущем вводе.

2. Когда метод `forward` вызывается итеративно, важно, чтобы форма маски внимания совпадала с общей длиной прошлых и текущих пар kv. Маска внимания должна иметь форму `(batch_size, past_kv_length + new_tokens_length)`. Обычно это обрабатывается внутри [generate()](/docs/transformers/main/ru/main_classes/text_generation#transformers.GenerationMixin.generate), но если вы хотите реализовать собственный цикл генерации с [Cache](/docs/transformers/main/ru/internal/generation_utils#transformers.Cache), имейте это в виду! Маска внимания должна содержать значения прошлых и текущих токенов.

## Реализация хранилища кэша

Кэши структурированы как список слоёв, где каждый слой содержит кэш ключей и кэш значений. Кэши ключей и значений — это тензоры с формой `[batch_size, num_heads, seq_len, head_dim]`.

Слои могут быть разных типов (например, `DynamicLayer`, `StaticLayer`, `StaticSlidingWindowLayer`), что в основном влияет на обработку длины последовательности и обновление кэша.

Самый простой вариант — это `DynamicLayer`, который расширяется по мере обработки новых токенов. Размер измерения длины последовательности (`seq_len`) увеличивается с каждым новым токеном:

```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)
```

Другие типы слоёв, такие как `StaticLayer` и `StaticSlidingWindowLayer`, имеют фиксированную длину последовательности, которая задаётся при создании кэша. Это делает их совместимыми с `torch.compile`. В случае `StaticSlidingWindowLayer` существующие токены вытесняются из кэша при добавлении нового токена.

Пример ниже демонстрирует, как создать цикл генерации с [DynamicCache](/docs/transformers/main/ru/internal/generation_utils#transformers.DynamicCache). Как обсуждалось, маска внимания представляет собой конкатенацию значений прошлых и текущих токенов.

```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,"
```