Генерация и KV-кэш

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

Модель обучена, теперь ей надо говорить. GPT генерирует авторегрессионноавторегрессионная модель (авторегрессия)EN autoregressive modelГенерирует последовательность по одному элементу, подавая каждый сгенерированный токен обратно на вход. Так работают GPT, Qwen, DeepSeek.: предсказала токен, приписала его к тексту, снова предсказала. Метод generate из прошлых уроков на каждом шаге прогоняет через модель всю последовательность заново. Чтобы написать тысячный токен ответа, модель снова обрабатывает предыдущие 999, хотя для 998 из них уже всё посчитала на прошлом шаге.

Решение — KV-кэшKV-кэшEN KV cache: хранить ключи и значения прошлых токенов и на каждом шаге считать только новый токен. Без него не работает ни один чат-бот: ни ChatGPT, ни DeepSeek, ни Qwen. Но у кэша есть цена — память. На длинном контексте кэш одного запроса весит столько же, сколько вся модель, а сервер обслуживает сотни запросов. С этой проблемы начинается целая ветка инноваций: MQA, GQA и MLA от DeepSeek (модуль 5.2).

Что пересчитывается зря

Вспомним причинное внимание (урок 4.4.2). Токен на позиции $i$ смотрит только на позиции $\le i$. Значит, его векторы на всех слоях зависят только от него и от прошлого. Когда мы приписываем к тексту новый токен, ничего из уже посчитанного не меняется.

Возьмём строку «Мороз и солнце». Наивная генерация делает так:

Шаг	Что подаём в модель	Что нужно на самом деле
1	«Мороз и солнце» (14 токенов)	всё — это промпт
2	«Мороз и солнце;» (15)	только «;»
3	«Мороз и солнце; » (16)	только пробел
100	113 токенов	только последний

Что именно нужно новому токену от прошлых? В формуле внимания $\mathrm{softmax}(qK^\top/\sqrt{d_h})\,V$ запрос $q$ — только у нового токена, а от прошлых нужны ключи $K$ и значения $V$ на каждом слое. Их и кэшируем. Запросы прошлых токенов не нужны никогда, а их выходы уже использованы. Поэтому кэш называется KV, а не QKV.

Сколько стоит генерация с кэшем и без

Посчитаем FLOPs для модели с $L$ слоями ширины $d$ (параметров в блоках $12Ld^2$, каждый параметр — 2 FLOPs на токен: умножение и сложение; внимание по $t$ позициям — $4td$ на слой: $2td$ на $qK^\top$ и $2td$ на взвешенную сумму $V$). LM-голову и эмбеддинги для простоты не считаем.

Без кэша, шаг $t$: прогоняем все $t$ токенов.

$$\text{FLOPs}_t \approx \underbrace{24 L d^2 \cdot t}_{\text{проекции и MLP}} + \underbrace{4 L d \cdot t^2}_{\text{внимание}}$$

С кэшем, шаг $t$: только новый токен, но внимание по всем $t$ ключам.

$$\text{FLOPs}_t \approx 24 L d^2 + 4 L d \cdot t$$

Что здесь написано: без кэша стоимость одного нового токена растёт квадратично по длине (из-за внимания) и линейно из-за проекций. С кэшем — линейно, а проекции и MLP вообще не зависят от длины. Сложите по шагам $t = 1 \dots N$ — получите $O(N^3)$ против $O(N^2)$ для всего ответа.

Для GPT-2 small ($L = 12$, $d = 768$; строка про 4096 токенов условная — настоящий контекст GPT-2 всего 1024):

Сгенерировать $N$ токенов	Без кэша	С кэшем	Во сколько раз
128	1.43 TFLOP	0.022 TFLOP	65
1024	102 TFLOP	0.19 TFLOP	530
4096	2270 TFLOP	1.0 TFLOP	2259

Тысяча токенов без кэша — сто триллионов операций. Видеокарта, которая в идеальных условиях выполняет десятки–сотни триллионов операций в секунду, потратила бы на одну арифметику около секунды, а на 4096 токенов — от десятков секунд до минуты. С кэшем та же тысяча токенов — 0.2 TFLOP, доли секунды.

Две фазы генерации: prefill и decode

