Multi-head Latent Attention (DeepSeek-V2)

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

В прошлом уроке мы увидели дилемму. GQA экономит память, но теряет качество: в экспериментах DeepSeek модель с GQA-8 отстала от MHA на 4 пункта MMLU. MQA экономит ещё сильнее и теряет ещё больше.

В мае 2024 года DeepSeek выпустила DeepSeek-V2 (236B параметров, из них 21B активных) с новым механизмом — Multi-head Latent Attention (MLA)многоголовое латентное внимание (MLA)EN multi-head latent attention (MLA)Предложено DeepSeek в DeepSeek-V2 (2024).. Его KV-кэш занимает 576 чисел на токен на слой. Это столько же, сколько у GQA с 2.25 группами, то есть меньше, чем у любой GQA-модели из таблицы прошлого урока. А по качеству MLA в экспериментах DeepSeek оказался лучше MHA. По сравнению с DeepSeek 67B модель V2 сократила KV-кэш на 93.3% и подняла максимальную пропускную способность генерации в 5.76 раза.

Дешёвый кэш означает дешёвый API: MLA вместе со смесью экспертов (модуль 5.3) позволила DeepSeek резко снизить цены, а вслед за ней подешевели и остальные китайские провайдеры. Эту историю мы разбираем в уроке 8.1.1, а здесь разберём саму MLA до последней матрицы. Сегодня её используют DeepSeek-V3, R1, Kimi K2 и другие модели.

Идея: кэшировать «сырьё», а не готовые K и V

Начнём с наблюдения. В DeepSeek-V2 вход слоя внимания $h_t$ имеет размер 5120. Из него MHA делает ключи и значения 128 голов по 128 чисел, то есть $2 \cdot 128 \cdot 128 = 32\,768$ чисел. Кэш в 6.4 раза больше, чем сам вход, из которого он получен линейными проекциями. Значит, в K и V много избыточности.

Можно было бы хранить сам $h_t$ и пересчитывать K и V на лету. MLA идёт дальше: сначала сжимает $h_t$ в маленький латентный вектор, а ключи и значения всех голов восстанавливает из него. Это низкоранговое сжатиенизкоранговое сжатиеEN low-rank compressionЗамена большой матрицы произведением двух узких (через «бутылочное горлышко» малой размерности). В MLA так сжимают ключи и значения в латентный вектор c_KV.:

$$c^{KV}_t = W^{DKV} h_t, \qquad k^C_t = W^{UK} c^{KV}_t, \qquad v^C_t = W^{UV} c^{KV}_t$$

• $W^{DKV}$ (down-projection) сжимает $d = 5120$ в $d_c = 512$.
• $W^{UK}$ и $W^{UV}$ (up-projections) разворачивают 512 чисел в ключи и значения всех голов: $128 \cdot 128 = 16\,384$ каждое.
• В кэш идёт только $c^{KV}_t$: 512 чисел вместо 32 768.

По сути, большая матрица ключей $W^K = W^{UK} W^{DKV}$ размера $16\,384 \times 5120$ заменена произведением двух узких с «горлышком» на 512. Похожий приём вы встретите в LoRA (урок 6.2.1), где так устроена добавка к весам. Только здесь узкое место нужно не ради экономии параметров, а ради того, чтобы промежуточный результат было дёшево хранить.

Запросы в DeepSeek-V2 тоже сжимаются: $c^Q_t = W^{DQ} h_t$ размера $d_c' = 1536$, затем $q^C_t = W^{UQ} c^Q_t$. На кэш это не влияет (запросы не кэшируются). По словам авторов, сжатие запросов экономит память на активациях при обучении.

Величина	Обозначение	DeepSeek-V2	DeepSeek-V3
Размер скрытого состояния	$d$	5120	7168
Слоёв	$L$	60	61
Голов	$n_h$	128	128
Размер головы	$d_h$	128	128
Латент KV	$d_c$	512	512
Латент запроса	$d_c'$	1536	1536
RoPE-часть головы	$d_h^R$	64	64

