MQA и GQA — делим ключи между головами

Зачем это нужно

В уроке 4.6.4 вы ускорили генерацию с помощью KV-кэшаKV-кэшEN KV cache: ключи и значения прошлых токенов не пересчитываются, а хранятся. За скорость платим памятью, и на длинном контексте этот счёт становится огромным.

Возьмём Llama 3.1 70B с её контекстом 128 тысяч токенов. Веса в BF16 занимают около 140 ГБ. KV-кэш одного запроса на всю длину контекста — ещё 40 ГиБ. И это уже с экономией из этого урока. Если бы модель использовала классическое многоголовое вниманиемногоголовое вниманиеEN multi-head attention (MHA), кэш одного такого запроса весил бы 320 ГиБ, более чем вдвое больше всей модели. Восемь пользователей с длинными документами заняли бы несколько серверов только под кэш.

Поэтому все современные LLM экономят на кэше. Этот модуль проходит путь от простых решений к изобретению DeepSeek: сначала MQA и GQA (этот урок), затем MLA (следующий). Первые два шага сделали в Google, а последний, самый смелый, — китайская лаборатория DeepSeek.

Сколько весит KV-кэш

На каждом слое для каждого токена хранится ключ и значение каждой KV-головы. Формула из урока 4.6.4, но теперь с числом KV-голов $G$, которое может быть меньше числа голов запросов $H$:

$$\text{KV-кэш} = 2 \cdot L \cdot G \cdot d_h \cdot T \cdot B \cdot \text{(байт на число)}$$

Здесь $L$ — число слоёв, $d_h$ — размер головы, $T$ — длина контекста, $B$ — батч, двойка — это K и V. Посчитаем для Llama 3 8B: $L = 32$, $G = 8$, $d_h = 128$, BF16 (2 байта):

$$2 \cdot 32 \cdot 8 \cdot 128 \cdot 2 = 131\,072 \text{ байт} = 128 \text{ КиБ на токен}$$

На контексте 8192 токена это ровно 1 ГиБ ($2^{30}$ байт), на 128K — 16 ГиБ. Сами веса модели — около 16 ГБ. На длинном контексте кэш одного запроса весит столько же, сколько вся модель.

Конфигурации реальных моделей (BF16, ГиБ = $2^{30}$ байт):

Модель	Слоёв $L$	Голов Q / KV	Кэш на токен	32K токенов
Llama 3 8B	32	32 / 8	128 КиБ	4.0 ГиБ
Llama 3 / 3.1 70B	80	64 / 8	320 КиБ	10 ГиБ
Llama 3.1 405B	126	128 / 8	504 КиБ	15.8 ГиБ
Qwen2.5-7B	28	28 / 4	56 КиБ	1.75 ГиБ
Qwen2.5-72B	80	64 / 8	320 КиБ	10 ГиБ
Qwen3-8B	36	32 / 8	144 КиБ	4.5 ГиБ
Qwen3-235B-A22B	94	64 / 4	188 КиБ	5.9 ГиБ
DeepSeek 67B (2023)	95	64 / 8	380 КиБ	11.9 ГиБ
Llama 3 70B, если бы MHA	80	64 / 64	2.5 МиБ	80 ГиБ

У всех моделей в таблице $d_h = 128$. Qwen3-235B-A22B — MoE-модель, но на размер кэша это не влияет: кэш определяется слоем внимания, а эксперты заменяют FFN (модуль 5.3). Поиграйте с калькулятором: переключите тип внимания с GQA на MHA и посмотрите, как растёт кэш.

Калькулятор KV-кэшаинтерактив

Архитектура

СлоёвГолов запросаKV-головРазмер головыMLA: d_c (латент)MLA: d_rope

Нагрузка

Длина контекстаБатч (запросов)

Та же модель с разными схемами внимания (на токен)

