KV-кэш

Наблюдение, из которого всё следует

При авторегрессионной генерации на шаге $n+1$ ключи и значения для позиций $1 \dots n$ уже были посчитаны на предыдущих шагах, и они не изменились: $k_i$ зависит только от токена $i$. Значит их можно сохранить.

$$\text{Attn}(q_{n+1}, \htmlData{k=cached}{K_{1:n}} \oplus \htmlData{k=fresh}{k_{n+1}}, \htmlData{k=cached}{V_{1:n}} \oplus \htmlData{k=fresh}{v_{n+1}})$$

Экономия арифметики огромна: без кэша проекции $K$ и $V$ пересчитываются для всех $n$ позиций на каждом шаге, с кэшем — для одной. Отношение растёт линейно по $n$ и при $n = 4096$ составляет примерно $\times 1366$.

Но за это платится памятью, и вот тут начинается интересное.

Сколько это байт

$\text{размер кэша} = 2 \cdot n \cdot L \cdot d_{\text{kv}} \cdot \text{байт}$

Двойка — за $K$ и $V$; $L$ — число слоёв; $d_{\text{kv}} = g \cdot d$ — суммарная ширина ключей. Ни одного множителя лишнего, и весь смысл в том, что $n$ входит линейно, а веса от $n$ не зависят вовсе.

Отсюда точка пересечения. Веса стека — примерно $12LD^2$ (четыре квадратные проекции внимания плюс MLP с четырёхкратным расширением), значит

$2nLd_{\text{kv}} = 12LD^2 \quad\Longrightarrow\quad n^\ast = \frac{6D^2}{d_{\text{kv}}} = 6D\cdot\frac{h}{g}$

Потяните длину контекста и посмотрите, когда синяя полоса перерастает серую.

7B: 32 слоя, 32 головы по 12870B: 80 слоёв, 64 головы по 128GPT-2: 12 слоёв, 12 голов по 64

длина контекста n

4096

групп ключей g

32

кэш

2048 МБ

веса

12288 МБ

сравняются при

24576 токенов

выигрыш против пересчёта

×1366

операций на элемент кэша

1 на элемент

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

Числа для $D = 4096$, $L = 32$, fp16:

конфигурация	на токен	при $n = 4096$	сравняются с весами
MHA, $g = 32$	$0.500$ МБ	$2.0$ ГБ	$24\,576$
GQA, $g = 8$	$0.125$ МБ	$0.5$ ГБ	$98\,304$
MQA, $g = 1$	$0.016$ МБ	$0.06$ ГБ	$786\,432$

Полмегабайта за токен — это та цифра, из-за которой длинный контекст дорог не по вычислениям, а по памяти на каждого пользователя одновременно.

Главное число: одна операция на элемент

Теперь самое важное, и это не про размер. Посчитаем арифметическую плотность внимания на шаге декодирования: сколько умножений-сложений приходится на один прочитанный из памяти элемент.

За шаг внимание в одном слое читает $2 n d_{\text{kv}}$ чисел кэша и делает $2nD$ умножений-сложений. Отношение:

$\frac{2nD}{2nd_{\text{kv}}} = \frac{D}{d_{\text{kv}}} = \frac{h}{g}$

При обычном MHA ($g = h$) это ровно единица. Одно умножение на одно прочитанное число.

Почему это приговор: у современных ускорителей отношение арифметической производительности к пропускной способности памяти — порядка ста–трёхсот операций на прочитанный байт. Значит при плотности $1$ вычислитель простаивает, ожидая память, и ускорение арифметики не даёт ничего. Генерация ограничена чтением кэша.

схема	$h/g$	плотность
MHA	$1$	$1$
GQA, $g = h/4$	$4$	$4$
GQA, $g = h/8$	$8$	$8$
MQA	$h$	$32$ при $h = 32$

Вот и полное объяснение, зачем существуют MQA и GQA — то, что урок 080 оставил без причины. Они не ускоряют арифметику (её столько же) и не уменьшают число голов (запросов по-прежнему $h$). Они делят объём читаемой памяти на $h/g$, и ровно во столько же раз поднимают плотность. Экономия кэша — приятное следствие, а причина в плотности.

Почему обучение — другая задача

Стоит сопоставить, потому что интуиция из обучения здесь подводит.

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

Причина в третьей строке. При обучении веса читаются один раз на весь батч, поэтому их чтение амортизируется; при генерации токен один, и каждое число весов и кэша используется ровно один раз. Умножение матрицы на вектор ограничено памятью по природе, а не из-за реализации.

Отсюда и батчирование запросов при сервинге: несколько независимых последовательностей одновременно возвращают плотность, потому что веса читаются один раз на всю группу. Кэш при этом не амортизируется — он у каждой последовательности свой, и именно он ограничивает, сколько запросов уместится.

Что ещё делают с кэшем

приём	идея	цена
GQA / MQA	группы вместо голов	немного выразительности
квантование кэша в int8	вдвое меньше байт	небольшая потеря качества
скользящее окно	хранить только последние $w$	забывание дальше $w$
PagedAttention	страницы вместо непрерывных буферов	сложность реализации
выгрузка в CPU	память дешевле	пропускная способность падает

Первые четыре снижают либо объём чтения, либо фрагментацию — то есть работают против того же узкого места. Последний пункт — противоположная ставка, и она проигрывает как раз потому, что проблема в пропускной способности, а не в ёмкости.

Итог

• Кэш превращает $O(n)$ пересчётов на шаг в $O(1)$ — выигрыш растёт линейно по контексту.
• Он растёт линейно по $n$, а веса не растут вовсе, поэтому существует вычислимая точка $n^\ast = 6D\,h/g$, где кэш становится больше модели.
• Арифметическая плотность внимания при декодировании равна $h/g$, то есть единице для MHA. Это и есть причина существования GQA и MQA.
• Генерация ограничена памятью, обучение — арифметикой. Приёмы для одного почти не помогают другому.