Но тут есть загвоздка. Если на каждом шаге декодирования восстанавливать K и V всего префикса из латентов, мы выполним огромное умножение $W^{UK} c^{KV}_j$ для каждого прошлого токена $j$. Память сэкономили, а вычисления выросли. Спасает алгебра.

Фокус: поглощение матриц

Посмотрим на оценку внимания (score, урок 4.4.2) головы $i$ между запросом $t$ и ключом $j$. Обозначим через $W^{UK}_i$ кусок матрицы $W^{UK}$, который восстанавливает ключ именно головы $i$: это матрица $128 \times 512$, так что $k^C_{j,i} = W^{UK}_i c^{KV}_j$. Подставим определения:

$$q_{t,i}^\top k^C_{j,i} = q_{t,i}^\top \left( W^{UK}_i \, c^{KV}_j \right) = \underbrace{\left( (W^{UK}_i)^\top q_{t,i} \right)^\top}_{\tilde q_{t,i} \,\in\, \mathbb{R}^{d_c}} c^{KV}_j$$

Умножение матриц ассоциативно — $(AB)C = A(BC)$ (урок 1.1.2), — поэтому скобки можно переставить. Посчитаем размеры: $q_{t,i}$ — 128 чисел, $W^{UK}_i$ — матрица $128 \times 512$, $c^{KV}_j$ — 512 чисел. Слева мы сначала делаем из 512 чисел ключ из 128 и потом скалярно умножаем его на запрос. Справа сначала из запроса делаем вектор $\tilde q$ из 512 чисел, а потом скалярно умножаем его на латент. Число получается одно и то же. Мини-проверка на числах: $q = (1, 2)$, $W = \begin{pmatrix} 1 & 0 & 2 \\ 0 & 1 & 1 \end{pmatrix}$, $c = (1, 1, 1)$. Слева $Wc = (3, 2)$ и $q \cdot (3, 2) = 7$; справа $W^\top q = (1, 2, 4)$ и $(1, 2, 4) \cdot c = 7$. Вместо того чтобы разворачивать каждый закэшированный латент в ключ, мы один раз переводим запрос в латентное пространство ($\tilde q_{t,i}$, 512 чисел) и умножаем его прямо на закэшированные $c^{KV}_j$. Матрица $W^{UK}$ «поглощена» запросом. Её можно даже заранее перемножить с $W^{UQ}$.

То же со значениями. Выход головы $i$ — это $o_{t,i} = \sum_j p_{t,j,i} \, W^{UV}_i c^{KV}_j = W^{UV}_i \left( \sum_j p_{t,j,i} \, c^{KV}_j \right)$. Сначала взвешиваем сами латенты, потом один раз разворачиваем результат. А поскольку дальше идёт выходная проекция $W^O$, произведение $W^O_i W^{UV}_i$ можно предвычислить. Это и есть поглощение матрицпоглощение матрицEN matrix absorption (weight absorption)Перестановка скобок в произведении матриц на инференсе MLA — W_UK переносится в запрос, W_UV в выходную проекцию, и внимание считается прямо по латентному кэшу без восстановления K и V..

MLA — это MHA при обучении и MQA при декодировании

Посмотрите, что получилось после поглощения. Все 128 голов скалярно умножают свои запросы на один и тот же вектор $c^{KV}_j$ и взвешивают его же как значение. Это MQA с одной общей «головой» размера 512!

Режим	Как считаем	На что похоже	Чем ограничен
Обучение и prefill	восстанавливаем K и V каждой головы (192 и 128 чисел)	MHA	вычислениями
Декодирование	внимание прямо по латентам, $W^{UK}$ и $W^{UV}$ поглощены	MQA с головой 512 + 64	памятью

Это наблюдение Су Цзяньлиня из статьи «Чем хорош MLA?». MLA позволяет выбирать форму вычислений под фазу. При обучении работает выразительная MHA с отдельными K и V у каждой головы. При декодировании работает экономная MQA, где каждое прочитанное из памяти число используется всеми 128 головами.