MHA/GQA: на токен хранится 2 · слои · KV-головы · размер головы чисел (K и V). MLA (DeepSeek-V2/V3): на токен хранится только латент c_KV размера d_c и общий RoPE-ключ размера d_rope — (d_c + d_rope) · слои чисел.

Почему важна не только память, но и скорость

На шаге декодирования модель генерирует один токен. Для этого она читает из видеопамяти все веса и весь KV-кэш, а считает на каждый прочитанный байт очень мало.

Прикинем для внимания одной головы. Чтобы обработать $T$ закэшированных токенов, нужно прочитать $2 \cdot T \cdot d_h$ чисел (в BF16 это $4 T d_h$ байт) и выполнить около $4 T d_h$ операций: скалярные произведения с ключами и взвешенную сумму значений. Выходит около 1 FLOP на байт. А у H100 SXM примерно 989 TFLOPS в BF16 при пропускной способности памяти около 3.35 ТБ/с, то есть порядка 300 FLOP на байт. Вычислитель простаивает, пока ждёт данные: внимание при декодировании ограничено памятью (memory-boundограничено памятью (memory-bound)EN memory-boundОперация, которая дольше ждёт данные из памяти, чем считает. Ускоряется, если читать меньше байт (квантование, GQA, MLA) или делать больше работы на каждый прочитанный байт (батчинг).).

Отсюда вывод: время, которое внимание тратит на шаге декодирования, почти пропорционально объёму кэша. Уменьшили кэш в 8 раз — ускорили внимание почти в 8 раз. Кроме кэша, на каждом шаге читаются и веса модели, но их читают один раз на весь батч. Поэтому чем больше батч и длиннее контексты, тем большую часть шага занимает именно чтение кэша. Подробно о compute-bound и memory-bound — в уроке 7.1.1.

MQA: одна пара K/V на все головы

Ноам Шазир (Noam Shazeer, один из авторов «Attention Is All You Need») в 2019 году предложил радикальный шаг: пусть у всех $H$ голов будут свои запросы, но общие ключи и значения. Это Multi-Query Attention (MQA)многозапросное внимание (MQA)EN multi-query attention (MQA)Все головы запросов делят одну голову ключей и одну голову значений. KV-кэш меньше в H раз, но качество заметно падает (Shazeer, 2019)..

$$\text{head}_i = \mathrm{softmax}\!\left(\frac{Q_i K^\top}{\sqrt{d_h}}\right) V, \qquad i = 1, \dots, H$$

Матрицы $W_K$ и $W_V$ сжимаются из $d \times H d_h$ до $d \times d_h$, кэш уменьшается в $H$ раз. Каждый прочитанный ключ теперь используется всеми $H$ головами, и вычислений на байт становится в $H$ раз больше.

MQA применяли в PaLM и StarCoder. Но у неё есть цена: головы больше не могут «смотреть» на разные аспекты ключей. DeepSeek честно измерила эту цену в приложении D.1 к статье о DeepSeek-V2. Три плотные модели примерно по 7B параметров (число слоёв подогнали так, чтобы параметров было поровну) обучили на 1.33T токенов:

Бенчмарк	MQA	GQA (8 групп)	MHA
MMLU	37.9	41.2	45.2
C-Eval	30.0	37.7	42.9
CMMLU	34.6	38.4	43.5
BBH	33.2	35.6	37.0

Что это за тесты: MMLU — вопросы с вариантами ответа по 57 предметам (урок 7.3.1), C-Eval и CMMLU — его китайские аналоги, BBH (BIG-Bench Hard) — набор трудных задач на многошаговое рассуждение. Числа — доля верных ответов в процентах. Разница в 7 пунктов MMLU между MQA и MHA — это огромно. Нужна золотая середина.

GQA: золотая середина

