Как устроен KV cache? Почему он bottleneck?

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

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

1. Что такое KV cache и зачем он нужен?

В авторегрессионном трансформере (например, GPT, LLaMA) генерация каждого следующего токена требует вычисления внимания ко всем предыдущим токенам. Без кэша на шаге t пришлось бы заново прогнать через все слои всю последовательность длины t, что даёт сложность O(t^2) на каждый шаг и O(t^3) на всю генерацию. KV cache решает эту проблему: на каждом шаге для нового токена вычисляются только его собственные Q (запрос), K (ключ) и V (значение). Затем его K и V добавляются в кэш, а для внимания используются все K и V из кэша (все предыдущие токены). Q-векторы не кэшируются, так как они нужны только для текущего токена.

Таким образом, KV cache превращает квадратичную сложность по длине контекста в линейную — каждый шаг требует O(1) относительно длины контекста (только константное количество новых вычислений и чтение всего кэша).

2. Как устроен KV cache?

Структура Для каждого слоя l и каждой головы внимания h хранятся два тензора:

• Keys форма [1, seq_len, head_dim
• Вики/Value|Values форма [1, seq_len, head_dim

На практике часто хранят объединённый кэш для всех голов в слое: batch_size, num_heads, seq_len, Вики/head_dim|head_dim. В инференсе с batch-запросами Вики/caching|кэш хранится отдельно для каждого элемента батча.

Процесс на шаге t (для одного слоя):

1. Вход: Вики/Hidden state|скрытое состояние нового токена h_t.
2. Линейные проекции: Q_t = h_t * W_Q, K_t = h_t * W_K, V_t = h_t * W_V.
3. Конкатенация с кэшем: K_cache = [K_cache; K_t], V_cache = [V_cache; V_t] (по оси seq_len).
4. Attention: score = softmax(Q_t @ K_cache^T / sqrt(d)) @ V_cache.
5. Результат идёт в FFN и далее.

Важно На шаге t Q-вектор взаимодействует со всеми t ключами, но K и V пересчитываются только для последнего токена. Без кэша пришлось бы пересчитать K и V для всех t токенов заново.

3. Размер KV cache (формула)

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

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

Где:

• num_layers — количество слоёв деко­дера (например, 80 для LLaMA-3-70B)
• num_heads — количество голов внимания (например, 64 для LLaMA-3-70B)
• head_dim — размерность одной головы (обычно 64 или 128)
• seq_len — текущая длина сгенерированной последовательности
• dtype_bytes — число байт на элемент (2 для FP16, 1 для INT8, 0.5 для INT4)

Множитель 2 — отдельно для K и V.

4. Пример расчёта (LLaMA-3-70B, 128k токенов)

Размер одного тензора K (или V) для одного слоя: 64 * 128 * 128 000 * 2 = 2 097 152 000 байт ≈ 2.1 ГБ

На все слои: 2.1 ГБ * 80 = 168 ГБ. Умножаем на 2 (K+V) → 336 ГБ.

Вывод Для одного запроса с контекстом 128k токенов KV cache занимает 336 ГБ — больше, чем веса самой модели (около 140 ГБ в FP16). Это категорически не помещается в память одного GPU (A100 80GB, H100 80GB). Даже с батч-размером >1 ситуация критическая.

Примечание В черновике упоминалось 3.3 ГБ — это ошибка округления. Правильный расчёт: 2 * 80 * 64 * 128 * 128 000 * 2 даёт ~335.5 ГБ. Меньшие значения могли получиться для меньшего seq_len или других параметров (например, LLaMA-3-8B: 32 слоя, 32 головы, 128 head_dim, 128k → ~33 ГБ — всё ещё много).

5. Почему KV cache — bottleneck?

5.1 Ограничение по памяти (capacity)

KV cache может легко превысить объём видеопамяти GPU, особенно при длинных контекстах и больших batch'ax. Это ограничивает максимальную длину контекста, batch size и количество параллельных запросов. Например, для обслуживания одного запроса с 128k токенами нужно ~340 ГБ — невозможно на одном GPU. Приходится использовать распределённую память (например, vLLM PagedAttention) или сжимать кэш.

5.2 Пропускная способность (bandwidth)

На каждом шаге генерации attention считывает весь KV cache из HBM (High Bandwidth Memory) в SRAM (shared memory) для вычисления softmax. Даже если размер кэша умещается в память GPU, его чтение становится узким местом: для seq_len = 128k и head_dim=128, num_heads=64, слоёв=80, один шаг требует чтения 336 ГБ. При пропускной способности HBM (например, 2 ТБ/с на A100) на чтение уходит ~170 мкс — это время почти равно времени вычисления. А при большом batch'е требуется чтение ещё большего объёма.

5.3 Влияние на latency

Чтение KV cache увеличивает latency шага генерации, особенно когда последовательность длинная. Это мешает приложениям реального времени (чат-боты, голосовые ассистенты).

5.4 Неэффективность стандартного хранения

Обычная реализация хранит кэш в непрерывном блоке памяти для каждого слоя. При изменении длины (увеличении на один токен) требуется копирование или предварительное выделение максимального размера (что приводит к трате памяти на padding). Решения вроде PagedAttention (vLLM) разбивают кэш на страницы фиксированного размера, динамически выделяя память и уменьшая фрагментацию.

6. Решения для смягчения bottleneck

6.1 Grouped Query Attention (GQA)

В GQA ключи и значения вычисляются не для каждой головы запроса (Query), а для группы голов. Например, при num_key_value_heads = 8 (вместо 64) размер KV cache уменьшается в 8 раз. Для LLaMA-3-70B это даёт 336 / 8 = 42 ГБ — уже помещается в один A100 80GB. За счёт некоторого снижения качества достигается значительная экономия. Multi-Query Attention (MQA) — крайний случай, когда используется одна общая голова для ключей и значений.

6.2 Квантизация KV cache

Хранение K и V в форматах с меньшей точностью (INT8, INT4). Квантизация может быть статической (применяется ко всем элементам) или адаптивной (например, per-token или per-channel). Это уменьшает размер в 2–4 раза (INT8: 42 ГБ → 21 ГБ; INT4: ~10.5 ГБ), но требует дополнительных этапов деквантизации при вычислениях, что может замедлить шаг.

6.3 PagedAttention (vLLM)

KV cache разбивается на блоки (страницы) фиксированного размера (обычно 16–32 токена). Страницы выделяются по мере необходимости, что резко сокращает фрагментацию и позволяет эффективнее использовать память при batch-обработке (разные запросы разной длины). Это не уменьшает объём кэша, но улучшает его утилизацию.

6.4 FlashAttention / FlashAttention-2

Хотя FlashAttention оптимизирует само вычисление attention (снижая количество чтений/записей HBM), он не уменьшает размер KV cache напрямую. Однако он позволяет быстрее обрабатывать существующий кэш, частично снимая bottleneck по bandwidth.

6.5 Удаление части кэша (eviction)

Для очень длинных контекстов (миллионы токенов) можно использовать стратегии вытеснения: хранить только последние N токенов, агрегировать важные токены в summary-векторы (например, H2O (Heavy Hitter Oracle), StreamingLLM). Но это может влиять на качество.

7. Сравнение подходов (таблица)

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

В Agentic RAG агент может делать много шагов: запрос к retrieval'ю, вызов инструмента, чтение длинного документа, итеративная генерация. Каждый шаг добавляет новые токены в контекст, и KV cache растёт. Если агент использует рекурсию или длинные цепочки мыслей, длина последовательности может достигать десятков-сотен тысяч токенов. Без эффективного управления KV cache такой агент будет работать крайне медленно или не помещаться в память. Поэтому знание механизмов GQA, квантизации и PagedAttention критично для проектирования масштабируемых агентных систем.

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

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

Инструменты Python (numpy, matplotlib), без фреймворков — только расчёты.

Шаги:

1. Определить класс ModelConfig с полями: num_layers, num_heads, head_dim, seq_len, dtype_bytes, num_key_value_heads (для GQA).
2. Реализовать функцию kv_cache_size(config) -> bytes.
3. Построить график зависимости размера от seq_len для трёх конфигураций: MHA (num_heads=64), GQA (num_key_value_heads=8), GQA+INT4.
4. Оценить время чтения всего кэша за один шаг: time = size / bandwidth (bandwidth взять 2 TB/s для A100).
5. Вывести таблицу для seq_len = 512, 4096, 32768, 131072.

Ожидаемый результат График, наглядно показывающий, как GQA и квантизация снижают размер, и осознание, что при 128k токенов даже с GQA+INT4 размер около 10 ГБ, что приемлемо.

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

Навигация

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