За поглощение есть цена. Оценки и взвешенная сумма теперь считаются в пространстве размера 512 (+64 на RoPE), а не 128 (+64), поэтому FLOPs внимания на один закэшированный токен примерно в 3–4 раза больше. При декодировании это не страшно: вспомните прошлый урок, там вычислитель всё равно простаивает в ожидании памяти. А там, где упираемся в вычисления (обучение, длинный prefill), выгоднее наивная форма. Какую форму где включать — инженерное решение. В простом официальном инференс-коде DeepSeek-V3 (inference/model.py) оба режима есть в одном классе. В режиме naive в кэш кладутся уже восстановленные K и V всех голов, как в MHA, а в режиме absorb — только латент и RoPE-ключ. По умолчанию включён attn_impl = "absorb".

Проблема: RoPE мешает поглощению

Всё было бы идеально, если бы не позиции. DeepSeek-V2 использует RoPEротационные позиционные эмбеддинги (RoPE)EN rotary position embedding (RoPE)RoPE предложен в Китае (Су Цзяньлинь и соавторы, компания Zhuiyi). Сейчас используется почти во всех LLM. из урока 4.5.2. RoPE поворачивает и запрос, и ключ на углы, зависящие от их позиций:

$$\mathrm{score}_{t,j} = \left( R_t \, q_{t,i} \right)^\top \left( R_j \, W^{UK}_i c^{KV}_j \right) = q_{t,i}^\top \, R_t^\top R_j \, W^{UK}_i \, c^{KV}_j$$

Матрица поворота $R_j$ зависит от позиции ключа и стоит между запросом и $W^{UK}_i$. Матрицы не коммутируют (в общем случае $AB \ne BA$, урок 1.1.2), поэтому $W^{UK}_i$ уже не «протащить» в запрос: для каждой позиции $j$ получалась бы своя матрица. Пришлось бы на каждом шаге восстанавливать и поворачивать ключи всего префикса, и выигрыш от MLA пропал бы. Авторы DeepSeek-V2 формулируют это прямо: RoPE несовместим с низкоранговым сжатием KV.

Развязанный RoPE

Решение DeepSeek — развязанный RoPEразвязанный RoPEEN decoupled RoPEПриём MLA — позиционная информация идёт через отдельную небольшую часть запроса и ключа (64 измерения в DeepSeek-V2/V3), а сжатая «контентная» часть остаётся без RoPE, чтобы работало поглощение матриц.. Голова получает две части:

• Контентная часть (128 измерений): $q^C_{t,i}$ и $k^C_{j,i}$ из латентов, без RoPE. Её можно поглощать.
• Позиционная часть (64 измерения): $q^R_{t,i} = \mathrm{RoPE}(W^{QR} c^Q_t)$ у каждой головы своя, а ключ $k^R_j = \mathrm{RoPE}(W^{KR} h_j)$ один на все головы, как в MQA. Его считают сразу из $h_j$ и кэшируют.

Запрос и ключ склеиваются: $q_{t,i} = [q^C_{t,i}; q^R_{t,i}]$, $k_{j,i} = [k^C_{j,i}; k^R_j]$ (квадратные скобки с точкой с запятой — склейка векторов: из 128 и 64 чисел получается 192). Скалярное произведение склеенных векторов — сумма скалярных произведений их частей, поэтому оценка распадается на сумму двух скалярных произведений, и масштаб берётся по полной размерности головы:

$$\mathrm{score}_{t,j,i} = \frac{ (q^C_{t,i})^\top k^C_{j,i} + (q^R_{t,i})^\top k^R_j }{ \sqrt{d_h + d_h^R} } = \frac{\ldots}{\sqrt{192}}$$

Первое слагаемое поглощается и считается по латентам, второе работает с маленьким общим ключом. В кэш на токен на слой идут $c^{KV}$ и $k^R$, всего 576 чисел ($512 + 64$).

