KV-cache

KV-cache (key-value cache) — это техника, которая позволяет LLM переиспользовать уже вычисленные векторы ключей и значений для прошлых токенов вместо того, чтобы пересчитывать их заново на каждом шаге генерации.

Это одна из ключевых оптимизаций для инференса: без нее генерация текста становится заметно медленнее, особенно при длинных контекстах и больших моделях.

В этой заметке я собрал основные идеи из статьи Sebastian Raschka про KV-cache и перевел их в более практический, инженерный язык.

Почему без кэша генерация медленная

Большинство LLM генерируют текст по одному токену за раз: модель принимает контекст, предсказывает следующий токен, добавляет его в последовательность и повторяет.

Пример:

prompt: “Time”
следующий токен: “flies”
затем: “fast”

Если модель не использует кэш, то на каждом шаге она заново обрабатывает весь уже сгенерированный контекст:

для “Time” вычисляет все необходимые промежуточные представления
для “Time flies” снова всё пересчитывает
для “Time flies fast” снова всё пересчитывает

То есть при генерации третьего токена модель уже дважды пересчитывает то, что было вычислено на предыдущих шагах.

С точки зрения attention это особенно заметно: каждую новую позицию модель сравнивает с предыдущими ключами и значениями. Но для уже увиденных токенов эти K/V-представления не меняются — они были рассчитаны раньше и не должны пересчитываться.

Что такое key и value в attention

В self-attention каждый токен преобразуется в три вектора:

Query (Q): то, что модель ищет в контексте
Key (K): то, что модель “помечает” как доступный для поиска
Value (V): содержимое, которое будет использоваться для агрегации

Для одного токена в каждом слое:

Q, K, V получаются линейными проекциями из эмбеддинга токена
далее модель вычисляет скоры attention между текущим query и всеми ключами
потом весирует значения по этим скорям

Уже вычисленные K и V для прошлых токенов зависят только от самих токенов и параметров модели, а не от того, что генерируется дальше.

Именно поэтому их можно сохранить и переиспользовать.

Идея KV-cache

KV-cache хранит уже вычисленные:

ключи K для каждого токена
значения V для каждого токена

на каждом слое и в каждом attention-head.

При генерации следующего токена модель:

вычисляет K и V только для нового токена
добавляет их в кэш
использует весь кэш вместе с новым токеном для attention

Вместо пересчета всей последовательности целиком модель работает только с:

новым токеном
уже сохраненными K/V для прошлого контекста

Именно это резко снижает вычислительную нагрузку.

Почему это важно в авторегрессионной генерации

LLM обычно генерируют текст автререгрессионно:

исходный prompt обрабатывается один раз
далее на каждом шаге выбирается один следующий токен
новый токен добавляется в последовательность

Если не использовать cache, то на каждом step мы пересчитываем историю заново. Это дает квадратичную зависимость по длине последовательности на уровне внимания.

С KV-cache мы в сущности уменьшаем повторную работу: каждая пара K/V вычисляется один раз и дальше используется много раз.

Упрощенно:

без кэша: O(n²) работы по накоплению
с кэшем: в каждом шаге мы добавляем только новые элементы, а старые переиспользуем

Это не магия: в реальных моделях это превращается в значимый выигрыш в latency и throughput, особенно в продакшене и при длинных диалогах, код-интерфейсах, генерации длинного текста и RAG-пайплайнах.

Идея на примере

Представим, что мы уже сгенерировали контекст:

“Time flies”

Теперь генерируем следующий токен “fast”.

Без cache модель заново делает:

вычисление key/value для “Time”
вычисление key/value для “flies”
затем для “fast”

С cache:

“Time” и “flies” уже лежат в кэше
вычисляется только key/value для “fast”
дальше attention смотрит сразу на кэш + новый токен

Это тот же результат математики, но без лишних вычислений.

Как выглядит реализация на уровне кода

Идея очень простая: хранить тензоры cache_k и cache_v и дополнять их новыми значениями.

Пример в духе кода из статьи:

import torch

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_in, d_out, num_heads):
        super().__init__()
        self.W_key = torch.nn.Linear(d_in, d_out)
        self.W_value = torch.nn.Linear(d_in, d_out)
        self.W_query = torch.nn.Linear(d_in, d_out)

        self.register_buffer("cache_k", None)
        self.register_buffer("cache_v", None)

    def reset_cache(self):
        self.cache_k = None
        self.cache_v = None

    def forward(self, x, use_cache=False):
        b, seq_len, d_in = x.shape

        keys_new = self.W_key(x)
        values_new = self.W_value(x)
        queries = self.W_query(x)

        if use_cache:
            if self.cache_k is None:
                self.cache_k = keys_new
                self.cache_v = values_new
            else:
                self.cache_k = torch.cat([self.cache_k, keys_new], dim=1)
                self.cache_v = torch.cat([self.cache_v, values_new], dim=1)

            keys = self.cache_k
            values = self.cache_v
        else:
            keys = keys_new
            values = values_new

        # далее обычный attention
        # scores = queries @ keys.transpose(-2, -1)
        # weights = softmax(scores, dim=-1)
        # output = weights @ values
        return keys, values

