Почему KV cache растет линейно с длиной контекста и как это оптимизировать?

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

KV cache хранит ключи и значения внимания для каждого уже обработанного токена, что позволяет не пересчитывать их при генерации каждого нового токена. Размер кэша прямо пропорционален длине контекста, числу слоёв, голов внимания и разрядности данных, что приводит к гигантским требованиям к памяти (сотни гигабайт для контекстов в 128 k токенов). Основные методы оптимизации включают Attention|Grouped Query Attention (GQA) и Attention|Multi-Query Attention (MQA), квантование KV cache до INT4/INT8, а также техники управления памятью, такие как PagedAttention и sliding window attention.

1. Термин: KV cache (кэш ключей и значений)

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

2. Почему KV cache растёт линейно с длиной контекста

Размер KV cache можно выразить формулой (для одного запроса без батча):

Size = 2 * num_layers * num_heads * head_dim * seq_len * dtype_size

Расшифровка компонентов

• 2 – отдельно храним K и V.
• num_layers – количество слоёв трансформера (например, 32, 80).
• num_heads – количество голов внимания в каждом слое.
• head_dim – размерность одной головы (обычно 64, 128).
• seq_len – текущая длина сгенерированной последовательности (с каждым новым токеном +1).
• dtype_size – размер одного элемента в байтах (FP16 = 2, INT8 = 1, FP32 = 4).

Отсюда видна прямая пропорциональность seq_len. Если контекст удваивается, память тоже удваивается – это линейный рост.

Пример: Llama‑3‑70B

• num_layers = 80
• num_heads = 64
• head_dim = 128
• dtype = FP16 (2 байта)

Для одного токена: 2 × 80 × 64 × 128 × 2 = 2 621 440 байт ≈ 2,6 МБ. Для 128 k токенов: 2,6 МБ × 128 000 ≈ 333 ГБ (без батча и без кэша для каждого слоя отдельно, что превышает память даже топовых GPU).

Если добавить батч B, размер умножается на B. Например, batch size = 4 даёт 1,3 ТБ.

3. Проблема для длинных контекстов и Agentic RAG

Agentic RAG часто требует обработки очень длинных контекстов: агент может читать множество документов, вызывать инструменты, анализировать историю диалога. Например, при multi‑hop генерации контекст легко достигает 50–100 k токенов. Без оптимизаций KV cache переполняет VRAM, что приводит к падению скорости (swap на CPU) или к ограничению длины контекста.

4. Оптимизация 1: Grouped Query Attention (GQA)

GQA вводит концепцию групп голов: несколько query-голов делят один набор key/value-голов. Если в Multi‑Head Attention (MHA) все головы независимы (num_key_value_heads = num_query_heads), то в GQA num_key_value_heads меньше, и каждая key/value-голова обслуживает группу query-голов.

Формула размера KV cache с GQA:

Size = 2 * num_layers * num_kv_heads * head_dim * seq_len * dtype_size

где num_kv_heads может быть, например, 8 вместо 64. Тогда размер кэша уменьшается в num_query_heads / num_kv_heads раз (в примере в 8 раз).

Пример: Llama‑3‑70B (GQA с 8 KV‑головами)

• num_kv_heads = 8
• head_dim = 128
• Размер на токен: 2 × 80 × 8 × 128 × 2 = 327 680 байт ≈ 0,33 МБ.
• На 128k: 0,33 × 128 000 ≈ 42 ГБ – уже помещается в современные GPU (A100 80 ГБ).

Торговля небольшое падение качества при сильной экономии памяти. GQA стала стандартом в Llama‑2/3, Mistral, Qwen‑2.5.

5. Оптимизация 2: Multi-Query Attention (MQA)

MQA – крайний случай GQA, когда num_kv_heads = 1. Все query-головы используют единую пару K и V. Экономия максимальная, но качество может заметно снизиться на сложных задачах, где требуется различать информацию от разных голов.

Сравнение MHA, GQA, MQA:

6. Оптимизация 3: Квантование KV cache

Квантование снижает разрядность хранимых значений: FP16 (16 бит) → INT8 (8 бит) → INT4 (4 бита). Современные методы, такие как NF4 (нормализованный float4) из QLoRA, сохраняют точность, используя блочное квантование с общими статистиками (например, absmax на блок из 32–128 элементов).

Экономия

• FP16 → INT8: уменьшение в 2 раза.
• FP16 → INT4: уменьшение в 4 раза.

Для Llama‑3‑70B c GQA (0,33 МБ на токен) квантование до INT4 даёт ~0,08 МБ на токен, а на 128k – около 10 ГБ.