С кэшем генерация распадается на две фазы, которые по-разному нагружают железо (подробно — в уроке 7.1.1):

flowchart LR
P["Промпт: 14 токенов"] --> A["Prefill: все 14 позиций одним проходом<br/>с причинной маской"]
A --> C[("KV-кэш: K и V<br/>14 позиций × L слоёв")]
A --> T1["Токен 15"]
T1 --> D["Decode: 1 токен за проход<br/>q нового токена × все K из кэша"]
C --> D
D -->|"дописать k, v в кэш"| C
D --> T2["Токен 16, 17, …"]

Prefill заполняет кэш за один проход, decode добавляет по одной строке на токен

• prefillprefill (предзаполнение)EN prefillПервая фаза генерации: весь промпт проходит через модель одним параллельным проходом, заполняется KV-кэш и получается первый токен. Упирается в вычисления. — промпт обрабатывается целиком, как при обучении: матрицы, причинная маска, все позиции параллельно. На каждый прочитанный вес приходится много арифметики (тем больше, чем длиннее промпт), и на длинных промптах GPU упирается в вычисления (ограничено вычислениямиограничено вычислениями (compute-bound)EN compute-boundОперация, в которой вычислитель загружен полностью, а память успевает подвозить данные. Ускоряется только более быстрым железом или меньшим числом операций.).
• Decodeдекодирование (фаза decode)EN decode phaseВторая фаза: один токен за один проход модели. На каждый токен заново читаются все (активные) веса и KV-кэш, поэтому фаза упирается в пропускную способность памяти. — по одному токену. Арифметики мало, а прочитать нужно все веса модели и весь кэш. Эта фаза упирается в скорость памяти (ограничено памятьюограничено памятью (memory-bound)EN memory-boundОперация, которая дольше ждёт данные из памяти, чем считает. Ускоряется, если читать меньше байт (квантование, GQA, MLA) или делать больше работы на каждый прочитанный байт (батчинг).), а не в FLOPs.

Такая генерация называется инкрементальным декодированиеминкрементальное декодированиеEN incremental decodingГенерация, при которой на каждом шаге модель обрабатывает только новый токен, а ключи и значения прошлых берёт из KV-кэша..

KV-кэш для одной головы на NumPy

Начнём с сути, без PyTorch. Функция получает один новый токен, дописывает его ключ и значение в кэш и считает внимание по всему кэшу. Маска не нужна: в кэше только прошлое и сам токен. Ровно так работала микро-GPT из урока 4.6.1, только там кэш был списком списков.

import numpy as np


def decode_step(x_t, W_q, W_k, W_v, cache):
    """Один новый токен x_t (вектор длины d): дописать его k, v в кэш и посчитать внимание по всему кэшу."""
    q, k, v = x_t @ W_q, x_t @ W_k, x_t @ W_v
    if "K" in cache:                                   # дописываем в конец по оси времени
        cache["K"] = np.vstack([cache["K"], k])        # vstack — приставить строку снизу: (t, 4) → (t+1, 4)
        cache["V"] = np.vstack([cache["V"], v])
    else:
        cache["K"], cache["V"] = k[None, :], v[None, :]   # первый токен: вектор → матрица из одной строки (1, 4)
    scores = cache["K"] @ q / np.sqrt(q.shape[0])      # (t,) — оценки нового запроса со всеми ключами
    w = np.exp(scores - scores.max())                  # устойчивый softmax (числитель)
    return (w / w.sum()) @ cache["V"]                  # веса × значения


def full_causal(X, W_q, W_k, W_v):                     # эталон из урока 4.4: все позиции сразу
    Q, K, V = X @ W_q, X @ W_k, X @ W_v
    s = Q @ K.T / np.sqrt(Q.shape[1])
    s = np.where(np.tril(np.ones((len(X), len(X)), dtype=bool)), s, -np.inf)
    w = np.exp(s - s.max(axis=1, keepdims=True))
    return (w / w.sum(axis=1, keepdims=True)) @ V


rng = np.random.default_rng(0)
X = rng.normal(size=(6, 8))                            # 6 токенов, d = 8
W_q, W_k, W_v = (rng.normal(size=(8, 4)) for _ in range(3))
cache = {}
step_by_step = np.stack([decode_step(x, W_q, W_k, W_v, cache) for x in X])
print("кэш K:", cache["K"].shape, " совпадает с полным пересчётом:",
      np.allclose(step_by_step, full_causal(X, W_q, W_k, W_v)))

