Стратегии кэширования

Ключевые и значения (KV) векторы используются для вычисления оценок внимания. Для авторегрессионных моделей оценки KV вычисляются каждый раз, поскольку модель предсказывает по одному токену за раз. Каждое предсказание зависит от предыдущих токенов, что означает, что модель выполняет одни и те же вычисления каждый раз.

Кэш KV хранит эти вычисления, чтобы их можно было переиспользовать без повторного пересчёта. Эффективное кэширование критически важно для оптимизации производительности модели, поскольку оно сокращает время вычислений и повышает скорость отклика. Более подробное объяснение работы кэша см. в документации Caching.

Transformers предлагает несколько классов Cache, реализующих различные механизмы кэширования. Некоторые из этих классов Cache оптимизированы для экономии памяти, тогда как другие разработаны для максимальной скорости генерации. Используйте таблицу ниже для сравнения типов кэша и выбора оптимального варианта для вашей задачи.

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

Кэш по умолчанию

Класс DynamicCache является классом кэша по умолчанию для всех моделей. Он позволяет динамически увеличивать размер кэша для хранения растущего количества ключей и значений по мере продвижения генерации.

Обратите внимание, что для моделей, использующих скользящее окно внимания (Mistral, Gemma2,…) или блочное внимание (Llama4), кэш перестанет расти, когда слои, использующие эти типы внимания, достигнут своего максимального размера (размера скользящего окна или блока).

Отключите кэш, задав use_cache=False в generate().

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

model.generate(**inputs, do_sample=False, max_new_tokens=20, use_cache=False)

Классы кэша также можно инициализировать заранее, а затем передать их в параметр past_key_values модели. Это может быть полезно для более тонкого управления или для более продвинутых сценариев использования, таких как кэширование контекста.

В большинстве случаев проще определить стратегию кэширования в параметре cache_implementation.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

past_key_values = DynamicCache(config=model.config)
out = model.generate(**inputs, do_sample=False, max_new_tokens=20, past_key_values=past_key_values)

Кэш фиксированного размера

По умолчанию DynamicCache не позволяет использовать большинство оптимизаций just-in-time (JIT), так как размер кэша не фиксирован. Оптимизации JIT позволяют минимизировать задержку ценой увеличения потребления памяти. Все перечисленные ниже типы кэша совместимы с оптимизациями JIT, такими как torch.compile, для ускорения генерации.

Кэш фиксированного размера (StaticCache) предварительно выделяет память под конкретный максимальный размер кэша для пар kv. Вы можете генерировать токены вплоть до максимального размера кэша, не изменяя его. Однако фиксированный (обычно большой) размер состояний ключей/значений означает, что в процессе генерации многие токены фактически маскируются, так как они не должны участвовать в вычислении внимания. Таким образом, этот подход позволяет легко compile стадию декодирования, но приводит к потере вычислительных ресурсов на токены в вычислении внимания. Как и во всём, здесь есть компромисс: он будет очень хорош, если вы генерируете несколько последовательностей примерно одинаковой длины, но может оказаться неоптимальным, если, например, у вас есть одна очень длинная последовательность, а затем только короткие (поскольку фиксированный размер кэша будет большим, для коротких последовательностей будет потрачено много лишнего). Убедитесь, что вы понимаете последствия, если используете этот метод!

Что касается DynamicCache, обратите внимание, что для моделей, использующих скользящее окно внимания (Mistral, Gemma2,…) или блочное внимание (Llama4), кэш никогда не будет больше размера скользящего окна/блока на слоях, использующих эти типы внимания, даже если указанная максимальная длина больше.

Включить StaticCache можно, задав cache_implementation="static" в generate(). Это также активирует автоматическую compilation этапа декодирования для стратегий жадного и выборочного декодирования.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="static")
tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Hello, my name is [Your Name], and I am a [Your Profession] with [Number of Years] of"

