Многоголовое внимание

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

В прошлом уроке каждый токен получал одно распределение внимания: одна строка softmax, одно взвешенное среднее значений. Это как задать соседям по предложению ровно один вопрос. Но токену обычно нужно узнать несколько вещей сразу.

Возьмём фразу «Сестра вчера подарила мне рыжую кошку, и теперь она спит» и токен «спит». Ему полезно знать:

• кто спит — нужно посмотреть на «она» (а через неё — на «кошку»);
• когда — посмотреть на «теперь»;
• что было прямо перед ним — соседний токен, это важно для грамматики.

Одна голова может разделить вес, скажем, 50 на 50 между «она» и «теперь». Но тогда на выходе будет смесь двух векторов значений, по половине каждого, и оба сигнала размоются. Нужны несколько независимых «взглядов», каждый со своим softmax. Это и есть многоголовое вниманиемногоголовое вниманиеEN multi-head attention (MHA) (multi-head attention, MHA): $H$ голов, каждая задаёт свой вопрос, а результаты потом собираются вместе.

После урока вы реализуете MHA на NumPy без единого цикла по головам, посчитаете, во что обходится внимание на длинном контексте, и увидите, откуда растут MQA, GQA и MLA из модуля 5.2.

Идея: H голов в маленьких подпространствах

Каждая голова вниманияголова вниманияEN attention headОдна независимая копия внимания со своими проекциями размера d_h = d_model / H. Головы работают параллельно и смотрят на разные отношения между токенами. — это полноценное scaled dot-product attention со своими матрицами, только уже: размер головы $d_h = d_{\text{model}} / H$.

$$\mathrm{head}_h = \mathrm{Attention}\big(XW_Q^{(h)},\ XW_K^{(h)},\ XW_V^{(h)}\big), \qquad h = 1, \dots, H$$ $$\mathrm{MultiHead}(X) = \mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_H)\, W_O$$

Что здесь написано: $H$ раз считаем обычное внимание, каждый раз со своими $W_Q^{(h)}, W_K^{(h)}, W_V^{(h)}$ размера $d_{\text{model}} \times d_h$. Выходы голов (по $d_h$ чисел на токен) ставим рядом, получаем снова $H \cdot d_h = d_{\text{model}}$ чисел, и умножаем на выходную проекциювыходная проекция (W_O)EN output projectionМатрица, которая смешивает склеенные выходы голов и возвращает вектор размера d_model. Эквивалентно сумме вкладов голов, у каждой — свой блок W_O. $W_O$.

В оригинальном трансформере $d_{\text{model}} = 512$, $H = 8$, значит $d_h = 64$.

Сколько это стоит

Неожиданно: столько же, сколько одна большая голова. Все $W_Q^{(h)}$ вместе — это $H$ матриц $d_{\text{model}} \times d_h$, то есть ровно одна матрица $d_{\text{model}} \times d_{\text{model}}$. То же для $W_K$ и $W_V$, плюс $W_O$. Итого в слое внимания без смещений:

$$4\, d_{\text{model}}^2 \ \text{параметров, независимо от } H$$

Для оригинального трансформера это $4 \cdot 512^2 = 1\,048\,576$, для GPT-2 small $4 \cdot 768^2 = 2\,359\,296$ (плюс 3072 смещения, итого 2 362 368, ровно как в чекпоинтечекпоинт (контрольная точка)EN checkpointСохранённое состояние обучения: веса модели, состояние оптимизатора, номер шага. Позволяет продолжить обучение после сбоя. GPT-2 — файле с сохранёнными весами). Число операций на оценки тоже не растёт: $H$ голов по $T^2 d_h$ — это те же $T^2 d_{\text{model}}$. Мы не покупаем головы, а делим ширину между ними: вместо одного наблюдателя с 512-мерным взглядом — восемь наблюдателей с 64-мерными.

Головы в реальных моделях

Модель	$d_{\text{model}}$	Голов $H$	$d_h$	Особенность
Transformer base (2017)	512	8	64	исходная статья
GPT-2 small	768	12	64
GPT-3 175B	12 288	96	128
Llama 3 8B	4096	32	128	GQA: 8 KV-голов
Qwen2.5-7B	3584	28	128	GQA: 4 KV-головы
DeepSeek-V3	7168	128	128 (+64 под RoPE)	MLA; $H \cdot d_h = 16\,384 > d_{\text{model}}$
Kimi K2	7168	64	128 (+64 под RoPE)	MLA, как у DeepSeek-V3, но голов вдвое меньше

