Key-Value кеш — отличный пример того, как недостаточно предложить и теоретически обосновать даже такую сильную архитектурную идею, как трансформеры в LLM.

Замечание 1. Трансформеры — это архитектурное решение, которое было впервые представлено в знаковой статье «Attention Is All You Need», опубликованной в 2017 году исследователями из компании Google Brain и Google Research.

Сердце трансформера — это механизм внимания (Attention) на основе трех матриц: KEYS (ключи), QUERIES (запросы), VALUES (значения). Для матриц подбираются такие параметры в процессе обучения, чтобы максимизировать способность модели выявлять и использовать смысловые связи между элементами текста (или других данных) независимо от их удаления (расположения) друг от друга.

Замечание 2. Цели обучения матриц внимания:
- QUERIES (запросы) — какая информация ищется в данном контексте;
- KEYS (ключи) — какая информация имеется и насколько она важна для текущего запроса;
- VALUES (значения) — извлечение смысла, что требуется для текущего запроса.

Замечание 3. Для повышения предсказательной способности трансформеров было предложено многоголовое внимание (Multi-Head Attention). Вместо одного общего механизма внимания используются несколько параллельных «голов». Каждая голова обладает своим собственным независимым набором матриц QUERIES, KEYS и VALUES меньшей размерности. Они обучаются параллельно, что позволяет модели одновременно извлекать различные типы специфических знаний из контента. Например, одни головы могут улавливать семантические связи (смысл слов), другие — иерархические и синтаксические (грамматическую структуру предложения), а третьи — родственные связи или знаки препинания.

К сожалению, любая теория рано или поздно «разбивается» о практику. Таким примером стала алгоритмическая сложность работы трансформеров — $O(n^2)$.

Квадратичная сложность возникает при перемножении матриц QUERIES и KEYS.
Две главные проблемы:

1. Вычислительная сложность — если длина текста была 2 ($2 \times 2 = 4$), если увеличится в два раза и станет 4 ($4 \times 4 = 16$);
2. Память — эти увеличивающиеся объемы данных нужно хранить в памяти GPU.

Кеширование матриц KEYS и VALUES

Процесс генерации с KV Cache разделяется на две четкие фазы: префилл и декодирование.

Фаза префилла (Prefill):

• Модель получает входной промпт (начальный текст);
• Обрабатывает всю последовательность промпта целиком;
• Вычисляет и сохраняет в кэш матрицы Key и Value для каждого токена промпта;
• Выполняется однократно перед началом генерации.

Фаза декодирования (Decoding):

• Берется последний сгенерированный токен как Query;
• Из кэша извлекаются все предыдущие Key и Value;
• Вычисляется внимание между новым Query и закэшированными Key;
• Генерируется следующий токен;
• Обновляется кэш добавлением Key и Value для нового токена.

На изображении схематически этот процесс изображен сверху первый прогон - Prefill,
и снизу фаза Decoding

_Схема работы технологии KV-Cache. Источник: Arxiv.

По шагам из картинки (фаза декодирования):

1. Запрос и Ключи: Приходит запрос 1x64 и умножается на ключи 64x5 (где 64x4 — это старый кэш ключей, а 64x1 — новый ключ текущего токена). Получается ответ 1x5 (веса внимания). Он никуда не сохраняется, а сразу идет в следующий шаг.
2. Значения: Этот ответ 1x5 умножается на значения 5x64 (где 4x64 — старый кэш значений, а 1x64 — новое значение текущего токена). Получается ответ 1x64 — это финальный результат, который уходит дальше на генерацию слова, в кэш он не сохраняется.
3. Что на самом деле идет в кэш: В кэш сохраняются не результаты умножения, а те самые новые синие кусочки: новый ключ (64x1) дописывается в кэш ключей, а новое значение (1x64) дописывается в кэш значений.

Когда придет следующий токен, матрицы в кэше уже будут иметь размеры 64x6 (для ключей) и 6x64 (для значений).

Замечание 4 Обратили внимание почему не кешируется queries? QUERIES не кэшируются, потому что нужен только один раз для вычисления связей внутри текущего текста. Как только связи найдены и сохранены в KV-кэш, старые QUERIES удаляются, так как для каждого следующего слова потребуется совершенно новый запрос.