flowchart LR
H["h_t (5120)"] --> DKV["W_DKV"] --> C["c_KV (512)"]
H --> KR["W_KR + RoPE"] --> KRv["k_R (64), общий для голов"]
C --> CACHE[("KV-кэш: 576 чисел")]
KRv --> CACHE
H --> DQ["W_DQ"] --> CQ["c_Q (1536)"]
CQ --> UQ["W_UQ: q_C, 128 голов × 128"]
CQ --> QR["W_QR + RoPE: q_R, 128 голов × 64"]
CACHE --> ATT["Внимание: q_C·k_C + q_R·k_R"]
UQ --> ATT
QR --> ATT
ATT --> O["W_UV, W_O → выход"]

Поток данных MLA в DeepSeek-V2/V3. В кэш попадают только латент и общий RoPE-ключ.

Сравнение MHA, GQA, MQA и MLA

Таблица 1 из статьи DeepSeek-V2, пересчитанная для их размеров ($n_h = 128$, $d_h = 128$), на один токен одного слоя:

Механизм	Формула кэша на слой	Чисел	Качество (по оценке авторов)
MHA	$2 n_h d_h$	32 768	сильное
GQA, 8 групп	$2 n_g d_h$	2 048	среднее
MQA	$2 d_h$	256	слабое
MLA	$d_c + d_h^R \approx \tfrac{9}{2} d_h$	576	сильнее MHA

Для целой модели DeepSeek-V3 (61 слой, BF16) это $61 \cdot 576 \cdot 2 = 70\,272$ байт, около 68.6 КиБ на токен. Для сравнения: у Qwen3-235B-A22B — 188 КиБ, у Llama 3.1 405B — 504 КиБ. Если бы V3 использовала MHA со 128 головами по 128, понадобилось бы 3.8 МиБ на токен, в 57 раз больше. На контексте 128K кэш V3 занимает около 8.6 ГиБ вместо почти 490 ГиБ.

Калькулятор 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) · слои чисел.

Откуда «−93.3%»: читаем статью внимательно

Цифра сравнивает DeepSeek-V2 с DeepSeek 67B, а не с MHA. У 67B было 95 слоёв и GQA с 8 KV-головами: $2 \cdot 95 \cdot 8 \cdot 128 = 194\,560$ чисел на токен. У V2: $60 \cdot 576 = 34\,560$. По числу элементов сокращение около 82%. В разделе о развёртывании авторы пишут, что дополнительно квантуют KV-кэш V2 в среднем до 6 бит. Наш пересчёт: $34\,560 \cdot 0.75$ байта против $194\,560 \cdot 2$ байт даёт как раз −93.3%. Прямо в статье такого расчёта нет, но похоже, что цифра из аннотации включает и архитектуру, и квантование. Это нормально, просто полезно понимать, что с чем сравнивают.

Сравнение с MHA при одинаковой остальной архитектуре авторы привели в приложении D.2. Там MoE-модели обучили с MHA и с MLA. У малых моделей (около 16B параметров) кэш MLA составил 15.6 тыс. чисел на токен против 110.6 тыс. у MHA, то есть 14%. У больших (около 250B) — 34.6 тыс. против 860.2 тыс., то есть 4%. По MMLU, BBH и CMMLU MLA оказалась лучше в обоих размерах. Единственное исключение — C-Eval у малой модели (50.9 против 51.6).

Почему MLA не хуже MHA?

Строгого доказательства нет, только эксперименты и интуиции.

• В отличие от GQA, головы не делят ключи буквально: у каждой головы своя матрица восстановления $W^{UK}_i$, поэтому ключи разные. Ограничение мягче: ключи всех 128 голов для одного токена — линейные функции одного и того же 512-мерного вектора $c^{KV}$.
• Размерность головы для оценок внимания больше: $128 + 64 = 192$ при обучении.
• Су Цзяньлинь в статьях «Чем хорош MLA?» поставил серию экспериментов. Среди факторов, которые объясняют выигрыш, он называет увеличенную размерность головы, частичный RoPE (позиции только в части измерений) и общие K/V.

Авторы работы TransMLA приводят ещё один аргумент: при одинаковом размере кэша MLA выразительнее GQA. Любую GQA-модель можно переписать как MLA с тем же кэшем, а обратное, по их анализу, неверно.

Практика

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