Две вещи видно сразу. Размер головы у современных моделей чаще всего 128, а растёт с размером модели в основном число голов. И $H \cdot d_h$ не обязано равняться $d_{\text{model}}$: у DeepSeek-V3 головы вместе шире модели, и $W_O$ просто переводит $16\,384 \to 7168$.

Как разрезать на головы: reshape и transpose

Хранить $H$ маленьких матриц и считать головы в цикле медленно. На практике делают одну большую проекцию $Q = XW_Q$ формы $(B, T, D)$ и договариваются, что столбцы $0 \dots d_h - 1$ принадлежат голове 0, следующие $d_h$ — голове 1 и так далее. Потом тензор перекладывают так, чтобы головы стали отдельной осью, как батч:

Шаг	Операция	Форма
1	$Q = XW_Q$ (то же для $K$, $V$)	$(B, T, D)$
2	reshape(B, T, H, d_h) — режем последнюю ось на $H$ кусков	$(B, T, H, d_h)$
3	transpose(0, 2, 1, 3) — головы вперёд	$(B, H, T, d_h)$
4	внимание во всех головах сразу: оценки	$(B, H, T, T)$
5	веса $\times V$	$(B, H, T, d_h)$
6	transpose(0, 2, 1, 3) и reshape(B, T, D) — склейка голов	$(B, T, D)$
7	$\times W_O$	$(B, T, D)$

Шаг 4 работает без изменений кода из прошлого урока: q @ np.swapaxes(k, -1, -2) умножает две последние оси, а оси $B$ и $H$ транслируются, как батч. Вся многоголовость — это перекладывание осей.

Главный баг: reshape сразу в (B, H, T, d_h)

Почему нельзя сделать reshape(B, H, T, d_h) за один шаг? reshape не двигает числа в памяти, а просто заново нарезает их подряд. Посмотрим на маленьком примере: 2 токена, $D = 6$, 3 головы по $d_h = 2$.

import numpy as np

B, T, D, H = 1, 2, 6, 3          # 2 токена, 6 признаков, 3 головы по d_h = 2
d_h = D // H
x = np.arange(B * T * D).reshape(B, T, D)
print(x[0])                       # токен 0: 0..5, токен 1: 6..11

right = x.reshape(B, T, H, d_h).transpose(0, 2, 1, 3)   # (B, H, T, d_h)
wrong = x.reshape(B, H, T, d_h)                         # та же форма, но...
for h in range(H):
    print(f"голова {h}: верно {right[0, h].tolist()}   неверно {wrong[0, h].tolist()}")

back = right.transpose(0, 2, 1, 3).reshape(B, T, D)     # склеиваем головы обратно
print("обратно совпало:", np.array_equal(back, x))

[[ 0  1  2  3  4  5]
 [ 6  7  8  9 10 11]]
голова 0: верно [[0, 1], [6, 7]]   неверно [[0, 1], [2, 3]]
голова 1: верно [[2, 3], [8, 9]]   неверно [[4, 5], [6, 7]]
голова 2: верно [[4, 5], [10, 11]]   неверно [[8, 9], [10, 11]]
обратно совпало: True

В правильном варианте голова 0 видит первые два признака каждого токена. В неправильном её «токен 1» — это кусок того же токена 0. Форма верная, код не падает, модель учится, но хуже, а загруженный чужой чекпоинт выдаёт мусор. Если же ошибиться только в одной из двух функций (разрезать неверно, а склеить верно, или наоборот), в причинной модели будущее начнёт просачиваться в прошлое. Первый тест задания ниже ловит именно этот баг, а тест на причинность — его вариант с одной ошибкой.

В PyTorch: view после transpose

transpose в PyTorch не копирует данные, а меняет шаги (strides), и тензор становится «несмежным». Вызов y.transpose(1, 2).view(B, T, D) упадёт с RuntimeError: view size is not compatible with input tensor's size and stride... Use .reshape(...) instead. Пишите .reshape(B, T, D) или .contiguous().view(B, T, D).

