Какие есть методы ускорения тест-тайм компьютинга? (KV-cache, speculative decoding)

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

Тест-тайм компьютинг (inference) — это этап, когда обученная модель генерирует ответ. Основные задержки возникают из-за автогрессивной природы трансформеров: каждый новый токен требует пересчёта внимания ко всем предыдущим. KV-cache устраняет повторные вычисления, сохраняя ключи и значения предыдущих токенов. Speculative decoding ускоряет генерацию, используя быструю draft-модель для предсказания нескольких токенов, которые затем верифицируются большой моделью параллельно. Дополнительные методы — квантованная верификация, Flash Attention, PagedAttention и batch inference — дают совокупный выигрыш в 2–5× без потери качества.

1. Почему тест-тайм компьютинг — узкое место

В авторегрессивных моделях (например, GPT, LLaMA) генерация каждого токена требует вычисления self-attention по всей предыдущей последовательности. Сложность внимания — O(n²·d), где n — длина контекста, d — размерность. При длинных контекстах (10k+ токенов) это становится доминирующим фактором задержки. Кроме того, каждый шаг генерации — это отдельный forward pass, что делает последовательную генерацию медленной.

Для Agentic RAG проблема усугубляется: агент может делать множество вызовов LLM (планирование, вызов инструментов, генерация ответа), и каждый вызов требует времени. Ускорение inference напрямую влияет на пользовательский опыт и стоимость.

2. KV-cache и его переиспользование

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

Как работает

1. На первом шаге (prefill) вычисляются K и V для всего входного промпта.
2. На каждом последующем шаге (decoding) вычисляются K и V только для нового токена, а attention к предыдущим токенам использует кэш.
3. Кэш хранится в памяти GPU и обновляется на каждом шаге.

Переиспользование KV-cache (KV-cache reuse) — это продвинутая техника, когда кэш от одного вызова модели используется для другого, если контекст частично совпадает. Например, в чат-интерфейсе системное сообщение и история диалога одинаковы для нескольких запросов. Можно сохранить кэш для общей части и достраивать только новый запрос.

Эффективность

• Экономия вычислений до 80% для длинных диалогов.
• Уменьшение latency в 2–4×.
• Основной недостаток — рост потребления памяти (кэш линейно растёт с длиной контекста). Для контекста 32k токенов и модели 7B параметров KV-cache может занимать ~1.5 GB на один запрос.

Реализация в популярных фреймворках

• vLLM — использует PagedAttention для эффективного управления KV-cache (разбивает кэш на страницы, избегая фрагментации).
• Hugging Face Transformers — поддерживает use_cache=True (по умолчанию включено).
• TensorRT-LLM — оптимизирует кэш для инференса на GPU.

3. Speculative decoding

Speculative decoding (спекулятивная декодировка) — метод, при котором быстрая draft-модель (обычно меньшего размера) генерирует несколько токенов (например, 5–10), а затем большая target-модель проверяет их все параллельно в одном forward pass. Если draft-модель угадала правильно, токены принимаются; если нет — target-модель корректирует.

Алгоритм

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

Ускорение 2–3× в среднем, до 5× на лёгких задачах (например, генерация кода). Зависит от качества draft-модели и степени совпадения распределений.

Пример реализации (псевдокод):

def speculative_decode(draft_model, target_model, prompt, K=5):
    # Draft: generate K tokens
    draft_tokens = draft_model.generate(prompt, max_new_tokens=K)
    # Target: compute logits for all positions
    target_logits = target_model.forward(prompt + draft_tokens)
    # Accept tokens where target probability > draft probability
    accepted = []
    for i in range(K):
        if target_logits[i] >= draft_logits[i]:
            accepted.append(draft_tokens[i])
        else:
            # resample from target distribution
            accepted.append(sample(target_logits[i]))
            break
    return accepted

Варианты

• Stochastic speculative decoding — использует rejection sampling для точного сохранения распределения target-модели.
• Greedy speculative decoding — принимает токены, если они совпадают с argmax target.
• Medusa — добавляет несколько головок к target-модели, которые предсказывают следующие токены без отдельной draft-модели.

Применимость в Agentic RAG особенно полезна, когда агент генерирует длинные ответы или вызывает инструменты с шаблонным синтаксисом (JSON, SQL). Draft-модель может быть специализирована на этих шаблонах.

4. Quantized verification (Quasar)