Реализуйте MLA без RoPE-части в двух формах: «как при обучении» (K и V восстанавливаются из латента) и «как на инференсе» (поглощённые матрицы, внимание прямо по кэшу). Тест проверит, что выходы совпадают до $10^{-9}$.

MLA: латентный кэш и поглощение матрицnumpy · нужен интернет6 тестов

import numpy as np


def causal_softmax(scores):
    """Готовый помощник: причинная маска + устойчивый softmax по последней оси. scores: (..., T, T)."""
    T = scores.shape[-1]
    future = np.triu(np.ones((T, T), dtype=bool), k=1)
    s = np.where(future, -np.inf, scores)
    s = s - s.max(axis=-1, keepdims=True)
    w = np.exp(s)
    return w / w.sum(axis=-1, keepdims=True)


def mla_naive(h, W_dkv, W_uk, W_uv, W_q, W_o, n_heads):
    """MLA «как при обучении» (без RoPE-части): восстанавливаем K и V из латента.

    h: (T, D) — входы слоя;  W_dkv: (D, d_c) — сжатие в латент c = h @ W_dkv;
    W_uk, W_uv: (d_c, H*d_h) — восстановление ключей и значений всех голов;
    W_q: (D, H*d_h);  W_o: (H*d_h, D).
    Возвращает (out, c): out — (T, D), c — (T, d_c), то есть ровно то, что кладётся в KV-кэш.
    Масштаб скоров — 1/sqrt(d_h), маска причинная.
    """
    T, D = h.shape
    d_h = W_q.shape[1] // n_heads
    # TODO: c = ...; K = ...; V = ...; Q = ...
    #       разбейте на головы: (T, H*d_h) -> (H, T, d_h)
    #       внимание по головам, склейте головы обратно (T, H*d_h) и умножьте на W_o
    ...


def absorb(W_q, W_uk, W_uv, W_o, n_heads):
    """«Поглощает» матрицы восстановления.

    Возвращает (W_qk, W_vo):
      W_qk[i] = W_q_i @ W_uk_i.T  — форма (H, D, d_c): скор головы i = (h W_qk[i]) · c
      W_vo[i] = W_uv_i @ W_o_i    — форма (H, d_c, D): вклад головы i = (веса @ c) @ W_vo[i]
    где W_q_i, W_uk_i, W_uv_i — столбцы головы i (срез i*d_h:(i+1)*d_h), W_o_i — строки головы i.
    """
    d_h = W_q.shape[1] // n_heads
    # TODO
    ...


def mla_absorbed(h, c, W_qk, W_vo, d_h):
    """MLA «как на инференсе»: внимание прямо по кэшу c, без восстановления K и V.

    h: (T, D) — входы (из них считаются запросы); c: (T, d_c) — латентный кэш;
    W_qk: (H, D, d_c); W_vo: (H, d_c, D); d_h — исходный размер головы (для масштаба 1/sqrt(d_h)).
    Возвращает out: (T, D).
    """
    # TODO: для каждой головы: scores = (h @ W_qk[i]) @ c.T / sqrt(d_h);
    #       веса = causal_softmax(scores); out += (веса @ c) @ W_vo[i]
    ...


# Проверьте себя кнопкой «Запустить»:
rng = np.random.default_rng(0)
T, D, H, d_h, d_c = 5, 32, 4, 8, 6
h = rng.normal(size=(T, D))
W_dkv = rng.normal(size=(D, d_c)) / np.sqrt(D)
W_uk, W_uv = rng.normal(size=(2, d_c, H * d_h)) / np.sqrt(d_c)
W_q = rng.normal(size=(D, H * d_h)) / np.sqrt(D)
W_o = rng.normal(size=(H * d_h, D)) / np.sqrt(H * d_h)
res = mla_naive(h, W_dkv, W_uk, W_uv, W_q, W_o, H)
ab = absorb(W_q, W_uk, W_uv, W_o, H)
if res is not None and ab is not None:
    out, c = res
    out2 = mla_absorbed(h, c, *ab, d_h)
    print("кэш на токен:", c.shape[1], "чисел вместо", 2 * H * d_h)
    print("совпадает:", out2 is not None and np.allclose(out, out2))

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