Grouped-Query Attention (GQA)внимание с группами запросов (GQA)EN grouped-query attention (GQA)Головы запросов делятся на G групп, у каждой группы своя пара K/V. G = H — это MHA, G = 1 — MQA. Стандарт Llama 3 и Qwen (Ainslie и др., 2023). (Ainslie и коллеги из Google, 2023) делит $H$ голов запросов на $G$ групп. Внутри группы головы делят одну пару K/V. Голова запроса $i$ использует KV-голову номер $\lfloor i / (H/G) \rfloor$ (скобки $\lfloor \cdot \rfloor$ — округление вниз, // в Python). Пример: $H = 4$, $G = 2$ — головы запросов 0 и 1 берут KV-голову 0, головы 2 и 3 — KV-голову 1.

flowchart TB
subgraph A["MHA: 4 KV-головы"]
  a1["Q1"] --> ak1["KV1"]
  a2["Q2"] --> ak2["KV2"]
  a3["Q3"] --> ak3["KV3"]
  a4["Q4"] --> ak4["KV4"]
end
subgraph B["GQA: 2 KV-головы"]
  b1["Q1"] --> bk1["KV1"]
  b2["Q2"] --> bk1
  b3["Q3"] --> bk2["KV2"]
  b4["Q4"] --> bk2
end
subgraph C["MQA: 1 KV-голова"]
  c1["Q1"] --> ck["KV"]
  c2["Q2"] --> ck
  c3["Q3"] --> ck
  c4["Q4"] --> ck
end

Четыре головы запросов и разное число KV-голов. Кэшируются только KV-головы.

Два крайних случая — старые знакомые: $G = H$ это MHA, $G = 1$ это MQA. Кэш и чтение памяти уменьшаются в $H/G$ раз. В статье GQA с 8 группами почти догнала MHA по качеству и почти догнала MQA по скорости.

Ещё одна практическая причина выбрать GQA — тензорный параллелизм (урок 5.4.3): большую модель режут по головам между, например, 8 GPU. Единственную KV-голову MQA приходится копировать на каждую GPU, и экономия частично теряется. В GQA с 8 KV-головами каждая GPU хранит ровно свою. Неудивительно, что $G = 8$ стал стандартом: так устроены все размеры Llama 3 и все плотные модели Qwen3.

Как получить GQA из готовой модели

Обучать модель с нуля ради другого внимания дорого. Авторы GQA предложили uptraining:

1. Взять обученный MHA-чекпоинт.
2. В каждой группе усреднить проекции $W_K$ и $W_V$ входящих в неё голов (mean-pooling). Это сработало лучше, чем взять одну голову группы или инициализировать заново.
3. Дообучить модель примерно на 5% исходного вычислительного бюджета.

Тот же приём «сжать готовую модель и чуть дообучить» вы встретите в следующем уроке: так переводят GQA-модели на MLA.

Как это выглядит в коде

Внутри GQA — обычное многоголовое внимание из урока 4.4.3. Отличие одно: проекции K и V выдают $G$ голов, а перед умножением их «размножают» до $H$. В коде Llama, Qwen и MiniMind это функция repeat_kv, эквивалент torch.repeat_interleave по оси голов.

repeat_interleave, а не repeat

Головы должны повторяться подряд: [KV0, KV0, KV1, KV1], а не [KV0, KV1, KV0, KV1]. Иначе голова запроса 1 получит ключи чужой группы. Код не упадёт, но модель, загруженная из чекпоинта, начнёт выдавать мусор. Тест в задании ниже ловит этот баг отдельно.

Важная деталь: копия нужна только для удобства вычислений. В кэше лежат $G$ голов, а эффективные ядра вроде FlashAttention (они считают внимание блоками, не собирая всю матрицу оценок в памяти) вообще не копируют данные: несколько голов запросов просто читают одну и ту же KV-голову. В PyTorch 2.5+ F.scaled_dot_product_attention(..., enable_gqa=True) принимает K и V с меньшим числом голов напрямую.

Практика

Задание в браузере

Реализуйте repeat_kv, причинное GQA и функцию размера KV-кэша. Тесты проверят крайние случаи MHA и MQA, правильное сопоставление голов, причинность и цифры для Llama 3.

GQA через повторение KV-головnumpy · нужен интернет7 тестов

import numpy as np


def repeat_kv(kv, n_rep):
    """Повторяет каждую KV-голову n_rep раз подряд.

    kv: массив формы (G, T, d) — G голов ключей (или значений).
    Возвращает массив (G * n_rep, T, d), где голова j результата — копия головы j // n_rep.
    Порядок важен: для G=2, n_rep=3 головы идут как [0, 0, 0, 1, 1, 1], а не [0, 1, 0, 1, 0, 1].
    """
    # TODO: одна строка с np.repeat (подумайте, по какой оси)
    ...


def gqa_attention(q, k, v):
    """Причинное grouped-query attention.

    q: (H, T, d) — H голов запросов;
    k, v: (G, T, d) — G голов ключей и значений, H делится на G.
    Возвращает (H, T, d). Масштаб 1/sqrt(d), причинная маска: позиция t видит только позиции <= t.
    При G == H это обычное MHA, при G == 1 — MQA.
    """
    H, T, d = q.shape
    G = k.shape[0]
    # TODO: 1) повторите k и v до H голов; 2) scores = q @ k^T / sqrt(d);
    #       3) замаскируйте будущее (-inf выше диагонали); 4) устойчивый softmax; 5) веса @ v
    ...


def kv_cache_bytes(n_layers, n_kv_heads, head_dim, seq_len, batch=1, bytes_per_elem=2):
    """Размер KV-кэша в байтах: K и V для каждого слоя, KV-головы, позиции и примера в батче."""
    # TODO
    ...


# Проверьте себя кнопкой «Запустить»:
rng = np.random.default_rng(0)
q = rng.normal(size=(8, 4, 16))
k = rng.normal(size=(2, 4, 16))
v = rng.normal(size=(2, 4, 16))
r = repeat_kv(np.arange(2).reshape(2, 1, 1), 3)
print(None if r is None else r.ravel())                      # ожидается [0 0 0 1 1 1]
out = gqa_attention(q, k, v)
print(None if out is None else out.shape)                    # ожидается (8, 4, 16)
gib = kv_cache_bytes(32, 8, 128, 8192)
print(gib if gib is None else gib / 2**30, "ГиБ")             # Llama 3 8B, 8K токенов: 1.0 ГиБ

ЗапуститьПроверитьПодсказкаСбросить

То же на PyTorch (запустите локально)

Модуль GQA, сравнение числа параметров MHA/GQA/MQA при размерах Llama 3 8B, проверка встроенного enable_gqa и конвертация MHA → GQA по рецепту статьи.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)