Quantized verification — метод, предложенный в работе Quasar (2024). Идея: использовать 4-битное квантование для фазы верификации в speculative decoding. Draft-модель генерирует токены, а target-модель в 4-битной точности проверяет их. Это даёт дополнительное ускорение за счёт меньшего объёма вычислений и памяти.

Как работает

1. Target-модель квантируется до 4-bit (например, с помощью GPTQ или AWQ).
2. Верификация выполняется на квантированной модели, что быстрее, чем FP16.
3. Если квантированная верификация даёт те же результаты, что и полная точность (что обычно верно для большинства токенов), ускорение достигается без потери качества.

Результаты

• Ускорение 1.28× относительно стандартного speculative decoding.
• Сохранение качества (perplexity меняется менее чем на 0.1).
• Дополнительная экономия памяти (4-bit вместо 16-bit).

Ограничения требует калибровочного датасета для квантования; возможна потеря качества на редких токенах.

5. Другие методы ускорения

5.1 Flash Attention

Flash Attention — алгоритм, который вычисляет attention без материализации полной матрицы внимания (O(n²) памяти). Использует tiling и recomputation, чтобы уложиться в быструю память (SRAM). Ускоряет как prefill, так и decoding, особенно для длинных контекстов. Ускорение 2–4× по сравнению с naive attention.

5.2 PagedAttention

PagedAttention — техника управления KV-cache, используемая в vLLM. Разбивает кэш на блоки (pages) фиксированного размера, что позволяет эффективно распределять память между запросами, избегать фрагментации и поддерживать continuous batching (объединение нескольких запросов в один батч на уровне шагов генерации).

5.3 Prefix caching

Prefix caching — сохранение KV-cache для общих префиксов (например, системного промпта). При повторном запросе с тем же префиксом кэш загружается из памяти, минуя повторные вычисления. Эффективно для чат-ботов и агентов с фиксированными инструкциями.

5.4 Batch inference

Batch inference — объединение нескольких независимых запросов в один батч. GPU эффективнее работает с большими батчами (лучше утилизация тензорных ядер). В Agentic RAG можно батчить вызовы от разных пользователей или параллельные вызовы внутри одного агента.

5.5 Continuous batching

Continuous batching — динамическое добавление новых запросов в батч по мере завершения старых. Позволяет достичь высокой пропускной способности (throughput) без увеличения latency для отдельных запросов.

6. Сравнение методов

Комбинирование лучшие результаты достигаются при совместном использовании нескольких методов. Например, vLLM применяет PagedAttention + continuous batching + Flash Attention. Добавление speculative decoding даёт дополнительный выигрыш.

7. Практические рекомендации для Agentic RAG

• Для агентов с длинными диалогами (multi-turn): обязательно используйте KV-cache reuse и prefix caching. Храните кэш для системного промпта и истории.
• Для агентов, генерирующих структурированные ответы (JSON, код): speculative decoding с draft-моделью, обученной на синтаксисе этих форматов.
• Для высоконагруженных систем: vLLM с PagedAttention и continuous batching. Добавьте Flash Attention, если GPU поддерживает.
• Для edge-устройств или ограниченной памяти: quantized verification и 4-bit квантование всей модели.
• Мониторинг: отслеживайте time-to-first-token (TTFT) и tokens-per-second (TPS). Speculative decoding улучшает TPS, но может увеличить TTFT из-за оверхеда draft-модели.

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

Задача Реализовать speculative decoding для небольшой модели (например, GPT-2 small как target и TinyStories-1M как draft) и измерить ускорение.

Инструменты

• Python, PyTorch, Transformers (Hugging Face)
• Датасет: случайные промпты из WikiText-2 (100 примеров)

Шаги:

1. Загрузите target-модель (distilgpt2) и draft-модель (gpt2 — на самом деле draft должна быть меньше, но для примера можно взять distilgpt2 как draft, а gpt2-medium как target; или использовать TinyStories).
2. Реализуйте функцию speculative_decode, которая принимает draft, target, промпт и K (число спекулятивных токенов).
3. Сравните время генерации 50 токенов с обычным декодированием (без speculative) и со speculative decoding (K=5).
4. Проверьте, что сгенерированные тексты совпадают (или близки) при использовании rejection sampling.
5. Постройте график ускорения в зависимости от K.

Ожидаемый результат Вы увидите ускорение в 1.5–2× для K=5. Если draft-модель слишком слабая, ускорение может быть меньше из-за частых rejection.

Расширение Добавьте quantized verification — квантируйте target-модель до 4-bit с помощью bitsandbytes и повторите эксперимент.

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

Навигация

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