Как speculative decoding взаимодействует с KV cache?

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

Speculative decoding (SD) — техника ускорения инференса LLM, при которой маленькая draft-модель генерирует несколько токенов, а большая target-модель проверяет их параллельно. Это создаёт tension в управлении KV cache: обе модели требуют собственных кэшей (Memory Overhead Ratio 1.4–2.0), а при отклонении (rejection) части токенов кэш target-модели нужно откатывать. Современные решения, такие как RelayCaching, предлагают переиспользовать KV cache между вызовами агентов в multi-turn сценариях, что особенно актуально для архитектур RAG|Agentic RAG.

1. Термины: speculative decoding и KV cache

Speculative decoding — метод ускорения генерации, при котором маленькая draft-модель (обычно в 2–10 раз меньше target) быстро генерирует последовательность из K токенов. Target-модель затем проверяет эти токены за один forward pass, принимая или отвергая каждый с помощью rejection sampling. Это позволяет за один шаг target-модели продвинуться на несколько токенов, снижая latency.

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

Draft-модель — лёгкая модель (например, 1–2 млрд параметров), используемая для быстрой генерации черновика. Target-модель — основная модель (например, 70B), которая валидирует черновик.

2. Проблема: две модели — два KV cache

При использовании SD одновременно работают две модели, каждая со своим KV cache. Это приводит к увеличению потребления памяти:

Memory Overhead Ratio = (KV_cache_target + KV_cache_draft) / KV_cache_target. Для batch > 1 и при использовании нескольких draft-моделей (например, в Medusa) этот коэффициент растёт.

Почему это проблема: в сценариях Agentic RAG, где несколько агентов могут параллельно вызывать LLM (каждый со своим SD), общий memory footprint может превысить доступную видеопамять, требуя offloading или уменьшения batch size.

3. Откат KV cache при rejection

Когда target-модель отвергает один или несколько токенов из черновика, необходимо откатить KV cache до состояния, соответствующего последнему принятому токену. Иначе при генерации следующего токена attention будет учитывать отвергнутые токены, что приведёт к некорректным результатам.

Процесс отката (rollback):

1. Draft-модель генерирует K токенов, сохраняя свой KV cache для этих токенов.
2. Target-модель выполняет один forward pass на всей последовательности (prefill + verification), используя свой KV cache, который уже содержит предыдущие принятые токены.
3. Для каждого из K токенов target-модель вычисляет вероятность. Если вероятность токена из draft ниже порога (обычно сравнивается с вероятностью target), токен отвергается.
4. Если отвергнут токен на позиции j, то все токены после j (включая j) отбрасываются. KV cache target-модели усекается до длины, соответствующей последнему принятому токену.

Реализация отката

• Copy-on-write: при каждом новом шаге создаётся копия KV cache, но физически данные не копируются до первого изменения. Это экономит память, но усложняет логику.
• Чекпоинты: периодически сохраняется состояние KV cache (например, каждые 10 токенов). При откате восстанавливается ближайший чекпоинт и достраивается до нужной длины.
• In-place rollback: в некоторых фреймворках (vLLM, TensorRT-LLM) KV cache хранится в виде списка тензоров, и откат реализуется простым изменением указателя длины (length truncation).

Пример псевдокода отката

def speculative_decode(draft_model, target_model, prompt, K=5):
    # prompt уже закэширован в KV cache target
    accepted_tokens = []
    for _ in range(K):
        draft_token = draft_model.generate_next(accepted_tokens)
        accepted_tokens.append(draft_token)
    # target verification
    target_logits = target_model.forward(accepted_tokens, kv_cache=target_kv_cache)
    rejected_pos = None
    for i, (draft_prob, target_prob) in enumerate(zip(draft_probs, target_probs)):
        if target_prob < draft_prob:  # упрощённое условие rejection
            rejected_pos = i
            break
    if rejected_pos is not None:
        # откатываем KV cache до длины len(prompt) + rejected_pos
        target_kv_cache.truncate(len(prompt) + rejected_pos)
        accepted_tokens = accepted_tokens[:rejected_pos]
    # продолжаем генерацию с последнего принятого токена

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

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

Как это работает

1. Первый агент (или первый вызов) генерирует ответ, сохраняя KV cache для всего контекста (промпт + сгенерированные токены).
2. Второй агент получает тот же контекст (или его часть) и может загрузить сохранённый KV cache, достроив только новые токены.
3. При использовании SD каждый агент может иметь свой draft-модель, но KV cache target-модели может быть общим, если draft-модель не меняет контекст.

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

