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

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

KV cache reuse — это техника оптимизации инференса LLM, при которой вычисленные Key-Value кэши (KV cache) предыдущих шагов диалога сохраняются и переиспользуются для последующих запросов в рамках одной сессии. В multi-turn диалогах это позволяет избежать повторного вычисления внимания для уже обработанных токенов, сокращая latency и вычислительные затраты. Реализация включает prefix caching, shared KV cache между агентами (RelayCaching) и механизмы инвалидации при изменении контекста.

1. Термин: KV cache и его роль в трансформерах

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

Формула attention с кэшем

Attention(Q, K, V) = softmax(Q * [K_cache; K_new]^T / sqrt(d_k)) * [V_cache; V_new]

Где K_cache, V_cache — сохранённые матрицы для предыдущих токенов, K_new, V_new — для текущего.

Зачем нужен KV cache

• Уменьшает сложность с O(L²) до O(L) на шаг декодирования (L — длина последовательности).
• Критичен для реального времени (чат-боты, ассистенты).

Размер кэша для модели LLaMA-2 7B с контекстом 4096 токенов и половинной точностью (float16) KV cache занимает ~1.5 ГБ на один запрос. Для multi-turn диалогов с длинной историей это становится узким местом.

2. Multi-turn диалог: проблема повторных вычислений

В multi-turn диалоге каждый новый запрос пользователя добавляется к истории (предыдущие сообщения + ответы модели). Без reuse каждый вызов LLM обрабатывает всю историю с нуля:

• Turn 1 пользователь: «Что такое RAG?» → модель вычисляет KV cache для всего промпта.
• Turn 2 пользователь: «А как оценить retrieval?» → модель снова вычисляет KV cache для всей истории (включая turn 1), хотя большая часть токенов уже была обработана.

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

3. KV cache reuse: концепция и преимущества

KV cache reuse — это сохранение кэша между вызовами модели в рамках одной сессии (или между агентами). Принцип:

1. После первого вызова сохраняем KV cache для общего префикса (история диалога).
2. При следующем запросе передаём сохранённый кэш и вычисляем только для новых токенов (новый запрос пользователя).
3. Обновляем кэш, добавляя новые Key/Value.

Преимущества

• Снижение latency в 2–5 раз для длинных диалогов.
• Экономия FLOPs (операций с плавающей точкой) — не нужно пересчитывать attention для старых токенов.
• Уменьшение пиковой памяти (если кэш хранится в CPU или разделяемой памяти).

Термин «Prefix caching» — частный случай reuse, когда кэш сохраняется для фиксированного префикса (например, системного промпта и истории). Используется в vLLM, TensorRT-LLM, SGLang.

4. Реализация: prefix caching (базовый подход)

4.1. Структура хранения

Кэш хранится в виде словаря, где ключ — хэш префикса (или идентификатор сессии), значение — тензоры (num_layers, 2, batch_size, seq_len, d_k) для Key и Value.

class KVCache:
    def __init__(self, max_seq_len=4096, dtype=torch.float16):
        self.cache = {}  # session_id -> (k_cache, v_cache)
        self.max_seq_len = max_seq_len

    def get_or_compute(self, session_id, prefix_tokens, model, compute_fn):
        if session_id in self.cache:
            k_cache, v_cache = self.cache[session_id]
            # Вычисляем только для новых токенов
            new_k, new_v = compute_fn(prefix_tokens[len(k_cache[0][0]):])
            # Конкатенируем
            k_cache = torch.cat([k_cache, new_k], dim=-2)
            v_cache = torch.cat([v_cache, new_v], dim=-2)
        else:
            k_cache, v_cache = compute_fn(prefix_tokens)
        self.cache[session_id] = (k_cache, v_cache)
        return k_cache, v_cache

4.2. Интеграция с инференс-движком

В vLLM prefix caching реализован на уровне PagedAttention: блоки KV-кэша (pages) хэшируются, и если такой блок уже есть в глобальном пуле, он переиспользуется. Это позволяет делить кэш между разными запросами с одинаковым префиксом.

Шаги реализации

1. Разбить последовательность на блоки фиксированной длины (например, 16 токенов).
2. Вычислить хэш каждого блока (содержимое + позиция).
3. При генерации нового запроса искать блоки с таким же хэшом в shared memory.
4. Если найдены — использовать их, иначе вычислить и добавить.

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

