Prompt Caching

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

Как работает кэширование подсказок#

LLM на основе трансформера отвечает в два этапа. На этапе предварительного заполнения каждый слой вычисляет вектор ключа и значения для каждого токена подсказки; при декодировании каждый новый токен учитывает все эти ключи и значения через механизм внимания. В пределах одного запроса KV-кэш хранит эти тензоры, чтобы не вычислять подсказку заново для каждого сгенерированного токена. Кэширование подсказок сохраняет их между запросами. Поскольку внимание в этих моделях причинно обусловлено, ключи и значения токена зависят только от предшествующих токенов; поэтому у двух подсказок с одинаковым префиксом ключи и значения для этого префикса совпадают, и сохранённую копию можно использовать вместо повторного вычисления.

В следующем примере показано повторное использование кэша с GPT-2 из библиотеки Transformers от Hugging Face. Сначала предварительно обрабатывается системная подсказка из 409 токенов, затем на основе кэшированного префикса обрабатываются только токены двух вопросов, а результат сравнивается с полным пересчётом:

import copy

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

# A long system prompt shared by every request, followed by different user questions
system = "You answer questions about a camera store. " + "Returns are accepted within 30 days with a receipt. " * 40
questions = [" Question: Can I return a lens?", " Question: Do you ship abroad?"]
prefix_ids = tokenizer(system, return_tensors="pt").input_ids

with torch.no_grad():
    # Prefill the shared prefix once and keep its key-value cache
    prefix_cache = model(prefix_ids).past_key_values
    for question in questions:
        ids = tokenizer(question, return_tensors="pt").input_ids
        # Cache hit: only the question's tokens run through the model
        cached = model(ids, past_key_values=copy.deepcopy(prefix_cache)).logits[0, -1]
        # No cache: the whole prompt is computed from scratch
        full = model(torch.cat([prefix_ids, ids], dim=1)).logits[0, -1]
        same = torch.allclose(cached, full, atol=1e-3)
        print(f"Computed {ids.shape[1]} of {prefix_ids.shape[1] + ids.shape[1]} tokens, same next-token scores: {same}")

Вывод содержит Computed 8 of 417 tokens, same next-token scores: True и Computed 7 of 416 tokens, same next-token scores: True: для каждого запроса модель обработала менее 2% его токенов и получила те же оценки следующих токенов, что и при полном вычислении.

Движки с открытым исходным кодом для обслуживания моделей применяют этот механизм автоматически. vLLM хранит KV-кэш в блоках фиксированного размера — так устроено размещение памяти в PagedAttention, — повторно использует блоки самого длинного совпадающего префикса для каждого нового запроса и вытесняет наименее используемые блоки, когда заканчивается память GPU. RadixAttention в SGLang хранит кэшированные префиксы в префиксном дереве, чтобы запросы с любой общей начальной частью могли повторно использовать кэш; в статье о SGLang сообщается о пропускной способности до 6.4 раза выше, чем у других современных движков инференса. Статья «Кэш подсказок» Ина Гима и его коллег, представленная на MLSys 2024, расширяет повторное использование за пределы точных префиксов, позволяя объявлять модули подсказок, например шаблоны и документы; сообщается о сокращении времени до первого токена от 8 раз на GPU до 60 раз на CPU.

Кэширование подсказок в API LLM#

Размещённые API предоставляют тот же механизм, но взимают за него плату. Кэширование подсказок OpenAI включено по умолчанию для поддерживаемых моделей, тогда как кэширование подсказок Claude от Anthropic включается с помощью поля cache_control — один раз для запроса или для четырёх блоков содержимого. Anthropic берёт за чтение из кэша не более одной десятой базовой цены входных данных, а запись в кэш по умолчанию со сроком хранения пять минут стоит 1.25 базовой цены; каждое попадание в кэш заново запускает пятиминутный срок. Оба поставщика сопоставляют только идентичный префикс: Anthropic сравнивает сначала инструменты, затем системную подсказку и сообщения, поэтому изменение определения инструмента делает недействительным весь следующий кэш; руководство OpenAI советует помещать стабильные инструкции и справочные материалы в начало, а пользовательские данные — в конец.

Наибольший выигрыш получают рабочие нагрузки с длинными общими префиксами:

• ИИ-агенты: На каждом шаге они повторно отправляют определения инструментов и всю историю — подсказку предыдущего шага, а также последнее действие и его результат.
• Чат-боты: Каждый ход диалога дополняет промпт предыдущего хода, поэтому история кэшируется вплоть до нового сообщения.
• Ответы на вопросы по документам: При множестве вопросов об одном длинном договоре, отчёте или репозитории обработанный документ используется повторно, а конвейеры поиска повторно используют его, когда снова извлекаются те же фрагменты.
• Промпты с несколькими примерами: Десятки фиксированных примеров вычисляются один раз и используются повторно для каждого входного запроса.

Модели компьютерного зрения с открытым словарём аналогичным образом повторно используют вычисления промптов. Ultralytics YOLOE один раз кодирует промпты классов, например «вилочный погрузчик», с помощью текстового энкодера MobileCLIP при вызове set_classes() и сохраняет полученные эмбеддинги в модели, поэтому при последующих предсказаниях текстовый энкодер вообще не запускается; save_prompt_embeddings() записывает их в файл NPZ, который следующий сеанс загружает с помощью load_prompt_embeddings(), вместо повторного запуска энкодера.

Кэширование промптов и связанные понятия#

• KV-кэш: Повторно использует ключи и значения между токенами одной генерации. Кэширование промптов повторно использует их в отдельных запросах.
• Семантическое кэширование: Сохраняет ответы целиком и возвращает один из них, когда новый запрос означает то же, что и предыдущий, полностью пропуская модель. При кэшировании промптов модель по-прежнему обрабатывает новую часть промпта и генерирует новый ответ.
• Сжатие промпта: Сокращает сам промпт, что удобно для текста, меняющегося при каждом запросе. Эти методы хорошо сочетаются: фиксированная часть промпта кэшируется, а переменная — сжимается.
• Промпт-инжиниринг: Определяет содержание промпта. Кэширование меняет только объём вычислений, которые сервер выполняет повторно, а попадание в кэш не меняет оценки следующих токенов, кроме погрешности округления с плавающей точкой.