Главная мысль здесь в том, что мы не заменяем сам attention, а лишь изменяем способ, каким модель получает K и V.

Обязательный паттерн: очистка кэша

Кэш нельзя просто держать между разными запросами или отдельными генерационными сессиями.

Если этого не делать, то новый запрос может начать “видеть” старый контекст, что приведет к неправильному поведению и запутанным результатам.

Поэтому обычно есть метод вроде:

model.reset_kv_cache()

или на уровне блока attention:

for block in model.trf_blocks:
    block.att.reset_cache()

Это особенно важно в продакшене, где один и тот же объект модели может использоваться для разных разговоров или запросов.

Почему нужен offset / current_pos

В простом коде при каждом новом вызове нужно помнить, на какой позиции в контексте мы находимся.

Например, если кэш уже содержит 10 токенов, а новый токен должен пойти после них, то позиционный индекс новых токенов должен начинаться не с 0, а с 10.

Иначе модель “смешает” старые и новые позиции и будет ошибочно учитывать их как перекрывающиеся.

Типичный паттерн:

self.current_pos = 0

и затем при генерации:

pos_ids = torch.arange(self.current_pos, self.current_pos + seq_len)
self.current_pos += seq_len

Это важно для корректного position embedding и для того, чтобы модель не путала временные шаги.

Почему кэш помогает в реальном инференсе

Кэш особенно полезен в системах, где:

пользователь пишет длинный диалог
модель должна быстро продолжать ответ
используется streaming generation
есть большие модели и длинный контекст

Если без кэша при каждом новом токене модель заново прогоняет весь контекст, то задержка быстро растет. С кэшем повторные вычисления исчезают, и скорость генерации становится заметно выше.

Это одна из причин, почему современные LLM-инференс-системы в продакшене используют кэширование K/V как стандартную практику.

Трейд-оффы

У KV-cache есть плюсы и недостатки

Плюсы

заметно ускоряет генерацию текста
уменьшает повторные вычисления
особенно полезен при длинных последовательностях

Минусы

увеличивает потребление памяти
с длинным контекстом кэш может значительно расти
добавляет сложность в реализации
кэш нужен только для inference, не для training

Для больших моделей с очень длинным контекстом память становится критичным фактором. Поэтому в продакшене часто используются оптимизации:

предаллокейшн памяти
sliding window для ограничения контекста
отсечение старых токенов
более специальные наборы K/V буферов

Оптимизации в продакшене

В учебном примере кэш часто делается через torch.cat. Это удобно и понятно, но на больших моделях это может быть неэффективно, потому что:

создаются новые тензоры на каждом шаге
возникает лишнее выделение памяти
возможна фрагментация памяти

Более продвинутые варианты:

Предаллоцировать буфер максимальной длины
Записывать новые значения в нужный слот вместо конкатенации
Использовать sliding window: держать только последние N токенов
Для длинных контекстов применять более сложные стратегии управления memory

То есть в реальных системах KV-cache — это не просто “один словарь с K и V”, а обычно продуманная структура памяти, обученная под модель и устройство.

Что важно запомнить

в autoregressive LLM прошлые токены не нужно пересчитывать заново
их K/V уже известны и стабильны
кэш позволяет переиспользовать их на каждом следующем шаге

Это превращает генерацию текста из повторного пересчета всей истории в incremental inference, где каждый новый токен добавляет лишь небольшую дополнительную работу.

Именно эта идея лежит в основе производительного инференса современных больших языковых моделей.

Итог

KV-cache — это один из самых важных инженерных приемов в LLM-инференсе.

Он решает одну и ту же проблему, с которой сталкивается любая генерация текста:

почему модель заново считает уже известный контекст?
почему скорость падает при длинных диалогах?

Ответ в том, что без кэша она фактически заново вычисляет старый контекст каждый раз.

KV-cache меняет ситуацию так, что:

прошлые ключи и значения сохраняются
новый токен обрабатывается отдельно
attention обращается к уже готовому контексту

Именно поэтому кэш считается фундаментальной оптимизацией для эффективной работы LLM в production.

Ссылки и рекомендуемое чтение

статья: Understanding and Coding the KV Cache in LLMs from Scratch
про attention: Understanding and Coding Self-Attention, Multi-Head Attention, Causal-Attention, and Cross-Attention in LLMs
по книге: Build a Large Language Model (From Scratch)
[lists/llm]

AI-warning

✨ Данная статья сгенерирована с помощью MAI-Code-1.1-Flash. Текст проверен человеком.