Кэш нужно инвалидировать, если:

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

Стратегии

• LRU (Least Recently Used) — вытеснять блоки, к которым дольше всего не обращались.
• По сессии — очищать весь кэш при завершении диалога.
• По хэшу — если хэш блока изменился (из-за редактирования), удалить все последующие блоки.

5. RelayCaching: переиспользование между агентами

RelayCaching (Chen et al., 2026) — метод, предложенный для Agentic RAG, где несколько агентов (или вызовов LLM) решают схожие подзадачи. Идея: если два запроса имеют общий префикс (например, одинаковый контекст из документов), их KV кэш можно переиспользовать, даже если они принадлежат разным сессиям.

Пример:

• Агент A: «Извлеки ключевые факты из документа D1».
• Агент B: «Суммируй документ D1». Оба запроса начинаются с одного и того же документа D1. RelayCaching позволяет вычислить KV кэш для D1 один раз и использовать его для обоих агентов.

Реализация

• Глобальный реестр кэшей с хэшами префиксов.
• При старте нового запроса проверяется, есть ли кэш для данного префикса.
• Если есть — кэш «передаётся» (relay) новому агенту.
• Требуется синхронизация между агентами (через shared memory или Redis).

Преимущества для Agentic RAG

• Снижение latency при параллельной работе нескольких агентов.
• Экономия GPU-памяти (один кэш на несколько вызовов).
• Возможность кэшировать результаты retrieval (документы) на уровне KV.

6. Практические аспекты и ограничения

6.1. Управление памятью

KV cache для длинных диалогов может занимать десятки гигабайт. Решения:

• Offloading — перемещать старые блоки в CPU RAM (с потерей скорости).
• Compression — использовать KV cache quantization (INT8, FP8) или sparse attention.
• Sliding window — хранить только последние N токенов (например, 2048).

6.2. Ограничения reuse

• Разные промпты — если префиксы не совпадают, reuse невозможен. Для диалогов с изменяющейся историей (например, пользователь удалил сообщение) нужна инвалидация.
• Batch-запросы — при batch-обработке разных сессий кэш нужно хранить отдельно для каждой, что увеличивает память.
• Безопасность — если кэш shared между пользователями, возможна утечка информации (нужно изолировать по сессиям).

6.3. Сравнение подходов

7. Связь с Agentic RAG

В Agentic RAG агенты часто делают несколько последовательных вызовов LLM с общим контекстом (например, план действий, результаты поиска). KV cache reuse позволяет:

• Ускорить цепочки вызовов (agent loops).
• Делить кэш между агентами, работающими над одной задачей (RelayCaching).
• Уменьшить latency при повторном использовании одних и тех же документов.

Пример архитектуры

1. Retriever находит 5 документов.
2. Агент-планировщик вызывает LLM для составления плана (кэш для документов сохраняется).
3. Агент-исполнитель вызывает LLM для каждого шага плана, используя тот же кэш документов.
4. RelayCaching передаёт кэш от планировщика к исполнителю.

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

Задача Реализовать простой multi-turn чат-бот с KV cache reuse на базе небольшой LLM (например, GPT-2 или TinyLlama).

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

Шаги:

1. Загрузите модель и токенизатор.
2. Напишите класс KVCacheManager, который хранит кэш для каждой сессии (словарь session_id -> (past_key_values)).
3. Реализуйте функцию generate_with_cache(session_id, new_input_ids, max_new_tokens):
   • Если кэш есть, передайте его в model.generate(..., past_key_values=cache).
   • Если нет, сначала вычислите кэш для всего промпта.
   • После генерации обновите кэш (добавьте новые Key/Value).
4. Создайте простой цикл диалога: пользователь вводит текст, модель генерирует ответ, кэш сохраняется.
5. Замерьте latency для 2-го, 3-го и т.д. turn’ов — сравните с режимом без reuse.

Ожидаемый результат Второй и последующие запросы обрабатываются в 2–3 раза быстрее первого (для длинной истории). Вы увидите, что past_key_values содержит всё больше токенов, но время генерации растёт медленно.

Дополнительно Реализуйте инвалидацию кэша при изменении системного промпта (например, добавьте команду /reset).

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

Навигация

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