кэш K: (6, 4)  совпадает с полным пересчётом: True

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

KV-кэш в нашей GPT на PyTorch

Добавим кэш к модели из model.py (урок 4.6.2), не меняя её: напишем отдельные функции, которые используют те же слои. Кэш — список словарей, по одному на блок; в словаре тензоры k и v формы $(B, H, t, d_h)$.

"""kv_generate.py — генерация с KV-кэшем для GPT из model.py и сравнение с наивной генерацией."""
import time

import torch
import torch.nn.functional as F

from model import GPT, GPTConfig


def attn_with_cache(attn, x, cache):
    """CausalSelfAttention из model.py, но с кэшем: cache — словарь слоя ({} в начале)."""
    B, T, C = x.shape
    H = attn.n_head
    q, k, v = attn.c_attn(x).split(C, dim=2)
    q, k, v = (t.view(B, T, H, C // H).transpose(1, 2) for t in (q, k, v))   # (B, H, T, d_h)
    if cache:                                       # дописываем новые K, V в конец по оси времени
        k = torch.cat([cache["k"], k], dim=2)
        v = torch.cat([cache["v"], v], dim=2)
    cache["k"], cache["v"] = k, v
    # prefill (кэш был пуст, T > 1): обычная причинная маска; decode (T == 1): маска не нужна
    y = F.scaled_dot_product_attention(q, k, v, is_causal=T > 1)
    return attn.c_proj(y.transpose(1, 2).reshape(B, T, C))


@torch.no_grad()
def forward_with_cache(model, idx, caches):
    past = caches[0]["k"].size(2) if caches[0] else 0           # сколько токенов уже в кэше
    T = idx.size(1)
    pos = torch.arange(past, past + T)                          # позиции продолжают счёт!
    x = model.wte(idx) + model.wpe(pos)
    for block, cache in zip(model.blocks, caches):
        x = x + attn_with_cache(block.attn, block.ln_1(x), cache)
        x = x + block.mlp(block.ln_2(x))
    return model.lm_head(model.ln_f(x))


@torch.no_grad()
def generate_cached(model, idx, max_new_tokens):
    """Жадная генерация с KV-кэшем (не длиннее block_size: дальше позиций нет)."""
    caches = [{} for _ in model.blocks]
    logits = forward_with_cache(model, idx, caches)             # prefill: весь промпт одним проходом
    for _ in range(max_new_tokens):
        next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        idx = torch.cat([idx, next_id], dim=1)
        logits = forward_with_cache(model, next_id, caches)     # decode: только новый токен
    return idx


@torch.no_grad()
def generate_naive(model, idx, max_new_tokens):
    """Жадная генерация без кэша: на каждом шаге вся последовательность заново."""
    for _ in range(max_new_tokens):
        logits, _ = model(idx)
        next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        idx = torch.cat([idx, next_id], dim=1)
    return idx


torch.manual_seed(0)
model = GPT(GPTConfig(vocab_size=137, block_size=1024, n_layer=4, n_head=4, n_embd=128)).eval()
prompt = torch.randint(137, (1, 16))
for n_new in (128, 512, 1000):
    t0 = time.perf_counter()
    a = generate_naive(model, prompt, n_new)
    t1 = time.perf_counter()
    b = generate_cached(model, prompt, n_new)
    t2 = time.perf_counter()
    print(f"{n_new:5d} новых токенов: без кэша {t1 - t0:6.2f} с, с кэшем {t2 - t1:5.2f} с, "
          f"ускорение ×{(t1 - t0) / (t2 - t1):4.1f}, токены совпали: {torch.equal(a, b)}")

Наш запуск на одном потоке процессора (добавьте в начало torch.set_num_threads(1), чтобы повторить; модель случайная: для замера скорости обучать её не нужно, а жадная генерация детерминирована):

  128 новых токенов: без кэша   1.45 с, с кэшем  0.46 с, ускорение × 3.1, токены совпали: True
  512 новых токенов: без кэша  18.80 с, с кэшем  4.00 с, ускорение × 4.7, токены совпали: True
 1000 новых токенов: без кэша  87.99 с, с кэшем  8.89 с, ускорение × 9.9, токены совпали: True

Токены совпали все до одного, а ускорение растёт с длиной, как и обещала формула. Но по FLOPs для нашей модели ($L = 4$, $d = 128$, промпт 16 токенов) выигрыш должен быть ×74, ×287 и ×576. Почему на деле всего ×3–10? Модель крошечная, и время шага с кэшем уходит не на арифметику, а на накладные расходы: вызовы Python, запуск десятков маленьких операций, чтение весов, torch.cat растущего кэша. Чем больше модель и длиннее текст, тем ближе реальный выигрыш к теоретическому.

Проверить кэш можно и на обученной модели из урока 4.6.2. Замените в kv_generate.py блок замера (всё, что идёт после функций) на такой:

ckpt = torch.load("gpt_ru.pt", weights_only=False)   # в чекпоинте лежит GPTConfig, поэтому weights_only=False
model = GPT(ckpt["cfg"])
model.load_state_dict(ckpt["model"])
model.eval()
stoi = {ch: i for i, ch in enumerate(ckpt["chars"])}
itos = dict(enumerate(ckpt["chars"]))

prompt = torch.tensor([[stoi[c] for c in "Модель "]])
n_new = model.cfg.block_size - prompt.size(1)        # промпт + новые токены ≤ block_size = 128
a = generate_naive(model, prompt, n_new)
b = generate_cached(model, prompt, n_new)
print("совпали:", torch.equal(a, b))
print("".join(itos[i] for i in b[0].tolist()))

совпали: True
Модель по словарь — по словарь — по словами — обучения в словаре — одна по словами в словаре — после нет.
После с обучении — пос

Без weights_only=False свежие версии PyTorch откажутся загружать файл: по умолчанию torch.load читает только тензоры, а у нас в чекпоинте ещё и объект GPTConfig. Заодно видно знакомое по уроку 4.1.3 поведение жадной генерации: модель зацикливается на «по словарь — по словарь». Кэш тут ни при чём, он лишь ускоряет вычисления.

Три ловушки, на которых спотыкаются все

1. Позиции должны продолжать счёт. Новый токен стоит на позиции past, а не 0. Если забыть сдвиг, модель решит, что каждый новый токен — первый в тексте. Код при этом не падает. У моделей с RoPE то же самое: угол поворота считается от абсолютной позиции, и без сдвига кэшированные ключи и новый запрос «повёрнуты» несогласованно.

2. is_causal=True при декодировании. Маска в F.scaled_dot_product_attention выравнивается по левому верхнему углу. При одном запросе и пяти ключах она разрешает смотреть только на первый ключ:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 1, 1, 8)                           # один новый токен
k, v = torch.randn(1, 1, 5, 8), torch.randn(1, 1, 5, 8)   # 5 токенов в кэше
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print("видит только первый токен:", torch.allclose(y, v[:, :, :1]))

видит только первый токен: True

Поэтому в attn_with_cache стоит is_causal=T > 1: при декодировании одного токена маска не нужна, а prefill мы делаем только с пустым кэшем. Если промпт приходится дописывать кусками в непустой кэш, маску строят вручную по правому нижнему углу, как в MiniMind (см. китайский угол ниже).

3. Склеивать по оси времени. У тензора $(B, H, T, d_h)$ время — ось 2. torch.cat(..., dim=1) молча склеит головы, и формы перестанут сходиться только на следующем слое — или, хуже, сойдутся.

Почему кэш упирается в block_size

У нашей GPT обучаемые позиционные эмбеддинги: позиций ровно block_size. Наивный generate обрезает контекст до последних block_size токенов и может генерировать бесконечно. С кэшем так нельзя: после обрезки у всех токенов сдвинутся позиции, и кэш станет недействительным, его пришлось бы пересчитать. Поэтому generate_cached генерирует не дальше block_size. Настоящие LLM решают это иначе: огромный контекст (128K–1M токенов) и специальные схемы вроде скользящего окна (внимание со скользящим окномвнимание со скользящим окном (SWA)EN sliding window attention (SWA)Токен видит только w последних токенов. Стоимость O(n·w) вместо O(n²), KV-кэш не больше w позиций на слой., модуль 5.5).

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

На каждом слое для каждого токена храним ключ и значение каждой головы:

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

• $2$ — ключи и значения;
• $L$ — число слоёв;
• $H_{kv}$ — число KV-голов (в классическом многоголовое вниманиемногоголовое вниманиеEN multi-head attention (MHA) равно числу голов $H$);
• $d_h$ — размер головы;
• $T$ — длина контекста, $B$ — число одновременных запросов (батч);
• байт на число: 2 в FP16/BF16, 1 в FP8.

Пример, GPT-2 small: $2 \cdot 12 \cdot 12 \cdot 64 \cdot 2 = 36\,864$ байт = 36 КиБ на токен. На полном контексте 1024 токена — 36 МиБ. Немного. Но посмотрите, что происходит с ростом модели:

Модель	$L$	$H_{kv} \times d_h$	Кэш на токен (BF16)	Пример
GPT-2 small (2019)	12	12 × 64	36 КиБ	1024 токена — 36 МиБ
GPT-3 175B (2020)	96	96 × 128	4.5 МиБ	2048 токенов — 9 ГиБ
Llama 2 7B (MHA)	32	32 × 128	512 КиБ	4096 токенов × 32 запроса — 64 ГиБ
Llama 3 8B (GQA)	32	8 × 128	128 КиБ	8192 токена — ровно 1 ГиБ
Qwen2.5-7B (GQA)	28	4 × 128	56 КиБ	32K токенов — 1.75 ГиБ
DeepSeek-V3 (MLA)	61	латент 512 + 64	≈ 69 КиБ	при 128 головах и 671B параметров

Строка про Llama 2 7B — главный урок. Веса модели в FP16 занимают около 13.5 ГБ. А сервер, который держит 32 разговора по 4096 токенов, тратит на кэш 64 ГиБ (≈ 69 ГБ) — примерно в пять раз больше, чем на саму модель. Поэтому при обслуживании LLM узкое место — не веса, а кэш: он определяет, сколько пользователей влезет на одну видеокарту и какой длины контекст им можно дать.

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

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

Как уменьшить кэш

Каждый множитель в формуле — рычаг. Индустрия потянула за все:

Рычаг	Идея	Где	Подробно
$H_{kv}$	меньше KV-голов: одна на все головы (MQAмногозапросное внимание (MQA)EN multi-query attention (MQA)Все головы запросов делят одну голову ключей и одну голову значений. KV-кэш меньше в H раз, но качество заметно падает (Shazeer, 2019).) или одна на группу (GQAвнимание с группами запросов (GQA)EN grouped-query attention (GQA)Головы запросов делятся на G групп, у каждой группы своя пара K/V. G = H — это MHA, G = 1 — MQA. Стандарт Llama 3 и Qwen (Ainslie и др., 2023).)	Llama 3, Qwen2.5/3	урок 5.2.1
$H_{kv} \cdot d_h$	сжать K и V всех голов в один латентный вектор (MLAмногоголовое латентное внимание (MLA)EN multi-head latent attention (MLA)Предложено DeepSeek в DeepSeek-V2 (2024).)	DeepSeek-V2/V3, Kimi K2	урок 5.2.2
байт	хранить кэш в FP8 вместо BF16	vLLM, SGLang	модуль 7.2
$T$	смотреть не на все токены: скользящее окно хранит только последние $W$ токенов, разреженное внимание читает только выбранные	MiMo-V2-Flash (окно 128 в пяти слоях из шести), DeepSeek-V3.2	урок 5.5.2
фрагментация	выделять память под кэш страницами, как ОС (PagedAttentionPagedAttention (страничный KV-кэш)EN PagedAttentionKV-кэш хранится блоками фиксированного размера (например, по 16 токенов), а таблица блоков связывает логические позиции с физическими — как страницы виртуальной памяти в ОС. Основа vLLM.)	vLLM	урок 7.1.2

У DeepSeek-V3 128 голов по 128 измерений. В классическом MHA кэш на токен был бы $2 \cdot 61 \cdot 128 \cdot 128 \cdot 2$ байт ≈ 3.8 МиБ, а MLA хранит только латент на 512 чисел и общий RoPE-ключ на 64: $(512 + 64) \cdot 61 \cdot 2 = 70\,272$ байта, в 57 раз меньше. Как это сделано без потери качества — центральная тема модуля 5.2.

Практика

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

Реализуйте шаг декодирования с кэшем и три функции для расчёта памяти: размер кэша, байты на токен по config.json реальной модели и максимальный контекст, который влезет в свободную память. Тесты сверят пошаговую генерацию с полным причинным вниманием и посчитают кэш GPT-2, Llama 3 8B, Qwen2.5-7B, Qwen3-8B и Qwen3-0.6B.

KV-кэш: шаг декодирования и память реальных моделейnumpy · нужен интернет7 тестов

import numpy as np


# ── Часть 1. Шаг декодирования с KV-кэшем (одна голова) ─────────────────────────
def decode_step(x_t, W_q, W_k, W_v, cache):
    """Внимание для ОДНОГО нового токена с KV-кэшем.

    x_t — вектор нового токена формы (d,); W_q, W_k, W_v — матрицы (d, d_h).
    cache — словарь; пустой {} в начале, затем cache["K"] и cache["V"] формы (t, d_h).

    1) q, k, v = x_t @ W_q, x_t @ W_k, x_t @ W_v
    2) допишите k и v в конец cache["K"] и cache["V"] (np.vstack); для первого токена — k[None, :]
    3) оценки = cache["K"] @ q / sqrt(d_h); устойчивый softmax; верните веса @ cache["V"]  → форма (d_h,)
    Маска не нужна: в кэше только прошлые токены и текущий.
    """
    # TODO: ваш код
    raise NotImplementedError("decode_step")


# ── Часть 2. Сколько весит KV-кэш ───────────────────────────────────────────────
def kv_cache_bytes(n_layers, n_kv_heads, head_dim, seq_len, batch=1, bytes_per_elem=2):
    """Размер KV-кэша в байтах: 2 (K и V) · слои · KV-головы · размер головы · длина · батч · байт на число."""
    # TODO: ваш код
    raise NotImplementedError("kv_cache_bytes")


def kv_bytes_per_token(config, bytes_per_elem=2):
    """Байт KV-кэша на один токен по конфигу в стиле Hugging Face (config.json).

    Обязательные ключи: num_hidden_layers, num_attention_heads, hidden_size.
    Необязательные: num_key_value_heads (если нет — это MHA: столько же, сколько голов запросов),
                    head_dim (если нет — hidden_size // num_attention_heads).
    """
    # TODO: ваш код (используйте kv_cache_bytes с seq_len=1, batch=1)
    raise NotImplementedError("kv_bytes_per_token")


def max_context(free_bytes, config, batch=1, bytes_per_elem=2):
    """Сколько токенов на КАЖДЫЙ из batch запросов поместится в free_bytes байт свободной памяти.

    Целое число (округление вниз, //).
    """
    # TODO: ваш код
    raise NotImplementedError("max_context")


# ── Проверка: кэш против полного пересчёта и размеры кэша реальных моделей ─────────
try:
    rng = np.random.default_rng(0)
    d, d_h, T = 8, 4, 5
    X = rng.normal(size=(T, d))
    W_q, W_k, W_v = (rng.normal(size=(d, d_h)) for _ in range(3))
    cache = {}
    outs = np.stack([decode_step(X[t], W_q, W_k, W_v, cache) for t in range(T)])
    print("выход по шагам:", outs.shape, " кэш K:", cache["K"].shape)

    MODELS = {
        "GPT-2 small": {"num_hidden_layers": 12, "num_attention_heads": 12, "hidden_size": 768},
        "Llama 2 7B (MHA)": {"num_hidden_layers": 32, "num_attention_heads": 32, "hidden_size": 4096},
        "Llama 3 8B (GQA)": {"num_hidden_layers": 32, "num_attention_heads": 32, "num_key_value_heads": 8,
                             "hidden_size": 4096},
        "Qwen2.5-7B (GQA)": {"num_hidden_layers": 28, "num_attention_heads": 28, "num_key_value_heads": 4,
                             "hidden_size": 3584},
        "Qwen3-0.6B (GQA)": {"num_hidden_layers": 28, "num_attention_heads": 16, "num_key_value_heads": 8,
                             "hidden_size": 1024, "head_dim": 128},      # head_dim ≠ 1024 // 16!
    }
    for name, cfg in MODELS.items():
        per_tok = kv_bytes_per_token(cfg)
        print(f"{name:17s} {per_tok / 1024:4.0f} КиБ на токен, {per_tok * 1024 / 2**20:4.0f} МиБ на 1024 токена")
except NotImplementedError as e:
    print("Сначала допишите функцию:", e)

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

Обратите внимание на конфиги Qwen3 в тестах: в них есть явный head_dim. У Qwen3-8B он совпадает с $\text{hidden\_size} / H = 4096 / 32 = 128$, а у Qwen3-0.6B — нет: $1024 / 16 = 64$, но head_dim равен 128. Посчитаете по формуле «ширина / число голов» — ошибётесь с кэшем вдвое. Читая чужой config.json, всегда проверяйте, нет ли там head_dim и num_key_value_heads.

Заодно сравните ответы: крошечная Qwen3-0.6B хранит 112 КиБ на токен — вдвое больше, чем Qwen2.5-7B (56 КиБ). Размер кэша зависит не от числа параметров, а от слоёв, KV-голов и head_dim.

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

Проверьте KV-кэш на настоящей модели. Библиотека transformers умеет генерировать с кэшем и без него (флаг use_cache). Возьмите маленькую Qwen2.5-0.5B, чтобы быстро скачалась:

"""kv_hf.py — генерация с кэшем и без на Qwen2.5-0.5B.  pip install transformers accelerate"""
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.bfloat16).to("cuda:0")
inputs = tok("Расскажи сказку про кота учёного.", return_tensors="pt").to("cuda:0")
cfg = model.config
head_dim = getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads   # явный head_dim — в приоритете
per_token = 2 * cfg.num_hidden_layers * cfg.num_key_value_heads * head_dim * 2
print(f"KV-кэш: {per_token / 1024:.0f} КиБ на токен")