Того же результата можно добиться через einsumeinsum (нотация Эйнштейна)EN Einstein summation (einsum)Запись тензорных операций строкой индексов, например 'bhqd,bhkd->bhqk'., например оценки всех голов: np.einsum('bhtd,bhsd->bhts', q, k). Это удобно читать, но разрезать на головы всё равно придётся.

Выходная проекция W_O

После склейки вектор токена — это $[\mathrm{head}_1 \mid \mathrm{head}_2 \mid \dots \mid \mathrm{head}_H]$: результаты голов лежат в отдельных «ячейках» и ещё никак не взаимодействовали. $W_O$ их смешивает и возвращает вектор размера $d_{\text{model}}$, который прибавляется к остаточному потоку (урок 3.4.2).

Есть полезный второй взгляд. Разрежем $W_O$ по строкам на $H$ блоков $W_O^{(h)}$ размера $d_h \times d_{\text{model}}$. Тогда по правилу умножения блочных матриц:

$$\mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_H)\, W_O = \sum_{h=1}^{H} \mathrm{head}_h\, W_O^{(h)}$$

Что здесь написано: выход MHA — сумма независимых вкладов голов. Каждая голова сама решает, что прочитать (через $Q$, $K$), что взять ($V$) и что записать в остаточный поток ($W_O^{(h)}$). Проверим численно:

import numpy as np

rng = np.random.default_rng(0)
T, D, H = 5, 8, 4
d_h = D // H
heads = rng.normal(size=(H, T, d_h))        # выходы 4 голов (как будто уже посчитаны)
W_o = rng.normal(size=(D, D))

concat = heads.transpose(1, 0, 2).reshape(T, D)          # (T, H·d_h): головы рядом
y1 = concat @ W_o                                        # склеить и умножить на W_O
y2 = sum(heads[h] @ W_o[h * d_h:(h + 1) * d_h] for h in range(H))  # сумма вкладов голов
print("Concat·W_O == сумма вкладов голов:", np.allclose(y1, y2))

Concat·W_O == сумма вкладов голов: True

Этот взгляд понадобится дважды. В интерпретируемости (разделе науки о том, как модель устроена изнутри) вклад каждой головы изучают отдельно (см. индукционные головы ниже). А в MLA от DeepSeek произведение $W_V^{(h)} W_O^{(h)}$ каждой головы перемножают заранее: так называемое поглощение матриц.

Что выучивают головы

Никто не говорит головам, на что смотреть, но после обучения они специализируются:

• BERT (Clark и соавт., 2019): есть головы, которые смотрят на следующий или предыдущий токен, на разделители и точки. Отдельные головы с высокой точностью находят прямое дополнение глагола, определители существительных, объекты предлогов и кореферентные упоминания («она» → «кошка»).
• Перевод с английского на русский (Елена Войта и соавт., 2019): важные головы делятся на позиционные (смотрят на соседа), синтаксические (конкретные грамматические связи) и головы редких слов. Остальные почти не нужны: если удалить 38 из 48 голов энкодера, качество падает всего на 0.15 BLEU.
• «Are Sixteen Heads Really Better than One?» (Michel и соавт., 2019): многие головы можно отключить уже после обучения без заметной потери качества, а некоторые слои можно сократить до одной головы.

Самый известный пример — индукционные головыиндукционная головаEN induction headПара голов, которая находит в контексте прошлое появление текущего токена и копирует токен, шедший после него ([A][B] … [A] → [B]). Механизм in-context learning (Olsson и др., 2022). (Olsson и соавт., Anthropic, 2022). Это пара голов в разных слоях. Первая записывает в каждый токен информацию о его предшественнике. Вторая ищет в контексте прошлое появление текущего токена и копирует то, что шло после него: $[A][B] \dots [A] \to [B]$. Если в тексте было «Пушкин родился в Москве», а потом снова встретилось «Пушкин», индукционная голова подскажет «родился». Это один из механизмов обучения в контекстеобучение в контекстеEN in-context learning (ICL)Модель решает новую задачу по примерам в промпте, без изменения весов. Ярко показано в статье о GPT-3 (2020).. Одной головой такой трюк не сделать: вторая голова ищет по информации, которую первая записала слоем ниже. Хороший пример того, зачем трансформеру и много голов, и много слоёв. Как найти такую пару голов по паттернам внимания и доказать её роль выключением, разберём в уроке 9.3.1.

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

Цена внимания: O(n²)

FLOPs