Держите тензоры генерации на CPU

Компиляторные бэкенды, такие как Neuron и TPU, преобразуют вашу модель в фиксированный граф вычислений. Цикл генерации поддерживает тензоры (выходная последовательность, attention_mask, position_ids), которые увеличиваются на один токен на каждом шаге. Компилятор пересобирает граф каждый раз, когда форма этих тензоров меняется на ускорителе, что замедляет генерацию.

generate() перемещает на устройство модели только те тензоры, которые потребляет forward, непосредственно перед каждым вызовом forward. Результат перемещается обратно, чтобы соответствовать устройству входных данных. Передавайте входные данные на CPU, чтобы избежать накопления тензоров на ускорителе в ходе выполнения цикла. Скомпилированный граф остаётся стабильным, а поскольку выход следует за устройством входа, сгенерированный результат остаётся на CPU.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", device_map="auto")

# leave the inputs on CPU instead of calling .to(model.device).
inputs = tokenizer("The French Bread Law states", return_tensors="pt")

# generate runs forward on the model device but returns the output on the input_ids device.
output = model.generate(**inputs, do_sample=False, max_new_tokens=20)
print(output.device)
cpu

Кэширование с выгрузкой

Кэш KV может занимать значительную часть памяти и становиться узким местом при генерации длинного контекста. Эффективные по памяти кэши делают ставку на снижение потребления памяти в ущерб скорости. Это особенно важно для больших языковых моделей (LLM) и в случаях, когда ваше оборудование ограничено по объёму памяти.

Размещение кэша на CPU экономит память GPU, перемещая KV-кэш всех слоёв модели, кроме одного, на CPU. Во время forward-прохода модели по слоям на GPU поддерживается только кэш текущего слоя. Система асинхронно предзагружает кэш следующего слоя и возвращает кэш текущего слоя на CPU после вычисления внимания.

Если у вас небольшая GPU и вы получаете ошибки нехватки памяти (OOM), стоит рассмотреть возможность offloading.

Вы можете заметить небольшое снижение пропускной способности генерации по сравнению с полным кэшем на устройстве, в зависимости от вашей модели и параметров генерации (размер контекста, количество генерируемых токенов, количество лучей и т. д.). Это связано с тем, что перемещение состояний ключей/значений туда и обратно требует определённых вычислительных затрат.

Функция offloading доступна как для DynamicCache, так и для StaticCache. Вы можете включить её, задав cache_implementation="offloaded" для динамической версии или cache_implementation="offloaded_static" для статической версии в GenerationConfig или generate(). Кроме того, вы можете создать собственный экземпляр DynamicCache или StaticCache с опцией offloading=True и передать этот кэш в generate или метод forward вашей модели (например, past_key_values=DynamicCache(config=model.config, offloading=True) для динамического кэша).

Обратите внимание, что два класса Cache, упомянутые выше, имеют дополнительный параметр при их прямом создании — offload_only_non_sliding. Этот дополнительный аргумент определяет, будут ли также выгружены слои, использующие скользящее окно/чанковое внимание (если таковые имеются). Поскольку эти слои обычно и так короткие, возможно, лучше не выгружать их, так как выгрузка может привести к снижению скорости. По умолчанию этот параметр установлен в False для DynamicCache и в True для StaticCache.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
inputs = tokenizer("Fun fact: The shortest", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=23, cache_implementation="offloaded")
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.

Пример ниже показывает, как можно переключиться на выгруженный кэш, если память закончится:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from accelerate import Accelerator

