Сколько памяти съест модель: как GQA, окна и MLA меняют KV-кэш

Знакомая ситуация: вы подняли готовую модель, веса спокойно уместились в память GPU, а сервер падает с out-of-memory уже на десятом параллельном запросе. Виноват не размер модели, а KV-кэш. Его объём определяется тем, как устроен слой внимания, и выбрать это вы уже не можете: решение принято на этапе обучения.

Чекпоинт фиксирует архитектуру

Конфиг и веса задают число голов запросов и key-value голов, размер головы, шаблон локального внимания и размерности латентного представления. Заменить MHA на GQA после обучения значит изменить модель, и качество, скорее всего, придётся добирать дообучением. Заменить полное внимание на скользящее окно означает выкинуть контекст, который модель привыкла видеть.

Поэтому практика такая: открыть конфиг (в Hugging Face это AutoConfig), посмотреть число слоёв, num_attention_heads, num_key_value_heads, head_dim и поля про sliding window. Одна оговорка: поле с окном может лежать в конфиге, но быть отключено соседним флагом, так что проверяйте и значение, и флаг. Некоторые модели чередуют локальные и глобальные слои, и истина всегда в коде реализации.

Как посчитать кэш на салфетке

Персистентная память живёт между шагами декодирования, и именно она определяет, сколько запросов вы обслужите одновременно. Временные буферы вроде тайлов attention-скоров после слоя можно освободить. Для стандартного кэша формула такая: 2 × байты на элемент × слои × токены × число KV-голов × размер головы. Двойка означает ключи плюс значения. Для батча умножаем на его размер.

Возьмём условную модель: 32 слоя, 32 головы, head_dim 128, fp16, контекст 8192 токена. С классическим MHA один запрос занимает 4 ГиБ. Если у модели GQA с 8 KV-головами, это 1 ГиБ. При MQA с одной KV-головой всего 128 МиБ. Разница между «влезает 4 запроса» и «влезает 16» возникает без изменения числа параметров.

Смешанные схемы считаем послойно. Пусть у такой же GQA-модели из 32 слоёв 24 локальных с окном 4096, а 8 полных, контекст 32 768 токенов. Полное внимание везде даст 4 ГиБ, а смешанное — около 1,4 ГиБ. Но общий кэш при этом не ограничен: каждый новый токен всё равно добавляет память в полных слоях. Растёт он медленнее, а не перестаёт расти.

Что за вариант внимания

• MHA: у каждой головы запросов своя пара ключ-значение. Это самый дорогой кэш, а на декоде ещё и самое тяжёлое чтение памяти.
• MQA и GQA: головы запросов делят общие KV-головы. Кэш и параметры K/V-проекций падают ровно во столько раз, во сколько уменьшилось число KV-голов.
• Скользящее окно: слой помнит только последние W токенов, кэш можно держать в кольцевом буфере фиксированного размера. Позиции при этом считаются по логическому номеру токена, а не по слоту буфера.
• MLA: вместо полных ключей и значений хранится сжатый латентный вектор. Кэш меньше, зато на декоде больше вычислений. Это обычно выгодно, потому что декод упирается в пропускную способность памяти, а не в арифметику.

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

Что может изменить рантайм

Удобно разделить четыре вопроса: кто на кого смотрит (полное или оконное внимание), сколько KV-представлений выучено (MHA/GQA/MQA), что кэшируется (KV-тензоры или латент) и как это исполняется. Первые три решает чекпоинт, четвёртое — движок. FlashAttention считает точное внимание тайлами и не пишет огромные матрицы скоров в память: временные буферы и трафик падают, а логический кэш и число KV-голов остаются прежними. Paged attention хранит кэш блоками, которые не обязаны лежать подряд, и улучшает аллокацию и шаринг, но сам по себе размер кэша не уменьшает. Аллокатор при этом обязан соблюдать правило вытеснения модели: нельзя выбрасывать старые записи из полного слоя только потому, что соседние слои оконные.

Практический вывод

Перед выбором чекпоинта прочитайте конфиг и прикиньте кэш на целевой контекст и батч. Если два кандидата близки по качеству, берите тот, у которого KV-головы или окна дешевле: это прямо превращается в число одновременных запросов на карту. И не забывайте про padding: при плотных батчах вы можете платить за пустые позиции, а ragged-реализация этого избегает.

Ограничение у метода простое: расчёт даёт оценку persistent-памяти, но не учитывает веса, активации и накладные расходы аллокатора. Цифры выше иллюстративные, на реальной модели подставляйте значения из её конфига.