for use_cache in (True, False):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    # **inputs распаковывает словарь (input_ids, attention_mask) в именованные аргументы;
    # min_new_tokens=512 не даёт остановиться раньше, чтобы замер был честным
    out = model.generate(**inputs, max_new_tokens=512, min_new_tokens=512, do_sample=False, use_cache=use_cache)
    torch.cuda.synchronize()
    print(f"use_cache={use_cache}: {512 / (time.perf_counter() - t0):.1f} токенов/с")

Чего ожидать: кэш у этой модели — 12 КиБ на токен (24 слоя, 2 KV-головы по 64). С кэшем генерация в разы быстрее, и разрыв растёт с max_new_tokens: попробуйте 128, 512 и 2048. Затем подставьте конфиг Qwen2.5-7B в функцию max_context из задания и прикиньте: сколько разговоров по 32K токенов поместится на одну карту с 32 ГБ после загрузки весов в BF16 (около 15 ГБ)? А на четыре карты, если веса разрезать между ними?

Что даёт GPU по сравнению с NumPy-версией: настоящая модель на 0.5B параметров, генерации на тысячи токенов без долгого ожидания и честный замер в токенах в секунду. Если карта ещё не настроена — «Ваша лаборатория (4×GPU)».

Китайский угол: KV-кэш как экзамен и как главный ресурс