def resilient_generate(model, *args, **kwargs):
    oom = False
    device = Accelerator().device
    torch_device_module = getattr(torch, device, torch.cuda)
    try:
        return model.generate(*args, **kwargs)
    except torch.OutOfMemoryError as e:
        print(e)
        print("retrying with cache_implementation='offloaded'")
        oom = True
    if oom:
        torch_device_module.empty_cache()
        kwargs["cache_implementation"] = "offloaded"
        return model.generate(*args, **kwargs)

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
prompt = ["okay "*1000 + "Fun fact: The most"]
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
beams = { "num_beams": 40, "num_return_sequences": 20, "max_new_tokens": 23, "early_stopping": True, }
out = resilient_generate(model, **inputs, **beams)
responses = tokenizer.batch_decode(out[:,-28:], skip_special_tokens=True)

Кэш с квантованием

Класс QuantizedCache снижает требования к памяти за счёт квантования значений KV до более низкой точности. В настоящее время QuantizedCache поддерживает три бэкенда квантования:

• hqq поддерживает типы данных int1, int2, int3, int4 и int8.
• quanto поддерживает типы данных int2 и int4. Это квантование по умолчанию.
• fp8 хранит значения KV в формате 8-битного числа с плавающей запятой (float8_e4m3fn). Он не требует дополнительных зависимостей и выполняет квантование на уровне тензора, а не на уровне группы, поэтому отклоняет параметры nbits, axis_key, axis_value, q_group_size и residual_length.

Квантование кэша может ухудшить задержку, если длина контекста небольшая и в памяти GPU достаточно места для генерации без включения квантования кэша. Постарайтесь найти баланс между эффективностью использования памяти и задержкой.

Включите QuantizedCache, задав cache_implementation="quantized" в GenerationConfig. Бэкенд квантования, а также дополнительные параметры квантования указываются в словаре cache_config, передаваемом в generate(). Используйте значения по умолчанию для этих дополнительных параметров, если только вы не сталкиваетесь с нехваткой памяти. В этом случае рассмотрите возможность уменьшения длины остатка.

Для бэкенда hqq мы рекомендуем установить параметры axis_key и axis_value в значение 1.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "hqq",
        "axis_key": 1,
        "axis_value": 1
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

Для бэкенда quanto мы рекомендуем установить параметры axis_key и axis_value в значение 0.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "quanto",
        "axis_key": 0,
        "axis_value": 0
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

Бэкенд fp8 не принимает дополнительных параметров и вызывает ошибку, если ему передан один из групповых.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Instruct-2507")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B-Instruct-2507", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="quantized", cache_config={"backend": "fp8"})
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it is powerful and emotional. I also like classical music because it is elegant and sophisticated. I enjoy

Кэш энкодер-декодера

EncoderDecoderCache предназначен для моделей с архитектурой «энкодер-декодер». Он управляет кэшами для само-внимания и кросс-внимания, обеспечивая хранение и извлечение предыдущих пар kv. Для энкодера и декодера можно независимо задать разные типы кэша.

Этот тип кэша не требует никакой настройки. Это простой обёрточный слой вокруг двух Cache, описанных выше, которые будут использоваться моделью независимо и напрямую.

Кэши, специфичные для модели

Некоторые модели имеют уникальный способ хранения прошлых kv-пар или состояний, который несовместим с другими классами кэша.

Модели Mamba, такие как Mamba, требуют специального кэша, поскольку в них отсутствует механизм внимания и состояния kv. В связи с этим они несовместимы с перечисленными выше классами Cache.

Итеративная генерация

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

Для итеративной генерации с кэшем начните с инициализации пустого класса кэша, после чего можно передавать новые промпты. Ведите учёт истории диалога с помощью шаблона чата.

В следующем примере показано использование Llama-2-7b-chat-hf. Если вы используете другую модель чат-формата, apply_chat_template() может обрабатывать сообщения иначе. В зависимости от того, как написан шаблон Jinja, он может отбрасывать важные токены. Для мультимодальных чат-моделей см. руководство iterative chatting with cache, чтобы узнать, как обрабатывать изображения или аудио.

Например, некоторые модели используют специальные токены <think> ... </think> в процессе рассуждений. При повторном кодировании они могут потеряться, что приведёт к проблемам с индексацией. Возможно, вам потребуется вручную удалить или скорректировать лишние токены в завершениях, чтобы сохранить стабильность.

