KV cache — кэш ключей и значений в трансформерах: ускорение генерации и экономия памяти

Префилл vs декод: где живёт кэш

* Prefill («чтение промпта») — модель прогоняет всю подсказку (N токенов) батчем, строит K/V для каждого токена и кладёт их в кэш. * Decode («пошаговая генерация») — на каждом шаге добавляется одна пара K/V для нового токена; Q текущего шага «смотрит» на весь кэш.

Интуитивно: префилл — это «построить библиотеку», декод — «брать книги с полки и класть одну новую».

Сколько памяти ест KV cache (формула)

Память пропорциональна длине контекста и числу «KV-голов»:

KV_bytes = n_layers × seq_len × n_kv_heads × head_dim × (K+V) × sizeof(dtype) = n_layers × seq_len × n_kv_heads × head_dim × 2 × sizeof(dtype)

Где: * n_layers — число блоков трансформера; * seq_len — текущая длина кэша (промпт + сгенерированные токены); * n_kv_heads — число голов, для которых хранится K/V; * head_dim — размерность головы внимания; * dtype — обычно FP16/BF16 (2 байта), реже FP8/INT8 (см. квантование ниже).

Пример (ориентир). Модель 32 слоя, n_kv_heads=8 (GQA), head_dim=128, seq_len=4096, dtype=fp16:
32 × 4096 × 8 × 128 × 2 × 2 ≈ 0,5 GiB на одну последовательность.

Важно: при beam_width > 1 KV кэш размножается на каждую ветку (если не используется «общий префикс»), а при батче B память растёт ~линейно по B.

MQA, GQA и экономия кэша

Классическое внимание хранит K/V на каждую из h голов (n_kv_heads = h). Современные варианты уменьшают кэш:

• MQA (Multi-Query Attention) — одна пара K/V на все Q-головы слоя → n_kv_heads = 1. Экономия в ~h раз по кэшу, часто с небольшой потерей качества.
• GQA (Grouped-Query Attention) — K/V общие для группы из g Q-голов → n_kv_heads = h / g. Компромисс между качеством и памятью.

Paged / блоковый кэш и «упаковка» последовательностей

На практике «выросший» кэш фрагментируется: пользователи приходят/уходят, последовательности разной длины. Чтобы не копировать гигантские массивы и не страдать от фрагментации: * применяют блочный (paged) KV cache — кэш разбивается на страницы/блоки, которые переиспользуются между запросами;

• sequence packing — несколько коротких последовательностей упаковывают в один тензор с масками (экономит префилл и увеличивает загрузку GPU);
• prefix caching — общий префикс (системный промпт, инструкция) кэшируется один раз и шарится между запросами, если позиции совместимы.

Эти идеи лежат в основе высокопроизводительных рантаймов и планировщиков декода.

Где ускорение, а где «узкие места»

• Prefill хорошо векторизуется на GPU: большой матричный GEMM.
• Decode — шаги по одному токену: скорость ограничена памятью (доступ к кэшу), а не только FLOPs. Правильная организация KV (paged, плотная компоновка, кэш-локальность) критична.

Стриминг и длинный контекст

Большое контекстное окно — это больше памяти под KV. Есть техники, которые удерживают скорость и качество:

• Скользящее окно (sliding-window attention) — хранить «последние M токенов» + «якоря» начала (attention sinks).
• RoPE-скейлинг/интерполяция — корректирует позиционные фазы, чтобы не деградировать на больших позициях при усечении.
• Резюме прошлого — компрессия далёкой истории в «свертки»/саммари и их периодическое обновление.
• Chunked prefill — «кормить» большие документы порциями, постепенно наращивая KV.

Компромисс: меньше памяти ↔ риск потери дальних зависимостей.

Оффлоад и гибридная память

Если GPU не хватает:

• KV на CPU/пинованной памяти — дешёво по объёму, но дороже по задержке (PCIe/NVLink); работают схемы «горячий хвост на GPU, холодный префикс на CPU».
• NVMe-подкачка — как крайняя мера; годится для архивных префиксов и оффлайн-задач.
• Микс dtype — хранить KV в FP8/INT8, а активные вычисления в FP16/BF16.

Квантование KV cache

Квантовать можно только кэш, не трогая веса слоёв:

• INT8/FP8 KV — часто даёт 2× экономию памяти с минимальным влиянием на качество.

* Тонкость — шкалы/кластеризация по головам/слоям, чтобы не «зажать» амплитуды. * Для чувствительных задач (код/арифметика) держите BF16 KV хотя бы на последних слоях.

Beam search, speculative decoding и KV

• Beam search: на разветвлении префикса KV клонится на каждый beam; по возможности держите shared-prefix KV и копируйте только «хвост».
• Speculative decoding: «быстрая» модель генерирует черновик, «медленная» проверяет; при принятии кусок K/V можно bulk-добавить в кэш основной модели; при откате — хвост отбрасывается (эффективно с блочным KV).

Частые причины OOM и просадок

• Недооценили память под KV: длинные контексты × батч × beam — линейный рост. Считайте по формуле и закладывайте запас.
• Фрагментация: «дырявые» аллокации → копирования и падение throughput. Лечится блоковым KV.
• Слишком высокий batch в декоде: узкое место — пропускная способность памяти, а не математика.
• Отсутствие prefix caching: системный промпт перерасходует префилл на каждый запрос.
• Агрессивное квантование KV: незаметные на коротких ответах артефакты «вылазят» на длинных и логико-арифметических задачах.

Практические рецепты

Мини-глоссарий

• Prefill — фазa построения KV для входного промпта.
• Decode — по-токенная генерация с использованием KV.
• Paged KV — блочный кэш, уменьшающий копирования и фрагментацию.
• MQA/GQA — варианты внимания с общими K/V на все или группы голов.
• Prefix caching — повторное использование KV для общих префиксов.
• Sliding window — внимание к последним M токенам + «якоря» начала.

См. также

• Контекстное окно
• RAG: извлечение и дополнение
• Top-k
• Top-p