В учебных курсах. В задаче 2 курса LLM-Beginner (Цю Сипэн, Фуданьский университет) KV-кэш — обязательный пункт, и проверяет его автотест kv_cache_equivalence: логиты с кэшем и без должны расходиться меньше чем на $10^{-4}$. Это ровно наш тест из NumPy-раздела. В README задачи перечислены типичные ошибки: позиция нового токена не продолжает историю (у них RoPE, поэтому углы поворота не совпадают с полным прогоном) и склейка кэша не по той оси. Авторы отдельно отмечают, что в разборе nanoGPT из их практикума KV-кэш только объясняют, но не реализуют, — поэтому здесь студенты пишут его сами.

В MiniMind (model_minimind.py) кэш склеивается по оси последовательности, а позиции RoPE начинаются с длины кэша (start_pos). Быстрый путь через scaled_dot_product_attention с is_causal используется только без кэша. С кэшем маска $-\infty$ накладывается вручную на последние столбцы матрицы оценок — правое нижнее выравнивание, о котором шла речь в ловушке 2.

В индустрии. DeepSeek-V2 начинается с проблемы кэша: по сравнению с DeepSeek 67B кэш стал меньше на 93.3%, а максимальная пропускная способность генерации выросла в 5.76 раза (arXiv:2405.04434). Судя по пересчёту, главный вклад — MLA (по числу хранимых чисел около −82%), а остальное, по-видимому, дало хранение кэша в среднем в 6 битах. Разбор этой цифры — в уроке 5.2.2. А сервис Kimi компании Moonshot построен на архитектуре Mooncake, где кэш — центральный ресурс: prefill и decode работают на разных машинах, а KV-кэш хранится и передаётся между ними как отдельный слой системы (arXiv:2407.00079). Подробно — в уроке 7.1.2.