Посчитаем операции одного слоя внимания на последовательность длины $T$ при ширине $D$. Операции считают во FLOPsFLOPs, число операций с плавающей точкойEN floating-point operations (FLOPs)Не путайте FLOPs (сколько операций нужно сделать) и FLOP/s, FLOPS (сколько операций железо делает в секунду). — числе арифметических действий над дробными числами; 1 TFLOP (терафлоп) = $10^{12}$ таких действий. Умножение матриц $(m \times k)(k \times n)$ стоит $2mkn$ FLOPs: в результате $m \cdot n$ чисел, каждое — скалярное произведение длины $k$, то есть $k$ умножений и примерно $k$ сложений. Пример: $(2 \times 3)(3 \times 4)$ — это $2 \cdot 2 \cdot 3 \cdot 4 = 48$ FLOPs.

• проекции $W_Q, W_K, W_V, W_O$: четыре умножения $(T \times D)(D \times D)$ — это $8TD^2$;
• оценки $QK^\top$ во всех головах: $2T^2D$;
• взвешивание $AV$: ещё $2T^2D$.

$$\text{FLOPs} \approx 8TD^2 + 4T^2D$$

Первое слагаемое растёт линейно по длине, второе — квадратично. Они равны при $T = 2D$. Возьмём слой шириной $D = 4096$, как у Llama 3 8B. Считаем его как обычное MHA; у настоящей Llama из-за GQA (о ней ниже) проекции $K$ и $V$ вчетверо дешевле, и квадратичная часть начинает доминировать ещё раньше.

Длина $T$	Проекции, TFLOP	Оценки и $AV$, TFLOP	Доля квадратичной части
1024	0.14	0.02	11%
8192	1.10	1.10	50%
32 768	4.40	17.6	80%
131 072	17.6	281	94%

Это квадратичная сложностьквадратичная сложность вниманияEN quadratic complexity (O(n²))Каждый из n токенов сравнивается с каждым — n² оценок на голову. При длинном контексте это главная стоимость внимания; отсюда FlashAttention, разреженное и линейное внимание.: на коротких текстах внимание дешёвое, на длинных съедает всё. (Причинная маска позволяет пропустить почти половину оценок, а полносвязные слои мы здесь не считали, но общая картина та же.)

Память

Наивная реализация хранит матрицу весов $(B, H, T, T)$ целиком. Для Llama 3 8B ($H = 32$, BF16BF16 (bfloat16)EN bfloat16 (brain floating point)16-битный формат: 1 бит знака, 8 бит порядка (как у FP32), 7 бит мантиссы. Диапазон как у FP32, точность — 2–3 значащие цифры. Стандарт обучения LLM. — формат дробного числа в 2 байта, подробно в уроке 5.4.1). Единицы: 1 КиБ = $2^{10}$ байт, 1 ГиБ = $2^{30}$ байт ≈ 1.07 млрд байт, 1 ТиБ = $2^{40}$ байт («двоичные» кило-, гига- и терабайты):

$$\begin{gathered} 32 \cdot 8192^2 \cdot 2 \text{ байт} = 4 \text{ ГиБ на слой}, \\ 32 \cdot 131\,072^2 \cdot 2 \text{ байт} = 1 \text{ ТиБ на слой} \end{gathered}$$

Терабайт на один слой одного запроса — так считать нельзя. Поэтому все современные системы используют FlashAttention (Tri Dao и соавт., 2022, arXiv:2205.14135). Он считает внимание блоками: берёт кусок запросов и кусок ключей, обновляет «бегущие» максимум и сумму экспонент для softmax и никогда не держит матрицу $T \times T$ в видеопамяти. Результат совпадает с формулой точно (до округления), память — $O(T)$, а работает быстрее, потому что гораздо реже читает видеопамять. FLOPs при этом остаются $O(T^2)$. F.scaled_dot_product_attention в PyTorch использует такие ядра автоматически. Как именно он устроен, вы разберёте и напишете сами в уроке 5.7.2.

Чтобы уйти от $O(T^2)$ в вычислениях, нужны другие идеи: разреженное внимание (NSA и DSA от DeepSeek) или линейное (Lightning Attention в MiniMax-01, Gated DeltaNet в Qwen3-Next, KDA в Kimi Linear). Это модуль 5.5. А при генерации, где каждый новый токен смотрит на все прошлые, спасает KV-кэш (урок 4.6.4).