Таким образом, использование KV-кэша позволяет снизить вычислительную сложность генерации одного очередного токена с квадратичной заместо $O(n^2)$ до линейной $O(n)$

Когда кеширования не достаточно

Казалось бы: кэшируй K и V, экономь ресурсы GPU. Однако на практике за кешем нужно следить, чтобы не было переполнения памяти и других неприятных ситуаций. В памяти GPU текст растет динамически. Видеокарта не знает, сколько слов напишет пользователь, поэтому выделяет память «с запасом» огромными сплошными кусками.

Это приводит к двум проблемам:

• Фрагментация: Память превращается в «швейцарский сыр». Свободные мегабайты есть, но они разбиты на мелкие осколки, и новый запрос туда не помещается.
• Избыточность: До 60-80% дорогой памяти видеокарты проистекает вхолостую, просто резервируя пустые места «на всякий случай».

Чтобы решить это, инженеры оптимизировали и железо, и математику алгоритмов.

Решение 1: Оптимизация на уровне железа (Управление памятью)

• PagedAttention (используется в vLLM): Борьба с фрагментацией. Технология заимствовала логику у операционных систем и стала нарезать кэш на маленькие фиксированные «страницы» (например, по 16 токенов). Они могут лежать в разных углах памяти, но связаны виртуальной таблицей. Это устранило пустоты и позволило упаковывать в один GPU в 2-4 раза больше запросов.

• FlashAttention: Борьба со скоростью шины памяти. Вместо того чтобы гонять огромные матрицы Q, K, V из медленной основной памяти GPU в сверхбыструю память ядер (SRAM) и обратно, этот алгоритм делит матрицы на мелкие блоки. Вычисления происходят прямо «на лету» внутри ядер, что ускоряет генерацию в 2-3 раза.

Решение 2: Оптимизация на уровне алгоритма (Сжатие кэша)

В стандартном подходе (MHA) у каждой головы внимания (их десятки) свои матрицы Ключей (K) и Значений (V) — кэш получается гигантским. Архитектуру изменили:

• MQA (Multi-Query): Все десятки голов запросов (Q) делят между собой всего одну общую пару K и V. Кэш уменьшается на 90%, но модель теряет качество.

• GQA (Grouped-Query): Идеальный компромисс. Головы Q делятся на группы (например, по 4 или 8 головок), и каждая группа делит одну общую пару K и V. Кэш сжимается в разы, а качество ответов не падает. Это стандарт для Llama 3 и всех современных LLM.

Подведем итоги

1. Проблема и базовое решение (KV-кеш): Механизм внимания в трансформерах имеет квадратичную сложность $O(n^2)$, что требует огромных вычислений и памяти при генерации текста. Чтобы это исправить, используется KV-кеш: матрицы Ключей (K) и Значений (V) уже обработанных токенов сохраняются в памяти, чтобы не пересчитывать их заново. Запросы (Q) не кешируются, так как нужны только для текущего шага.
2. Проблемы наивного кеширования: Поскольку длина генерируемого текста заранее неизвестна, видеокарта выделяет память «с запасом» сплошными блоками. Это приводит к сильной фрагментации памяти («эффект швейцарского сыра») и простому резервированию до 60-80% дорогой VRAM впустую.
3. Оптимизация на уровне «железа»: Для борьбы с неэффективной работой памяти инженеры внедрили PagedAttention (разбивает кеш на фиксированные «страницы», как операционная система, устраняя фрагментацию) и FlashAttention (ускоряет вычисления в 2-3 раза, обрабатывая матрицы мелкими блоками прямо в сверхбыстрой памяти SRAM ядер GPU, минуя медленную шину).
4. Алгоритмическое сжатие кеша: Чтобы радикально уменьшить объем данных, архитектуру изменили с стандартной (MHA) на MQA (все головы запросов делят одну общую пару K и V) и GQA (головы запросов делятся на группы, и каждая группа делит свою пару K и V). GQA позволяет сжать кеш в разы без потери качества и является современным стандартом (например, в Llama 3).