• Снижение latency в multi-turn диалогах (до 40% ускорения).
• Уменьшение memory overhead, так как не нужно хранить отдельные кэши для каждого агента.
• Возможность параллельной работы агентов с общим кэшем (требует механизмов блокировок).

Ограничения

• RelayCaching эффективен только если контекст не меняется между вызовами (или меняется незначительно).
• При rejection токенов в SD кэш target-модели может быть частично инвалидирован, что требует синхронизации между агентами.

5. Влияние на производительность

Компромиссы

• Увеличение memory overhead может ограничить batch size, снижая throughput.
• Rollback добавляет накладные расходы, особенно при частых rejection (например, при низком качестве draft-модели).
• RelayCaching требует дополнительной инфраструктуры для хранения и синхронизации кэшей.

6. Практические реализации

vLLM — популярный фреймворк для инференса LLM. Поддерживает SD через механизм "speculative decoding" с автоматическим управлением KV cache. Draft-модель может быть загружена как отдельный worker, а KV cache хранится в PagedAttention, что упрощает откат (достаточно изменить page table).

TensorRT-LLM — оптимизированный движок от NVIDIA. Предлагает "speculative decoding" с поддержкой Medusa и Eagle. KV cache управляется через блоки фиксированного размера, откат реализован через изменение индексов.

Hugging Face Transformers — базовая реализация SD (generate с параметром assistant_model). KV cache хранится в виде tuple тензоров, откат выполняется через срезы (slicing). Не оптимизирован для production.

Пример конфигурации в vLLM

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    speculative_model="lmsys/vicuna-7b-v1.5",  # draft модель
    num_speculative_tokens=5,
    use_v2_block_manager=True  # для эффективного управления KV cache
)

7. Связь с Agentic RAG

В архитектурах Agentic RAG несколько агентов (например, retriever, re-ranker, generator, fact-checker) могут вызывать LLM многократно. Каждый вызов может использовать SD для ускорения. Управление KV cache становится критичным:

• Multi-turn диалоги: агент генерации может вызываться несколько раз с растущим контекстом. RelayCaching позволяет избежать повторного prefill.
• Параллельные агенты: если два агента используют один и тот же контекст (например, общий набор документов), можно разделять KV cache, но нужно учитывать блокировки при rejection.
• Агенты с разными draft-моделями: могут возникать конфликты при откате кэша, если один агент отклонил токены, а другой ещё использует старый кэш.

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

1. Retrieval агент получает запрос, ищет документы, генерирует список релевантных чанков (используя SD).
2. Generator агент получает запрос + чанки, генерирует ответ (тоже с SD).
3. Если Generator отвергает часть черновика, его KV cache откатывается. При этом Retrieval агент мог уже начать новый поиск, используя общий кэш контекста — требуется синхронизация.

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

Задача Реализовать симуляцию speculative decoding с управлением KV cache и откатом для двух моделей (draft и target) в сценарии multi-turn диалога.

Инструменты Python, PyTorch, transformers (можно использовать маленькие модели, например, GPT-2 и DistilGPT-2).

Шаги:

1. Загрузите draft-модель (DistilGPT-2) и target-модель (GPT-2).
2. Реализуйте класс KVCache, который хранит тензоры ключей и значений, поддерживает методы append, truncate, copy.
3. Реализуйте функцию speculative_decode(prompt, draft_model, target_model, K=3):
   • Draft генерирует K токенов, сохраняя свой кэш.
   • Target проверяет токены, используя свой кэш.
   • При rejection откатывайте кэш target и возвращайте принятые токены.
4. Добавьте симуляцию multi-turn: в каждом раунде добавляйте новый пользовательский запрос к истории, используйте RelayCaching (сохраняйте кэш target между раундами).
5. Измерьте latency и memory usage для разных K и сравните с обычной генерацией.

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

• Вы увидите, что SD ускоряет генерацию, но при частых rejection (например, если draft-модель плохая) overhead отката может нивелировать выигрыш.
• RelayCaching даст дополнительное ускорение в multi-turn сценарии (prefill только для новых токенов).
• Вы сможете оценить Memory Overhead Ratio для ваших моделей.

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

Навигация

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