Мост к MQA, GQA и MLA

При генерации модель хранит ключи и значения всех прошлых токенов, чтобы не пересчитывать их заново для каждого нового токена. Это KV-кэшKV-кэшEN KV cache. В MHA у каждой головы свои $K$ и $V$, поэтому на токен на слой нужно $2 \cdot H \cdot d_h$ чисел. Для Llama 3 8B в MHA-варианте это $2 \cdot 32 \cdot 128 = 8192$ числа на слой, а на все 32 слоя в BF16 — $8192 \cdot 32 \cdot 2$ байт = 512 КиБ на токен. На контексте 32K это 16 ГиБ на один запрос, примерно столько же, сколько весят все веса модели.

Головы во многом дублируют друг друга, так что можно поделиться ключами и значениями:

flowchart LR
A["MHA: у каждой головы свои K и V"] --> B["MQA: одна пара K/V на все головы (кэш ÷ H)"]
A --> C["GQA: одна пара K/V на группу голов (Llama 3: 32 → 8)"]
A --> D["MLA: K и V всех голов сжаты в один латентный вектор (DeepSeek-V2)"]

Три способа сократить KV-кэш многоголового внимания

• MQAмногозапросное внимание (MQA)EN multi-query attention (MQA)Все головы запросов делят одну голову ключей и одну голову значений. KV-кэш меньше в H раз, но качество заметно падает (Shazeer, 2019). (Shazeer, 2019): все головы запросов используют одну голову ключей и значений. Кэш меньше в $H$ раз, но качество может заметно просесть.
• GQAвнимание с группами запросов (GQA)EN grouped-query attention (GQA)Головы запросов делятся на G групп, у каждой группы своя пара K/V. G = H — это MHA, G = 1 — MQA. Стандарт Llama 3 и Qwen (Ainslie и др., 2023). (Ainslie и соавт., 2023): головы делятся на $G$ групп, у группы — своя пара K/V. Llama 3 8B: 32 головы запросов, 8 KV-голов, кэш 128 КиБ на токен вместо 512.
• MLAмногоголовое латентное внимание (MLA)EN multi-head latent attention (MLA)Предложено DeepSeek в DeepSeek-V2 (2024). (DeepSeek-V2, 2024): ключи и значения всех голов восстанавливаются из одного латентного вектора, и кэшируется только он: 512 + 64 числа на токен на слой у DeepSeek-V3. Для MHA той же ширины (128 голов по 128) было бы $2 \cdot 128 \cdot 128 = 32\,768$, почти в 57 раз больше.

Подробно, с кодом и честными таблицами качества — в модуле 5.2: MQA и GQA и MLA. Внутри всех трёх — то самое многоголовое внимание, которое вы сейчас напишете.

Практика

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

Реализуйте split_heads, merge_heads, multi_head_attention и score_matrix_bytes. Функция attention из прошлого урока дана готовой. Тесты проверят раскладку голов, обратимость склейки, совпадение с эталоном-циклом по головам, причинность и размер матрицы внимания для Llama 3.

Многоголовое внимание на NumPynumpy · нужен интернет8 тестов

import numpy as np


# ── Готовые функции из прошлого урока (их писать не нужно) ──────────────────────
def softmax(x):
    x = x - x.max(axis=-1, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)


def attention(q, k, v, causal=True):
    """Scaled dot-product attention для массивов (..., T, d). Возвращает (out, weights)."""
    T, d = q.shape[-2], q.shape[-1]
    scores = q @ np.swapaxes(k, -1, -2) / np.sqrt(d)          # (..., T, T)
    if causal:
        scores = np.where(np.tril(np.ones((T, T), dtype=bool)), scores, -np.inf)
    w = softmax(scores)
    return w @ v, w


# ── Ваша часть ─────────────────────────────────────────────────────────────────
def split_heads(x, n_heads):
    """(B, T, D) -> (B, H, T, D // H).

    Голова h получает столбцы h*d_h : (h+1)*d_h каждого токена, где d_h = D // H.
    Сначала reshape в (B, T, H, d_h), потом поменять местами оси T и H.
    """
    # TODO
    ...


def merge_heads(x):
    """(B, H, T, d_h) -> (B, T, H * d_h) — обратная операция к split_heads."""
    # TODO: сначала вернуть оси в порядок (B, T, H, d_h), потом reshape
    ...