Китайские термины: «KV-кэш», «инкрементальное декодирование», «предварительное заполнение» — prefill.

Проверь себя

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

1 Почему K и V прошлых токенов можно хранить и не пересчитывать?

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

2 Сколько КиБ занимает KV-кэш GPT-2 small (12 слоёв, 12 голов по 64) на один токен в FP16?

3 Llama 2 7B (MHA: 32 слоя, 32 головы по 128), BF16. Сколько ГиБ кэша нужно на 32 запроса по 4096 токенов?

4 Что верно про генерацию с кэшем?Можно выбрать несколько ответов

Prefill обрабатывает весь промпт за один проходНа шаге decode в модель подаётся только новый токенСтоимость внимания на шаге decode растёт линейно с длиной контекстаКэш хранит запросы Q всех прошлых токеновDecode обычно упирается в скорость памяти, а не в FLOPs

5 Вы декодируете один токен с кэшем на 100 позиций и вызываете F.scaled_dot_product_attention(q, k, v, is_causal=True). Что произойдёт?

Всё правильно: маска закроет будущие позицииОшибка формы: q и k разной длиныНовый токен будет смотреть только на первый токен кэша, генерация испортится без ошибокВнимание станет равномерным

6 Во сколько раз KV-кэш DeepSeek-V2 меньше, чем у DeepSeek 67B, если он сокращён на 93.3%? Округлите до целого.

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