Полная MLA на PyTorch: с развязанным RoPE

Модуль MLA в миниатюре: сжатие запросов, латент $c^{KV}$, общий RoPE-ключ, наивная и поглощённая формы. В конце — численное доказательство, что «повесить» RoPE на сжатые ключи нельзя.

import torch
import torch.nn as nn

torch.manual_seed(0)
torch.set_default_dtype(torch.float64)  # float64 — чтобы сравнение было точным


def rope(x, pos, base=10000.0):
    """RoPE: вращаем пары координат (0,1), (2,3), … на угол pos·θ_i. x: (..., T, d)."""
    d = x.shape[-1]
    theta = base ** (-torch.arange(0, d, 2) / d)
    ang = pos[:, None] * theta[None, :]                      # (T, d/2)
    cos, sin = ang.cos(), ang.sin()
    x1, x2 = x[..., 0::2], x[..., 1::2]
    out = torch.empty_like(x)
    out[..., 0::2] = x1 * cos - x2 * sin
    out[..., 1::2] = x1 * sin + x2 * cos
    return out


class MLA(nn.Module):
    """Multi-head Latent Attention в миниатюре (имена как в DeepSeek-V2)."""

    def __init__(self, d=64, n_h=4, d_h=16, d_c=8, d_cq=24, d_r=8):
        super().__init__()
        self.n_h, self.d_h, self.d_c, self.d_r = n_h, d_h, d_c, d_r
        self.W_DQ = nn.Linear(d, d_cq, bias=False)           # сжатие запроса
        self.W_UQ = nn.Linear(d_cq, n_h * d_h, bias=False)   # «контентная» часть запроса
        self.W_QR = nn.Linear(d_cq, n_h * d_r, bias=False)   # RoPE-часть запроса (своя у каждой головы)
        self.W_DKV = nn.Linear(d, d_c, bias=False)           # сжатие K и V в латент c_KV  → в кэш
        self.W_KR = nn.Linear(d, d_r, bias=False)            # общий RoPE-ключ k_R        → в кэш
        self.W_UK = nn.Linear(d_c, n_h * d_h, bias=False)    # восстановление ключей
        self.W_UV = nn.Linear(d_c, n_h * d_h, bias=False)    # восстановление значений
        self.W_O = nn.Linear(n_h * d_h, d, bias=False)

    def heads(self, x, dim):                                  # (T, H*dim) -> (H, T, dim)
        return x.view(x.shape[0], self.n_h, dim).transpose(0, 1)

    def queries(self, h, pos):
        c_q = self.W_DQ(h)
        return self.heads(self.W_UQ(c_q), self.d_h), rope(self.heads(self.W_QR(c_q), self.d_r), pos)

    def cache(self, h, pos):
        """Всё, что MLA хранит на токен: c_KV (d_c чисел) и k_R (d_r чисел)."""
        return self.W_DKV(h), rope(self.W_KR(h), pos)

    def attend(self, scores, T):
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool), 1)
        return (scores / (self.d_h + self.d_r) ** 0.5).masked_fill(mask, float("-inf")).softmax(-1)

    def forward_naive(self, h):
        """Как при обучении: восстанавливаем K и V каждой головы из c_KV."""
        T = h.shape[0]
        pos = torch.arange(T, dtype=torch.float64)
        q_c, q_r = self.queries(h, pos)
        c_kv, k_r = self.cache(h, pos)
        k_c, v = self.heads(self.W_UK(c_kv), self.d_h), self.heads(self.W_UV(c_kv), self.d_h)
        scores = q_c @ k_c.transpose(1, 2) + q_r @ k_r.T     # k_R общий для всех голов
        o = self.attend(scores, T) @ v                       # (H, T, d_h)
        return self.W_O(o.transpose(0, 1).reshape(T, -1))

    def forward_absorbed(self, h):
        """Как на инференсе: внимание прямо по кэшу, K и V не восстанавливаются."""
        T = h.shape[0]
        pos = torch.arange(T, dtype=torch.float64)
        q_c, q_r = self.queries(h, pos)
        c_kv, k_r = self.cache(h, pos)
        W_UK = self.W_UK.weight.view(self.n_h, self.d_h, self.d_c)   # (H, d_h, d_c)
        W_UV = self.W_UV.weight.view(self.n_h, self.d_h, self.d_c)
        q_lat = q_c @ W_UK                                    # (H, T, d_c): W_UK «ушла» в запрос
        scores = q_lat @ c_kv.T + q_r @ k_r.T
        o_lat = self.attend(scores, T) @ c_kv                 # (H, T, d_c): значения = сам латент
        o = o_lat @ W_UV.transpose(1, 2)                      # W_UV можно заранее влить в W_O
        return self.W_O(o.transpose(0, 1).reshape(T, -1))