class GQA(nn.Module):
    """Grouped-Query Attention: n_heads голов запросов, n_kv_heads голов ключей/значений."""

    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        assert n_heads % n_kv_heads == 0
        self.h, self.g, self.dh = n_heads, n_kv_heads, d_model // n_heads
        self.wq = nn.Linear(d_model, n_heads * self.dh, bias=False)
        self.wk = nn.Linear(d_model, n_kv_heads * self.dh, bias=False)  # в H/G раз уже, чем в MHA
        self.wv = nn.Linear(d_model, n_kv_heads * self.dh, bias=False)
        self.wo = nn.Linear(n_heads * self.dh, d_model, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.h, self.dh).transpose(1, 2)  # (B, H, T, dh)
        k = self.wk(x).view(B, T, self.g, self.dh).transpose(1, 2)  # (B, G, T, dh) ← это и кэшируется
        v = self.wv(x).view(B, T, self.g, self.dh).transpose(1, 2)
        rep = self.h // self.g
        k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)  # (B, H, T, dh)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.wo(y.transpose(1, 2).reshape(B, T, -1))


# 1) Сколько параметров внимания в одном слое при d_model = 4096, 32 головы (как в Llama 3 8B)
for name, g in [("MHA", 32), ("GQA-8", 8), ("MQA", 1)]:
    n = sum(p.numel() for p in GQA(4096, 32, g).parameters())
    print(f"{name:6s} {n / 1e6:5.1f} M параметров, KV-голов: {g}")