Итоги

• Из-за причинной маски ключи и значения прошлых токенов не меняются. KV-кэш хранит их, и новый токен стоит $O(t \cdot d)$ на внимание вместо полного пересчёта.
• Генерация делится на prefill (весь промпт одним проходом) и decode (по токену, упирается в память).
• Правильный кэш даёт точно те же токены, что и наивная генерация. Ловушки: позиции продолжают счёт, is_causal при декодировании, склейка по оси времени.
• Кэш весит $2 \cdot L \cdot H_{kv} \cdot d_h \cdot T \cdot B \cdot$ байт. На длинном контексте и большом батче он больше самой модели.
• MQA, GQA, MLA, FP8-кэш, скользящее окно и PagedAttention — рычаги этой формулы.

Домашняя работа

Модуль 4.6 завершают две домашние работы.

• s4-hw02-attention — внимание на PyTorch: ядро softmax(QKᵀ/√d)V с причинной маской (сверка с SDPA), MultiHeadAttention, GQA, инкрементальное декодирование с KV-кэшем и расчёт его памяти; на уровне Boss — RoPE.
• s4-hw03-gpt-pytorch — GPT уже в современном виде, как Llama, Qwen и MiniMind (урок 4.5.3): RMSNorm, RoPE, SwiGLU, GQA, weight tying, сэмплирование с температурой, top-k и top-p; на уровне Boss — KV-кэш. Сигнатуры в ДЗ свои, но model.py и kv_generate.py из этого модуля — хорошая шпаргалка: структура модели и проверка «с кэшем = без кэша» те же.

Проверка из корня репозитория:

python homework/check.py s4-hw02-attention
python homework/check.py s4-hw03-gpt-pytorch

Дальше: перед ступенью 5 соберём в одну карту инструменты, которыми пользуются практики «На вашей машине»: Hugging Face, accelerate и torchrun, серверы инференса, терминальные помощники и типичные ошибки — Инструменты LLM-инженера.