m = MLA()
h = torch.randn(10, 64)
a, b = m.forward_naive(h), m.forward_absorbed(h)
print("макс. расхождение naive vs absorbed:", f"{(a - b).abs().max().item():.1e}")
print("кэш на токен: MLA", m.d_c + m.d_r, "чисел, MHA", 2 * m.n_h * m.d_h, "чисел")

# Почему RoPE нельзя повесить на k_C: пробуем «поглотить» W_UK, когда ключ повёрнут
pos = torch.arange(10, dtype=torch.float64)
q_c, _ = m.queries(h, pos)
c_kv, _ = m.cache(h, pos)
k_c = m.heads(m.W_UK(c_kv), m.d_h)
true = rope(q_c, pos) @ rope(k_c, pos).transpose(1, 2)              # честный RoPE на обеих сторонах
W_UK = m.W_UK.weight.view(m.n_h, m.d_h, m.d_c)
fake = (rope(q_c, pos) @ W_UK) @ c_kv.T                              # «поглощение» без поворота ключа
print("ошибка «поглощения» при RoPE на k_C:", f"{(true - fake).abs().max().item():.2f}")

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

макс. расхождение naive vs absorbed: 1.1e-16
кэш на токен: MLA 16 чисел, MHA 128 чисел
ошибка «поглощения» при RoPE на k_C: 1.23

Первая строка — это машинный ноль: две формы MLA математически тождественны. Последняя строка показывает, что при RoPE на сжатых ключах поглощение даёт другие оценки внимания. Именно поэтому позиции вынесены в отдельные 64 измерения.

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

Поднимите этот модуль до настоящих размеров DeepSeek-V3: $d = 7168$, 128 голов, $d_h = 128$, $d_c = 512$, $d_h^R = 64$. Один такой слой — около 187 млн параметров, в BF16 это примерно 0.4 ГБ, одна карта держит его легко. Замерьте через torch.cuda.Event время шага декодирования при длине кэша 4K, 32K и 128K в двух формах. Наивная форма из нашего модуля на каждом шаге разворачивает весь префикс в K и V (при 128K это около 8.6 ГБ временных тензоров K и V в BF16 на один слой), поглощённая считает прямо по латенту размером 576 чисел на токен. Вы увидите своими глазами, почему в официальном коде DeepSeek по умолчанию стоит attn_impl = "absorb". Для сравнения посчитайте кэш всей модели на 128K токенов для MLA и для GQA с 8 KV-головами при тех же 61 слое: 8.6 ГиБ против 30.5 ГиБ. Это разница между «на одной 32-гигабайтной карте помещаются кэши трёх таких запросов» и «кэш одного запроса занимает почти всю карту, и для весов места уже не остаётся». Настройка карт — на странице «Ваша лаборатория (4×GPU)»; как честно мерить время на GPU (с синхронизацией перед замером) — в уроке 4.6.4.

Китайский угол: как MLA разошлась по индустрии

DeepSeek. MLA появилась в DeepSeek-V2 (май 2024) и перешла без изменений в DeepSeek-V3 и R1. В официальном конфиге V3: kv_lora_rank = 512, q_lora_rank = 1536, qk_rope_head_dim = 64, 128 голов. Название lora_rank напоминает, что в основе — низкоранговое разложение. Разреженное внимание DSA в V3.2 тоже работает поверх MLA: дешёвый индексатор выбирает top-k токенов, и полное MLA считается только по ним (урок 5.5.2).

