Что такое KV cache reuse в multi-turn диалогах и как его реализовать?

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

KV cache reuse — техника ускорения инференса LLM в многошаговых диалогах (multi-turn), при которой кэш ключей]] (K) и значений (V) из предыдущих шагов сохраняется и переиспользуется для общих частей промпта (system prompt, история диалога). Это позволяет избежать повторного вычисления attention для одних и тех же токенов, снижая latency и вычислительные затраты. Реализация требует управления памятью, сегментации кэша и корректной инвалидации при изменении контекста. RelayCaching расширяет идею на multi-agent сценарии, где несколько агентов разделяют общий префикс.

1. Термины и базовые понятия

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

Multi-turn диалог — последовательность запросов (turns) пользователя и ответов модели, где каждый новый запрос включает историю предыдущих обменов. Пример: пользователь задаёт вопрос, модель отвечает, затем пользователь уточняет — новый запрос содержит system prompt + всю предыдущую переписку.

Prefix caching — частный случай KV cache reuse, при котором кэш сохраняется для фиксированного префикса (например, system prompt), общего для всех запросов в диалоге.

RelayCaching — техника, предложенная в контексте agentic RAG, где несколько агентов (LLM-вызовов) используют один и тот же префикс (например, инструкции для агента) и могут разделять KV cache, передавая его между вызовами.

2. Зачем нужен KV cache reuse

Основные мотивации:

• Снижение latency: в длинных диалогах история может занимать тысячи токенов. Без reuse каждый новый запрос требует полного пересчёта attention для всей истории, что линейно растёт с длиной контекста. С reuse время обработки нового запроса практически не зависит от длины истории (только для нового токена).
• Экономия compute: уменьшение числа операций умножения матриц в attention слоях. Для диалога из N шагов с длиной истории L токенов экономия составляет ~O(N*L^2) операций.
• Увеличение пропускной способности: сервер может обслуживать больше параллельных диалогов при том же количестве GPU.

3. Как работает KV cache в LLM (кратко)

В трансформере decoder-only каждый слой содержит multi-head attention. На шаге t:

1. Входной токен x_t преобразуется в Query (Q_t), Key (K_t), Value (V_t) через линейные проекции.
2. Attention вычисляется как softmax(Q_t * [K_1..K_t]^T / sqrt(d)) * [V_1..V_t].
3. Если сохранить K_1..K_{t-1} и V_1..V_{t-1} в кэше, то на шаге t нужно вычислить только K_t, V_t, а затем attention с кэшированными матрицами.

Без кэша пришлось бы заново вычислить K_1..K_t и V_1..V_t для всех токенов.

Размер кэша: для модели с L слоями, H головами, размерностью d_head и длиной последовательности T: ~ 2 * L * H * T * d_head * precision_bytes. Например, для LLaMA-7B (32 слоя, 32 головы, d_head=128, T=4096, fp16) ~ 232324096128*2 = 2.1 GB.

4. Проблема multi-turn диалогов

Рассмотрим диалог:

• Turn 1: пользователь: "Какая погода?" → модель: "Солнечно."
• Turn 2: пользователь: "А завтра?" → модель: "Дождь."

Промпт для turn 2 (с историей):

System: Ты помощник.
User: Какая погода?
Assistant: Солнечно.
User: А завтра?

Без reuse модель вычисляет K,V для всех 20+ токенов заново. С reuse можно сохранить K,V для "System: Ты помощник. User: Какая погода? Assistant: Солнечно." и вычислить только для "User: А завтра?".

Проблема: история может меняться (пользователь редактирует предыдущее сообщение, system prompt динамический). Нужна стратегия инвалидации.

5. Подходы к KV cache reuse

5.1 Prefix caching (кэширование префикса)

Самый простой подход: кэш сохраняется для общего префикса (system prompt + фиксированная часть истории). При каждом новом turn кэш префикса остаётся неизменным, добавляется только новый сегмент.

Реализация: хранить кэш в виде списка сегментов, каждый сегмент соответствует части промпта. При новом запросе проверяется, совпадает ли начало промпта с уже закэшированным префиксом. Если да — используем кэш, иначе сбрасываем.

5.2 Sliding window cache

Для очень длинных диалогов (например, 100k токенов) кэш всей истории может не поместиться в память. Используется скользящее окно: кэшируются только последние N токенов. Старые токены отбрасываются. Это компромисс между точностью (модель может забыть ранний контекст) и памятью.

5.3 Segment caching (сегментное кэширование)

Промпт разбивается на логические сегменты: system prompt, каждый turn пользователя, каждый ответ модели. Каждый сегмент кэшируется отдельно. При новом turn можно переиспользовать все сегменты, кроме последнего (если он изменился). Это даёт максимальную гибкость, но требует более сложного менеджера кэша.