# 2) Встроенная поддержка GQA в PyTorch ≥ 2.5 даёт то же самое без копирования голов
m = GQA(256, 8, 2)
x = torch.randn(2, 10, 256)
B, T, _ = x.shape
q = m.wq(x).view(B, T, 8, 32).transpose(1, 2)
k = m.wk(x).view(B, T, 2, 32).transpose(1, 2)
v = m.wv(x).view(B, T, 2, 32).transpose(1, 2)
y = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)
y = m.wo(y.transpose(1, 2).reshape(B, T, -1))
print("enable_gqa совпадает:", torch.allclose(y, m(x), atol=1e-5))

# 3) «Uptraining» по рецепту статьи GQA: усредняем K- и V-головы внутри группы
mha = GQA(256, 8, 8)
gqa = GQA(256, 8, 2)
with torch.no_grad():
    for src, dst in [(mha.wk, gqa.wk), (mha.wv, gqa.wv)]:
        w = src.weight.view(2, 4, 32, 256)          # (группы, голов в группе, dh, d_model)
        dst.weight.copy_(w.mean(dim=1).reshape(64, 256))
    gqa.wq.weight.copy_(mha.wq.weight)
    gqa.wo.weight.copy_(mha.wo.weight)
err = (gqa(x) - mha(x)).norm() / mha(x).norm()
print(f"относительная ошибка после mean-pool: {err:.2f} — дальше нужно дообучение")

Ожидаемый вывод (PyTorch 2.5+):

MHA     67.1 M параметров, KV-голов: 32
GQA-8   41.9 M параметров, KV-голов: 8
MQA     34.6 M параметров, KV-голов: 1
enable_gqa совпадает: True
относительная ошибка после mean-pool: 0.87 — дальше нужно дообучение

🖥 На вашей машине (4×GPU)

Проверьте формулу кэша на живой модели. Загрузите Qwen3-8B в BF16 через transformers (около 16 ГБ — одна карта) и под torch.no_grad() прогоните через неё 2000 токенов: out = model(ids, use_cache=True). Оставьте только кэш (cache = out.past_key_values, затем del out — в out лежат ещё и логиты, они весят больше кэша) и сравните torch.cuda.memory_allocated() до и после. Прирост должен совпасть с нашей формулой: 144 КиБ на токен, то есть примерно 281 МиБ на 2000 токенов. Потом замерьте время одного шага декодирования при длине контекста 1K и 32K: оно растёт вместе с кэшем, как и обещает оценка memory-bound. Третий опыт — на вашей GPT из урока 4.6.2: замените в ней слой внимания на модуль GQA из кода выше и обучите три модели с 8, 2 и 1 KV-головой (число голов запросов не меняйте). Удобно запустить их одновременно, по одной на карту: CUDA_VISIBLE_DEVICES=0 python train.py ..., CUDA_VISIBLE_DEVICES=1 .... Сравните val loss и скорость генерации. Это ваш собственный маленький вариант таблицы MQA/GQA/MHA из статьи DeepSeek. На CPU формулу кэша можно проверить только на игрушечных моделях, а на карте вы видите её на настоящей 8B-модели и чувствуете, как время шага декодирования упирается в чтение кэша из видеопамяти. Настройка карт — на странице «Ваша лаборатория (4×GPU)».

Обратите внимание: GQA экономит и параметры (−38% в слое внимания), но главное — кэш, который меньше в 4 раза. Большая ошибка после mean-pool на случайных весах ожидаема: случайные головы ничем не похожи друг на друга. Именно поэтому после усреднения модель дообучают.