Moonshot. Kimi K2 (1T параметров, 32B активных) взяла архитектуру DeepSeek-V3 вместе с MLA, но уменьшила число голов со 128 до 64, чтобы удешевить инференс. Kimi Linear сочетает MLA с линейным вниманием KDA.

Инфраструктура. В первый день Open Source Week (24 февраля 2025) DeepSeek открыла FlashMLA — ядро декодирования MLA для GPU архитектуры Hopper (H800/H100; Hopper и Blackwell — названия поколений видеокарт NVIDIA для дата-центров, NPU — ускорители нейросетей, у Huawei это Ascend). Репозиторий развивается дальше: с релиза 30.09.2026 он обслуживает DeepSeek-V4.1 на GPU NVIDIA поколения Blackwell и на NPU Huawei Ascend, а поддержка Hopper и прежних моделей (V3, V3.2, V4.0) осталась в зафиксированной старой версии (так сказано в README репозитория). По опубликованной DeepSeek суточной статистике сервиса V3/R1 (февраль 2025) 56.3% входных токенов попали в KV-кэш на дисках. Хранить кэш на дисках выгодно, когда он маленький, и MLA делает его маленьким.

Перевод готовых моделей. В 2025 году вышли две работы китайских исследователей, которые переводят готовые модели на MLA коротким дообучением: TransMLA (для GQA-моделей вроде Llama и Qwen) и MHA2MLA (для MHA-моделей). Это тот же приём «сжать и дообучить», что и uptraining в GQA из прошлого урока.

Что почитать. У Су Цзяньлиня — вторая половина статьи о кэше и цикл «Чем хорош MLA?» с экспериментами.

Проверь себя

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

1 Сколько чисел на токен на один слой хранит MLA в DeepSeek-V2 ($d_c = 512$, $d_h^R = 64$)?

2 Скольким группам GQA (при $d_h = 128$) соответствует кэш MLA в DeepSeek-V2? Ответ — число с точностью до сотых.

3 Почему RoPE нельзя применить к ключам, восстановленным из латента, если мы хотим поглощать $W^{UK}$?

RoPE не работает с размерностью 512RoPE увеличивает KV-кэш вдвоеМатрица поворота зависит от позиции ключа и встаёт между запросом и $W^{UK}$, поэтому общее произведение нельзя предвычислитьRoPE несовместим с softmax

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

После поглощения все головы скалярно умножают свои запросы на один и тот же латент, как в MQARoPE-ключ $k^R$ у каждой головы свойСжатие запросов ($d_c' = 1536$) уменьшает KV-кэшПри обучении и prefill выгодно восстанавливать K и V (форма MHA), при декодировании — работать с латентами

5 Сколько КиБ занимает KV-кэш DeepSeek-V3 на один токен в BF16 (61 слой, 576 чисел на слой)? Округлите до десятых.

6 DeepSeek-V2 сократила KV-кэш на 93.3% по сравнению с…

DeepSeek 67B (GQA, 95 слоёв); часть выигрыша, по нашему пересчёту, — квантование кэша до ~6 битMHA-версией самой себяLlama 3 70BGPT-4

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

Итоги

• MLA кэширует не K и V, а латентный вектор $c^{KV}$ (512 чисел) и маленький общий RoPE-ключ (64): 576 чисел на токен на слой.
• Поглощение матриц переносит $W^{UK}$ в запрос и $W^{UV}$ в выход. При декодировании MLA работает как MQA по латентам, при обучении — как MHA.
• RoPE ломает поглощение, поэтому позиции вынесены в отдельную развязанную часть головы.
• Кэш DeepSeek-V3 — около 69 КиБ на токен против 188–504 КиБ у GQA-моделей сопоставимого размера, при качестве не хуже MHA.

Дальше: вторая половина эффективности DeepSeek — смесь экспертов. Начнём с основ MoE.