import torch
from transformers import AutoTokenizer,AutoModelForCausalLM, DynamicCache, StaticCache

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

user_prompts = ["Hello, what's your name?", "Btw, yesterday I was on a rock concert."]

past_key_values = DynamicCache(config=model.config)

messages = []
for prompt in user_prompts:
    messages.append({"role": "user", "content": prompt})
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
    input_length = inputs["input_ids"].shape[1]
    outputs = model.generate(**inputs, do_sample=False, max_new_tokens=256, past_key_values=past_key_values)
    completion = tokenizer.decode(outputs[0, input_length: ], skip_special_tokens=True)
    messages.append({"role": "assistant", "content": completion})

Обрезка кэша

Используйте crop(), чтобы откатить токены из кэша, например, когда кандидатное продолжение отклоняется при итеративной генерации. Передайте количество удаляемых токенов в виде отрицательного целого числа.

# Remove the last three tokens from the cache.
past_key_values.crop(-3)

Слои со скользящим окном и линейным вниманием по мере работы отбрасывают прошлые состояния, поэтому по умолчанию откатываться некуда. Вызовите activate_past_recording() перед прямым проходом, чтобы они сохраняли эти состояния. Без этого crop вызывает RuntimeError, как только слой со скользящим окном заполняет своё окно, а для слоёв с линейным вниманием — сразу.

# Enable this before the forward pass if you need to roll back tokens.
past_key_values.activate_past_recording()

crop(0) не очищает кэш. Он не удаляет токены из закэшированной последовательности и является операцией без эффекта для слоёв полного внимания. Для слоёв скользящего окна и линейного внимания он может отбросить закэшированную информацию, которая не будет использована в следующем вызове модели. Это снижает потребление памяти без отката токенов.

                    cached positions
                    ┌───┬───┬───┬───┬───┬───┐
before crop(-3)     │ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │
                    └───┴───┴───┴───┴───┴───┘
                                ╰─────┬─────╯
                                 rolled back

after crop(-3)
                    ┌───┬───┬───┐
full attention      │ 0 │ 1 │ 2 │    keeps the whole prefix
                    └───┴───┴───┘
                            ┌───┐
sliding_window=2            │ 2 │    keeps sliding_window - 1
                            └───┘

Оба слоя теперь сообщают о длине последовательности 3, и генерация возобновляется с позиции 3.

Передача 0 не удаляет токены из последовательности. Для слоёв полного внимания это операция без эффекта, но для слоёв скользящего окна и линейного внимания она всё равно отбрасывает записанные состояния, которые не будут прочитаны при следующем проходе вперёд, что освобождает память без отката.

Предзаполнение кэша (кэширование префикса)

В некоторых ситуациях вы можете заполнить Cache парами kv для определённого префиксного промпта и повторно использовать его для генерации различных последовательностей.

Пример ниже инициализирует StaticCache, а затем кэширует начальный промпт. Теперь вы можете сгенерировать несколько последовательностей из предзаполненного промпта.

import copy
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache

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

# Init StaticCache with big enough max-length (1024 tokens for the below example)
# You can also init a DynamicCache, if that suits you better
prompt_cache = StaticCache(config=model.config, max_cache_len=1024)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to(model.device.type)
# This is the common prompt cached, we need to run forward without grad to be able to copy
with torch.no_grad():
     prompt_cache = model(**inputs_initial_prompt, past_key_values = prompt_cache).past_key_values

prompts = ["Help me to write a blogpost about travelling.", "What is the capital of France?"]
responses = []
for prompt in prompts:
    new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to(model.device.type)
    past_key_values = copy.deepcopy(prompt_cache)
    outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20)
    response = tokenizer.batch_decode(outputs)[0]
    responses.append(response)

print(responses)