Вызов при квантовании во время вычисления внимания нужно деквантовать K и V обратно в FP16 (или вычислять в INT8 с помощью специальных операций). Современные библиотеки (vLLM, TensorRT‑LLM) поддерживают это прозрачно с минимальным замедлением.

7. Оптимизация 4: PagedAttention и vLLM

PagedAttention – техника, реализованная в системе vLLM, которая разбивает KV cache на блоки фиксированного размера (например, по 16 токенов) и управляет памятью аналогично виртуальной памяти в ОС: блоки могут располагаться не непрерывно, что устраняет фрагментацию и позволяет более гибко выделять память для разных запросов.

Преимущество можно обслуживать несколько запросов с различной длиной контекста, не резервируя максимально возможный кэш для каждого. Эффективность использования памяти возрастает на 20–40% по сравнению с традиционной плоской аллокацией.

Дополнительно vLLM поддерживает копирование KV cache между запросами при shared prefixes (например, в чат‑ботах с системным промптом), что ещё сильнее экономит память.

8. Оптимизация 5: Sliding Window Attention

Sliding window attention (оконное внимание) ограничивает контекст, который видит модель, скользящим окном фиксированной ширины W. KV cache хранит только последние W токенов, поэтому его размер не зависит от общей длины контекста (становится константным O(W)).

Пример: Mistral‑7B использует W = 4096. При генерации 100k токенов кэш всё равно ~0,33 МБ × 4 096 ≈ 1,3 ГБ.

Недостаток модель «забывает» токены за пределами окна. Для задач, где нужна информация из начала документа (например, инструкции в длинном диалоге), это неприемлемо. Компромисс – комбинировать с attention sink или делать частичный глобальный слой.

9. Оптимизация 6: Offloading KV cache на CPU/диск

При экстремально длинных контекстах (сотни тысяч токенов) даже с GQA и квантованием KV cache может не поместиться в VRAM. Тогда часть кэша (самые старые токены) выгружается на CPU DRAM или даже на SSD. При обращении к ним во время генерации блоки загружаются обратно.

Торговля резкое увеличение latency (в 10–100 раз) из‑за медленных каналов (PCIe, NVMe). Применяется редко, в основном в исследовательских сценариях или когда скорость не критична.

Пример: библиотека FlexGen позволяет распределять KV cache между GPU, CPU и SSD.

10. Сравнительная таблица методов оптимизации

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

• Используйте model zoo с поддержкой GQA (Mistral, Llama‑3‑8B/70B, Qwen‑2.5).
• Квантуйте KV cache до INT8 или INT4. В production – через vLLM или TensorRT‑LLM.
• Настройте PagedAttention – она уже встроена в vLLM, просто включите --max-model-len и --gpu-memory-utilization.
• Если агент работает с контекстом >100k, рассмотрите sliding window + attention sink (например, как в InfLLM или StreamingLLM).
• Для самых длинных сценариев (чтение миллиона токенов) – offloading на CPU/SSD, но только если latency не критична (например, nightly batch анализ).

12. Связь с архитектурой Agentic RAG

В Agentic RAG агент часто делает несколько шагов, на каждом из которых контекст увеличивается (добавляются результаты поиска, вызовов инструментов). KV cache неизбежно растёт. Если агент параллельно обрабатывает несколько запросов (multi‑turn, multi‑agent), память становится критическим ресурсом. Поэтому грамотная оптимизация KV cache – ключевое требование для масштабируемых агентов.

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

Задача написать скрипт, который симулирует рост KV cache для нескольких архитектур (MHA, GQA, MQA с квантованием) и визуализирует экономию памяти при разных длинах контекста.

Инструменты Python, PyTorch (или просто numpy для моделирования), matplotlib.

Шаги:

1. Определите функцию kv_cache_size(layers, kv_heads, head_dim, seq_len, bits).
2. Постройте графики для:
   • Llama‑3‑70B (MHA, GQA 8, INT4)
   • Mistral‑7B (GQA 8, FP16)
   • Qwen‑2.5‑72B (GQA 8, INT8)
3. На одном графике отложите размер кэша от seq_len (до 1 млн токенов).
4. Добавьте горизонтальную линию, показывающую VRAM типичного GPU (например, 80 ГБ).
5. Вычислите, при какой длине контекста каждая конфигурация «ломается».

Ожидаемый результат увидите, что без оптимизаций даже 50k токенов не помещаются в 80 ГБ, а GQA+INT4 позволяют обрабатывать >500k.

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

Навигация

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