def multi_head_attention(X, W_q, W_k, W_v, W_o, n_heads, causal=True):
    """Многоголовое self-attention.

    X: (B, T, D); W_q, W_k, W_v, W_o: (D, D); D делится на n_heads.
    Возвращает (out, weights): out — (B, T, D), weights — (B, H, T, T).
    Шаги: проекции -> split_heads -> attention во всех головах сразу -> merge_heads -> @ W_o.
    """
    # TODO
    ...


def score_matrix_bytes(batch, n_heads, seq_len, bytes_per_elem=2):
    """Сколько байт занимает матрица весов внимания (B, H, T, T) одного слоя, если её хранить целиком."""
    # TODO
    ...


# Проверьте себя кнопкой «Запустить»:
x = np.arange(2 * 3 * 4).reshape(1, 2, 12)          # B=1, T=2, D=12
s = split_heads(x, 3)
print(None if s is None else s.shape)                # ожидается (1, 3, 2, 4)
print(None if s is None else s[0, 1])                # голова 1: [[4 5 6 7] [16 17 18 19]]
rng = np.random.default_rng(0)
X = rng.normal(size=(2, 5, 8))
W = [rng.normal(size=(8, 8)) / np.sqrt(8) for _ in range(4)]
res = multi_head_attention(X, *W, n_heads=2)
if res is not None:
    print(res[0].shape, res[1].shape)                # ожидается (2, 5, 8) (2, 2, 5, 5)
b = score_matrix_bytes(1, 32, 8192)
print(b if b is None else b / 2**30, "ГиБ")          # Llama 3 8B, 8K токенов: 4.0 ГиБ на слой

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

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

Модуль MHA в стиле nanoGPT и проверка на встроенном nn.MultiheadAttention: копируем в него наши веса и сравниваем выходы.

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

torch.manual_seed(0)


class MultiHeadAttention(nn.Module):
    """Причинное многоголовое self-attention в стиле GPT-2 / nanoGPT."""

    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0, "d_model должен делиться на n_heads"
        self.h, self.dh = n_heads, d_model // n_heads
        # одна большая матрица D×D на все головы сразу; bias=False — без смещения, как в формулах урока
        self.wq = nn.Linear(d_model, d_model, bias=False)
        self.wk = nn.Linear(d_model, d_model, bias=False)
        self.wv = nn.Linear(d_model, d_model, bias=False)
        self.wo = nn.Linear(d_model, d_model, bias=False)    # выходная проекция W_O

    def forward(self, x):
        B, T, D = x.shape
        # (B, T, D) -> (B, T, H, d_h) -> (B, H, T, d_h)
        q = self.wq(x).view(B, T, self.h, self.dh).transpose(1, 2)
        k = self.wk(x).view(B, T, self.h, self.dh).transpose(1, 2)
        v = self.wv(x).view(B, T, self.h, self.dh).transpose(1, 2)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)  # (B, H, T, d_h)
        y = y.transpose(1, 2).reshape(B, T, D)                       # склеиваем головы
        return self.wo(y)


d_model, n_heads = 64, 8
ours = MultiHeadAttention(d_model, n_heads)
print("параметров:", sum(p.numel() for p in ours.parameters()), "= 4·d² =", 4 * d_model**2)

# Сверяемся со встроенным nn.MultiheadAttention: копируем в него те же веса
ref = nn.MultiheadAttention(d_model, n_heads, bias=False, batch_first=True)   # batch_first: вход (B, T, D)
with torch.no_grad():                        # копирование весов — не часть обучения, градиенты не нужны
    # встроенный модуль хранит W_Q, W_K, W_V одной матрицей, сложенной «стопкой» по строкам
    ref.in_proj_weight.copy_(torch.cat([ours.wq.weight, ours.wk.weight, ours.wv.weight]))
    ref.out_proj.weight.copy_(ours.wo.weight)

x = torch.randn(2, 10, d_model)
future = torch.ones(10, 10, dtype=torch.bool).triu(1)        # True — запрещено (так принято в nn.MultiheadAttention)
y_ref, _ = ref(x, x, x, attn_mask=future, need_weights=False)
print("совпадает с nn.MultiheadAttention:", torch.allclose(ours(x), y_ref, atol=1e-5))

Ожидаемый вывод:

