Что такое prefix caching и когда он эффективен?

Краткий тезис

Prefix caching — это техника оптимизации инференса LLM, при которой кэшируется KV-кэш (ключи и значения внимания) для общего начального блока токенов (префикса), например, одного system prompt или набора few-shot примеров. При повторении этого префикса в новом запросе система не вычисляет KV-кэш заново, а использует сохранённый, что резко сокращает время до первого токена (TTFT) и затраты на вычисления. Техника особенно эффективна в сценариях, где множество запросов разделяют один и тот же префикс: массовые чат-боты с единым system prompt, мульти-пользовательские приложения, batch-обработка с одинаковыми инструкциями. Провайдеры (Anthropic, vLLM) сообщают об ускорении TTFT до 10 раз и снижении затрат до 90%.

1. Термины: KV-кэш и prefix caching

KV-кэш (Key-Value cache) — это структура данных, которая хранит матрицы ключей $K$ и значений $V$ для каждого слоя трансформера на каждом шаге генерации. При авторегрессивном декодировании модель вычисляет внимание между текущим токеном и всеми предыдущими. Без кэша пришлось бы пересчитывать всю последовательность заново. KV-кэш позволяет хранить промежуточные результаты и добавлять к ним новый токен.

Prefix caching — это расширение идеи кэширования: мы кэшируем KV-кэш для префикса (одной или нескольких последовательностей токенов), который многократно используется. Если новый запрос начинается с того же префикса, система отображает сохранённый KV-кэш и достраивает его только для новой части (суффикса). Техника называется также prefix sharing или prompt caching.

2. Как работает prefix caching

2.1 Базовый механизм

Трансформер состоит из $L$ слоёв. Для каждого слоя вычисляются матрицы $K$ и $V$ размером $(n_{[text](/wiki/text){seq}}, d_{[text](/wiki/text){head}})$, где $n_{[text](/wiki/text){seq}}$ — длина последовательности. При prefix caching:

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

2.2 Поиск по общему префиксу

Системы (например, vLLM с механизмом PagedAttention, SGLang с RadixAttention) хранят несколько сегментов KV-кэша в виде дерева (radix tree). Каждый узел — это последовательность токенов и соответствующий KV-кэш. При получении нового запроса система обходит дерево, находит самый длинный общий префикс (longest common prefix) и использует его KV-кэш. Если префикс обрывается, кэш достраивается только для различающейся части.

2.3 Пример с кодом (упрощённая симуляция)

class PrefixCache:
    def __init__(self):
        self.cache = {}  # префикс -> KV-кэш (набор тензоров для всех слоёв)

    def compute_prefix_kv(self, prefix: list[int], model):
        # только для первого запроса с таким префиксом
        with torch.no_grad():
            output = model(prefix, use_cache=True)
        kv = output.past_key_values  # список кортежей (K, V) по слоям
        self.cache[tuple(prefix)] = kv
        return kv

    def lookup(self, prompt: list[int], model):
        # поиск самого длинного общего префикса
        for i in range(len(prompt), 0, -1):
            prefix = prompt[:i]
            if tuple(prefix) in self.cache:
                suffix = prompt[i:]
                kv = self.cache[tuple(prefix)]
                # достраиваем суффикс
                with torch.no_grad():
                    output = model(suffix, past_key_values=kv, use_cache=True)
                return output
        # нет общего префикса — считаем весь промпт
        return self.compute_prefix_kv(prompt, model)

На практике, кэш хранится в виде деревьев, а не простого словаря, чтобы эффективно работать с частичными совпадениями.

3. Когда prefix caching эффективен?

3.1 Основные сценарии

3.2 Какие префиксы выгодны

• Длинные префиксы: чем больше доля общего префикса в общей длине промпта, тем больше выгода.
• Высокая частота повторений: если один префикс встречается тысячи раз в минуту (APIGW → множество юзеров), окупаются затраты на его хранение в памяти.
• Стабильные префиксы: system prompt может меняться редко; если он меняется каждую минуту — кэш часто сбрасывается.

4. Метрики и выгоды

4.1 Метрики производительности

• TTFT (Time To First Token) — время от получения промпта до первого сгенерированного токена. Prefix caching уменьшает TTFT пропорционально длине префикса (обычно в 2–10 раз, по данным провайдеров).
• TPOT (Time Per Output Token) — время на генерацию каждого последующего токена. Не меняется, так как суффикс вычисляется как обычно.
• Общая задержка (end-to-end latency) = TTFT + (кол-во токенов) * TPOT. Выигрыш в TTFT даёт вклад в общее время, особенно для коротких суффиксов.