Китайский угол: как Qwen и DeepSeek пришли к GQA и почему DeepSeek пошла дальше

Qwen. Все плотные модели Qwen2.5 и Qwen3 используют GQA, причём группы бывают крупнее, чем у Llama: у Qwen2.5-7B 28 голов запросов на 4 KV-головы (7 к 1), у флагманской MoE-модели Qwen3-235B-A22B — 64 на 4 (16 к 1). В итоге кэш на токен у 235-миллиардной Qwen3 (188 КиБ) меньше, чем у Llama 3 70B (320 КиБ).

DeepSeek. Первая крупная модель DeepSeek LLM 67B (2023) тоже использовала GQA с 8 группами. Параметры, сэкономленные на K и V, обычно отдают в ширину FFN, а DeepSeek, по словам авторов, ради качества вложила их в глубину: у модели 95 слоёв. А затем DeepSeek поставила эксперимент из таблицы выше и увидела, что GQA и MQA заметно проигрывают MHA на сложных бенчмарках. Так появилась цель: кэш меньше, чем у GQA, а качество не хуже MHA. Ответ — MLA, тема следующего урока.

Что почитать на китайском. Су Цзяньлинь (автор RoPE) в статье «Предельное перетягивание каната между кэшем и качеством» проводит линию MHA → MQA → GQA → MLA одним рассуждением. На наш взгляд, это один из лучших текстов по теме модуля на любом языке. Читайте его с машинным переводом и глоссарием курса. А в MiniMind GQA реализована в несколько строк: по умолчанию 8 голов запросов и 4 KV-головы, функция repeat_kv в model/model_minimind.py.

Проверь себя

Проверь себя6 вопросов·

1 Сколько КиБ занимает KV-кэш Llama 3 8B на один токен? 32 слоя, 8 KV-голов по 128, BF16.

2 У модели 64 головы запросов и 8 KV-голов. Какую KV-голову использует голова запроса номер 20 (нумерация с нуля)?

3 Какие утверждения верны?Можно выбрать несколько ответов

GQA с числом групп $G = H$ — это обычное MHAMQA уменьшает размер матрицы $W_Q$GQA уменьшает и объём KV-кэша, и объём данных, читаемых из памяти на каждом шаге декодированияGQA ускоряет обучение ровно в $H/G$ раз

4 Почему при генерации одного длинного ответа внимание упирается в пропускную способность памяти?

Потому что softmax очень дорогойПотому что GPU не умеет умножать матрицы на векторыНа каждый новый токен нужно прочитать весь KV-кэш, а вычислений на прочитанный байт мало (порядка 1 FLOP/байт для MHA)Потому что KV-кэш хранится на диске

5 Во сколько раз GQA в Qwen3-235B-A22B (64 головы запросов, 4 KV-головы) уменьшает KV-кэш по сравнению с MHA той же формы?

6 Как авторы GQA превращают обученную MHA-модель в GQA?

Усредняют K- и V-проекции голов внутри группы и дообучают модель на ~5% исходного бюджетаОставляют первую голову каждой группы, остальные удаляют, без дообученияОбучают модель с нуляКвантуют кэш в INT4

Проверить ответы

Итоги

• KV-кэш $= 2 \cdot L \cdot G \cdot d_h \cdot T \cdot B \cdot$ байт. На длинном контексте он сравним с весами модели или больше их.
• Декодирование ограничено памятью: время внимания на каждом шаге пропорционально объёму кэша, который нужно прочитать.
• MQA ($G = 1$) сжимает кэш в $H$ раз, но заметно теряет в качестве. GQA ($1 < G < H$) — стандартный компромисс, обычно $G = 8$ или меньше.
• MHA-модель можно перевести в GQA усреднением голов и коротким дообучением.

Дальше: DeepSeek нашла способ сделать кэш меньше, чем у GQA, и при этом не потерять качество MHA. Это Multi-head Latent Attention — следующий урок.