параметров: 16384 = 4·d² = 16384
совпадает с nn.MultiheadAttention: True

Обратите внимание на маску: у nn.MultiheadAttention булево True значит «запрещено», поэтому здесь triu(1), а не tril(). Об этой разнице в соглашениях мы предупреждали в прошлом уроке.

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

Почувствуйте $O(T^2)$ на одной RTX 4080. Скрипт сравнивает наивную формулу и F.scaled_dot_product_attention на слое размером с Llama 3 8B (32 головы по 128) и растущей длине:

"""Память внимания: наивная формула против F.scaled_dot_product_attention.

Запуск: python attn_memory.py               (длины по умолчанию)
        python attn_memory.py 4096 8192     (свои длины)
"""
import math
import sys
import time

import torch
import torch.nn.functional as F

assert torch.cuda.is_available(), "нужна CUDA-сборка PyTorch и видеокарта"
dev = "cuda:0"
B, H, d_h = 1, 32, 128                       # один слой внимания размером с Llama 3 8B


def naive(q, k, v):
    T = q.size(-2)
    s = q @ k.transpose(-2, -1) / math.sqrt(d_h)            # (B, H, T, T) — та самая матрица T×T
    allowed = torch.ones(T, T, dtype=torch.bool, device=q.device).tril()
    return s.masked_fill(~allowed, float("-inf")).softmax(-1) @ v


def sdpa(q, k, v):
    return F.scaled_dot_product_attention(q, k, v, is_causal=True)  # FlashAttention-ядро


lengths = [int(a) for a in sys.argv[1:]] or [2048, 4096, 8192, 16384, 32768]
for T in lengths:
    q, k, v = (torch.randn(B, H, T, d_h, device=dev, dtype=torch.bfloat16) for _ in range(3))
    for name, fn in [("naive", naive), ("sdpa", sdpa)]:
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats(dev)
        base = torch.cuda.memory_allocated(dev)
        try:
            fn(q, k, v)                                      # прогрев
            # GPU работает асинхронно: Python ставит задачу в очередь и бежит дальше.
            # synchronize ждёт, пока карта всё досчитает, — иначе замер времени будет враньём
            torch.cuda.synchronize(dev)
            t0 = time.perf_counter()
            y = fn(q, k, v)
            torch.cuda.synchronize(dev)
            ms = (time.perf_counter() - t0) * 1000
            peak = (torch.cuda.max_memory_allocated(dev) - base) / 2**30
            print(f"T={T:6d}  {name:5s}  пик памяти {peak:6.2f} ГиБ  время {ms:8.1f} мс")
            del y
        except torch.cuda.OutOfMemoryError:
            print(f"T={T:6d}  {name:5s}  не хватило памяти (OOM)")

Чего ожидать (оценка по формулам, точные цифры зависят от версии PyTorch и драйвера). У наивной версии пик памяти растёт в 4 раза при каждом удвоении $T$: при $T = 8192$ одновременно живут несколько матриц по 4 ГиБ (оценки, маскированная копия, веса), а при $T = 16\,384$ каждая такая матрица весит 16 ГиБ, и 32 ГБ карты не хватит. У SDPA пик — сотни мегабайт даже на 32K: хранятся только $Q, K, V$ и выход. Время SDPA тоже растёт примерно квадратично. FlashAttention экономит память, но не FLOPs. Весь прогон должен занять от нескольких секунд до минуты (тоже оценка).

Хотите задействовать все четыре карты? Запустите четыре копии скрипта, каждую на своей карте и со своими длинами. Переменная CUDA_VISIBLE_DEVICES оставляет процессу одну видеокарту, и внутри скрипта она видна как cuda:0:

CUDA_VISIBLE_DEVICES=0 python attn_memory.py 2048 4096 &
CUDA_VISIBLE_DEVICES=1 python attn_memory.py 8192 &
CUDA_VISIBLE_DEVICES=2 python attn_memory.py 16384 &
CUDA_VISIBLE_DEVICES=3 python attn_memory.py 32768 65536 &
wait

Чем это лучше CPU-версии: на процессоре те же длины считались бы в десятки раз дольше, а главное — на карте с жёсткими 32 ГБ вы своими глазами видите OOM, ради которого придумали FlashAttention. Как настроить карты — на странице «Ваша лаборатория (4×GPU)».