4.2 Экономическая выгода (стоимость)

• Anthropic сообщает, что prefix caching может снизить стоимость инференса до 90%, потому что префикс не требует FLOPs и времени GPU.
• vLLM с PagedAttention автоматически кэширует KV-кэш для повторяющихся последовательностей (включая префиксы) в рамках одного батча.

4.3 Нагрузка на память

Хранение KV-кэша для префиксов увеличивает потребление памяти GPU (или CPU + GPU transfer). Например, один префикс длиной 2000 токенов с 32 слоями и 32 головами требует ~2000 × 32 × (2×d_head) × precision байт. Для модели 7B с d=4096, d_head=128 это около 2000 × 32 × 2 × 128 × 2 байт = ~32 МБ (FP16). Тысяча таких префиксов — уже 32 ГБ. Поэтому системы используют механизмы вытеснения (LRU, TTL) и сегментирование.

5. Реализации в современных фреймворках

5.1 PagedAttention (vLLM)

• vLLM использует страничную организацию KV-кэша (аналогично виртуальной памяти). При prefix caching внутри батча общие префиксы автоматически делятся между запросами через одни и те же страницы (page). Это эффективно при batch inference, но не обязательно сохраняется между разными батчами.

5.2 RadixAttention (SGLang)

• SGLang реализует radix tree для хранения KV-кэша всех ранее встреченных последовательностей. При новом промпте система находит наибольший общий префикс (prefix match) с любым из сохранённых и использует его кэш. Поддерживает кэширование между запросами от разных пользователей.

5.3 Prompt Caching (Anthropic API)

• API Anthropic (Claude) предоставляет опцию cached_prefix: пользователь явно помечает фрагмент промпта, который гарантированно будет повторяться (например, system prompt). Платформа кэширует KV для этого фрагмента и выставляет пониженные тарифы на эти токены.

5.4 Пример настройки в SGLang

import sglang as sgl

@sgl.function
def chat(system, user_message):
    system += user_message   # префикс system+user может кэшироваться
    ...

При каждом вызове SGLang автоматически сравнивает префикс с деревом и использует кэш, если он уже есть.

6. Ограничения и подводные камни

7. Сравнение с другими техниками оптимизации

Prefix caching лучше всего сочетается с FlashAttention (быстрое вычисление внимания для суффикса) и используется в связке с continuous batching в таких системах, как vLLM.

8. Пет-проект для закрепления

Задача

Реализовать симуляцию prefix caching на основе простой модели (например, малый GPT с Hugging Face) и измерить ускорение TTFT при повторении префиксов разной длины.

Инструменты

• Python, PyTorch, Transformers
• Для ускорения — библиотека transformers с past_key_values
• Jupyter notebook для анализа

Шаги

1. Выбор модели — distilgpt2 или TinyStories-1M (маленькая, чтобы уместить на CPU/GPU).
2. Сбор тестовых запросов — создать набор из 100 промптов: первые 200 токенов одинаковы (system prompt), последние 20–50 различаются (суффиксы).
3. Без кэша — для каждого запроса передать полный промпт в модель, замерить TTFT (время от вызова до первого нового токена) и TPOT.
4. С кэшем — предварительно вычислить KV-кэш для префикса (200 токенов) и сохранить его. При каждом запросе передавать past_key_values= с сохранённым кэшем и суффиксом, замерить TTFT и TPOT.
5. Результаты — построить графики: TTFT vs длина суффикса, ускорение (во сколько раз), процент экономии времени. Дополнительно замерить потребление памяти (сколько байт занимает кэш для префикса).
6. Варьировать длину префикса — от 50 до 500 токенов, чтобы увидеть, при каком отношении префикс/суффикс кэш окупается.

Ожидаемый результат

• Вы получите чёткую зависимость: ускорение TTFT в 2–5 раз для типичного соотношения (префикс 80% длины промпта).
• Поймёте, какой объём памяти тратится на кэш одного префикса (для distilgpt2 — порядка 1–2 МБ на 200 токенов).
• Увидите, что TPOT не меняется — подтверждение теории.

9. Связь с другими вопросами

Техника prefix caching часто обсуждается вместе с другими методами ускорения (speculative decoding, prompt compression) и способами батчизации (continuous batching).

10. Навигация

• Предыдущий: 441
• Следующий: 443
• Индекс: 00. Индекс разборов

Навигация

• Предыдущий: 441
• Следующий: 443
• Индекс: 00. Индекс разборов