6. Реализация KV cache reuse (пример на Python)

Рассмотрим реализацию с использованием библиотеки Hugging Face transformers и ручным управлением кэшем.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "microsoft/phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

class KVCacheManager:
    def __init__(self):
        self.cache = None  # хранит tuple of (past_key_values, token_ids)
        self.prefix_ids = None

    def generate(self, prompt, use_cache=True):
        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
        if use_cache and self.cache is not None:
            # проверяем, что начало промпта совпадает с закэшированным префиксом
            prefix_len = self.cache[1].shape[1]  # длина закэшированных токенов
            if torch.equal(input_ids[0, :prefix_len], self.cache[1][0]):
                # используем кэш, передаём только новые токены
                new_ids = input_ids[:, prefix_len:]
                if new_ids.shape[1] == 0:
                    # если промпт полностью совпадает (например, повторный запрос)
                    new_ids = torch.tensor([[tokenizer.eos_token_id]], device=model.device)
                out = model.generate(
                    new_ids,
                    past_key_values=self.cache[0],
                    max_new_tokens=100,
                    do_sample=False
                )
                # обновляем кэш
                full_ids = torch.cat([self.cache[1], out], dim=1)
                self.cache = (model.past_key_values, full_ids)
                return tokenizer.decode(out[0], skip_special_tokens=True)
        # иначе полный пересчёт
        out = model.generate(input_ids, max_new_tokens=100, do_sample=False)
        self.cache = (model.past_key_values, input_ids)
        return tokenizer.decode(out[0], skip_special_tokens=True)

manager = KVCacheManager()
# Первый turn
print(manager.generate("System: Ты помощник.\nUser: Привет!\nAssistant:"))
# Второй turn (история та же + новый запрос)
print(manager.generate("System: Ты помощник.\nUser: Привет!\nAssistant: Здравствуйте!\nUser: Как дела?\nAssistant:"))

Важно: в реальных фреймворках (vLLM, TensorRT-LLM) управление кэшем встроено и оптимизировано. Приведённый код — иллюстрация принципа.

7. RelayCaching — reuse между агентами

В архитектуре Agentic RAG несколько агентов (LLM-вызовов) могут использовать общий system prompt (например, инструкции для агента, описание инструментов). RelayCaching позволяет сохранить KV cache для этого общего префикса и передавать его между вызовами разных агентов, экономя compute.

Пример сценария:

• Агент A: "Ты аналитик. Используй инструменты: поиск, калькулятор."
• Агент B: "Ты аналитик. Используй инструменты: поиск, калькулятор." (тот же префикс)

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

Реализация RelayCaching требует:

• Идентификации общего префикса (может быть задан явно или вычислен через longest common prefix).
• Хранения кэша в глобальном хранилище (in-memory, Redis).
• Механизма блокировок, чтобы избежать гонок при параллельных вызовах.

8. Инвалидация кэша

Кэш нужно сбрасывать, когда:

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

Стратегии инвалидации:

• Полный сброс: при любом изменении контекста.
• Частичная инвалидация: сбрасывается только изменённый сегмент и все последующие.
• Lazy invalidation: проверка совпадения префикса при каждом запросе (как в примере выше).

9. Компромиссы и ограничения

Ограничения:

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

10. Инструменты и фреймворки

• vLLM: встроенная поддержка prefix caching (--enable-prefix-caching). Автоматически кэширует общие префиксы между запросами.
• TensorRT-LLM: поддерживает KV cache reuse через in-flight batching и paged attention.
• Hugging Face TGI: поддержка prefix caching с флагом --prefix-caching.
• SGLang: фреймворк с продвинутым управлением кэшем, включая RadixAttention (дерево префиксов).
• OpenAI API: не предоставляет прямого доступа к кэшу, но автоматически переиспользует его для одинаковых префиксов в рамках сессии.

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

Задача: Реализовать чат-бота с KV cache reuse на основе любой open-source LLM (например, Phi-2 или TinyLlama). Бот должен поддерживать multi-turn диалог и показывать ускорение по сравнению с полным пересчётом.

Инструменты: Python, Hugging Face Transformers, timeit.

Шаги:

1. Загрузите модель и токенизатор.
2. Реализуйте класс ChatSession, который хранит past_key_values и историю токенов.
3. Метод chat(user_input):
   • Формирует полный промпт (system + история + новый запрос).
   • Если начало промпта совпадает с закэшированным префиксом — использует кэш.
   • Иначе — полный пересчёт и обновление кэша.
4. Замерьте время ответа для 5-10 turns с reuse и без.
5. Выведите статистику: среднее время, ускорение.

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

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

Навигация

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