Китайский угол: сколько голов нужно — ответ DeepSeek и Kimi

DeepSeek-V3: 128 голов. Для MHA это было бы очень дорого по KV-кэшу, но в MLA кэш не зависит от числа голов: кэшируется один латентный вектор на токен. Поэтому DeepSeek может позволить себе много голов, причём вместе они шире самой модели ($128 \cdot 128 = 16\,384$ против $d_{\text{model}} = 7168$).

Kimi K2: вдвое меньше. Moonshot построила K2 по образцу DeepSeek-V3 (MLA и смесь экспертов MoE из модуля 5.3), но сократила число голов до 64 (arXiv:2507.20534). Аргумент из техотчёта K2: при одинаковом числе токенов удвоение голов давало лишь небольшое улучшение валидационного loss (0.5–1.2%). А на контексте 128K переход от 64 к 128 головам увеличивает FLOPs инференса на 83%. Сверьте с нашей формулой: при длинном контексте доминирует слагаемое $4T^2 \cdot H d_h$, а оно пропорционально числу голов. (Рост 83%, а не 100%, потому что слои экспертов и часть проекций от числа голов не зависят.) Для агентной модели, которая живёт на длинных контекстах, это плохой размен. Вывод для инженера: число голов — не «чем больше, тем лучше», а компромисс между качеством и стоимостью длинного контекста. Подробнее о K2 — в треке «Китайская линза».

Что почитать и посмотреть. Су Цзяньлинь в статье «Предельное перетягивание каната между кэшем и качеством: от MHA, MQA, GQA к MLA» проводит одной линией весь путь от многоголового внимания этого урока до MLA. Лучше прочитать её перед модулем 5.2. А вторая лекция Ли Хунъи о self-attention (YouTube) объясняет многоголовость на картинках. Китайские термины здесь дословные: «многоголовое внимание», «голова внимания».

Проверь себя

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

1 Сколько параметров в слое многоголового внимания без смещений при $d_{\text{model}} = 768$ и 12 головах?

2 Как правильно превратить x формы $(B, T, D)$ в головы формы $(B, H, T, d_h)$?

x.reshape(B, H, T, d_h)x.reshape(B, T, H, d_h).transpose(0, 2, 1, 3)x.transpose(0, 2, 1).reshape(B, H, T, d_h)x.reshape(B, T * H, d_h)

3 Сколько ГиБ занимает матрица весов внимания одного слоя при $H = 32$, $T = 8192$, BF16 (2 байта), батч 1?

4 FLOPs слоя внимания ≈ $8TD^2 + 4T^2D$. При какой длине $T$ квадратичная часть сравняется с проекциями, если $D = 4096$?

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

MHA с $H$ головами имеет столько же параметров, сколько одна голова ширины $d_{\text{model}}$Каждая голова работает только со своим куском $d_h$ спроецированных признаков$W_O$ не смешивает головы, он только меняет размерностьМногие головы обученной модели можно удалить с небольшой потерей качестваБольше голов всегда строго лучше

6 Как называется пара голов, которая продолжает повторяющийся в контексте паттерн $[A][B] \dots [A] \to [B]$? (англ. или рус. термин)

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

Итоги

• Одна голова — один softmax, один «вопрос» на токен. $H$ голов смотрят на разные отношения одновременно, каждая в своём подпространстве размера $d_h = d_{\text{model}} / H$.
• Параметров столько же, сколько у одной большой головы: $4d_{\text{model}}^2$. Головы не покупают, а делят ширину.
• Реализация — перекладывание осей: reshape(B, T, H, d_h) → transpose → внимание → обратно → $W_O$. Выход MHA — сумма вкладов голов.
• Стоимость слоя $\approx 8TD^2 + 4T^2D$: после $T = 2D$ доминирует квадратичная часть. FlashAttention убирает $O(T^2)$ по памяти, но не по FLOPs.
• Свои K и V у каждой головы делают KV-кэш дорогим. Отсюда MQA, GQA и MLA (модуль 5.2), а число голов — инженерный компромисс (Kimi K2: 64 вместо 128).

Дальше: внимание — только половина блока трансформера. Добавим полносвязный слой, остаточные связи и нормализацию — Блок трансформера. Домашняя работа по вниманию (s4-hw02-attention) ждёт вас в конце модуля 4.6, когда из этих деталей соберётся целая GPT.