Ли Хунъи, «Машинное обучение» 2026: FlashAttention, KV-кэш и позиционные эмбеддинги

О конспекте

Пересказ своими словами. Главный источник — разборы лекций в блоге vincentxuu/quidproquo (по слайдам inference.pdf, pos.pdf, субтитрам и учебному Colab; ссылки на статьи сверены с arXiv), а также конспекты по субтитрам из dimthink/awesome-ai-podcast-notes. Формулы записаны нами. Обзор курсов и заданий — здесь.

Коротко

• Спрашивайте о цене ускорения. Ли предлагает рамку на две лекции: меняет ли метод результат внимания (точный или приближённый)? Требует ли он обучать/переделывать модель? И что ещё за него платим?
• FlashAttention точен, подходит к любой модели и почти бесплатен: GPU считает быстро, а медленно перевозит данные между большой памятью (HBM) и крошечной быстрой (SRAM). Если разрешить себе «временно неверный» softmax и потом домножать на поправку, внимание считается за один проход без записи матрицы весов.
• KV-кэш избавляет от пересчёта, но съедает память: в примере с архитектурой Gemma 2 27B — около 0,72 МБ на токен, на A100 80 ГБ влезает лишь ~114 тыс. токенов. Отсюда GQA, MLA, скользящее окно, attention sink, обрезка кэша.
• Кэширование префикса — это деньги: кэшированные входные токены в API в несколько раз дешевле, поэтому всё неизменное в промпте агента ставят в начало, а изменяющееся (дату!) — в конец.
• Позиции: важны относительные расстояния; RoPE кладёт их во вращение $q$ и $k$. Контекст расширяют, «сжимая» углы низкочастотных пар (PI, NTK, YaRN, LongRoPE), а работы NoPE и DroPE показывают, что каузальная модель может обходиться без позиционного кодирования.

Как смотреть вместе с курсом

Лекция	Когда смотреть	Наш урок
FlashAttention	вместе с уроком 5.7.2 (перед ним — 5.7.1 про устройство GPU)	5.7.1, 5.7.2
KV-кэш	после 4.6.4, вместе с модулем 5.2 и уроком 7.1.2	4.6.4, 5.2.1, 5.2.2, 7.1.2
Позиции	после 4.5.2, вместе с 5.5.1	4.5.2, 5.5.1, 5.5.2

Курс предполагает, что вы уже понимаете устройство Transformer (лекция «Анатомия большой языковой модели» из курса осени 2025 года). Спекулятивное декодирование в этом году не читали — Ли отсылает к лекции 16 курса 2024 года (наш конспект).

---

Лекция «Ускоряем генерацию, часть 1: FlashAttention»

Видео на YouTube · слайды inference.pdf, с. 1–28, учебный Colab.

Префилл и декодирование. Подсказка проходит модель целиком за один раз (prefill), затем токены порождаются по одному (decode). Для токена 4 внимание — это $a_i = q_4 \cdot k_i$, softmax до $\hat a_i$, взвешенная сумма $v_i$.

Метафора Ли. Вычислительные блоки GPU — толпа «многоруких эльфов», считают мгновенно; верстак — встроенная SRAM, крошечная (обычно с десяток мегабайт); склад — HBM (на A100 — 80 ГБ). Всё, что считается, надо перенести со склада на верстак и обратно, и именно перевозка — узкое место. Главное ограничение: на верстак нельзя класть ничего, что растёт с длиной последовательности $L$ (в агентах $L$ — сотни тысяч), поэтому ключи режут на блоки по $N$.

Наивный softmax — много походов на склад. Чтобы вычесть максимум (ради численной устойчивости), надо сначала просмотреть все $a_i$; потом посчитать экспоненты; потом знаменатель; потом нормировать; потом взвешенную сумму — каждый шаг заново читает все блоки.

Трюк: «сначала ошибёмся, потом поправим». Примем максимум первого блока $d_1$ за глобальный и посчитаем частичную сумму $s_1$. Встретив в следующем блоке больший максимум $d_2$, не перечитываем старое — домножаем:

$$s_k = s_{k-1}\, e^{\,d_{k-1}-d_k} + \sum_{i \in \text{блок } k} e^{\,a_i - d_k}, \qquad d_k = \max(d_{k-1},\ \max_{i \in \text{блок } k} a_i).$$

Ли называет это «всего лишь одной формулой» — но весь FlashAttention повторяет её снова и снова. Следующий вопрос лектора: обязательно ли сначала вычислить веса внимания, чтобы взять взвешенную сумму? Нет: читаем вместе с ключами и значения $v$ и сразу копим выход с «неправильными» весами, а затем поправляем и его:

$$O_k = O_{k-1}\,\frac{s_{k-1}}{s_k}\, e^{\,d_{k-1}-d_k} + \sum_{i \in \text{блок } k} \frac{e^{\,a_i - d_k}}{s_k}\, v_i .$$

После последнего блока $O$ в точности равен обычному $\sum_i \hat a_i v_i$, а матрица весов так и не была записана в память. Побочный эффект: в Hugging Face с FlashAttention нельзя достать веса внимания для анализа — их просто нет.

Проверим на числах

Пусть $a = (1, 3 \mid 2, 4)$ — два блока. Блок 1: $d_1 = 3$, $s_1 = e^{-2} + e^{0} \approx 1{,}135$. Блок 2: $d_2 = 4$, $s_2 = 1{,}135 \cdot e^{-1} + e^{-2} + e^{0} \approx 0{,}418 + 0{,}135 + 1 = 1{,}553$. Напрямую: $\sum_i e^{a_i - 4} = e^{-3} + e^{-1} + e^{-2} + e^{0} \approx 0{,}050 + 0{,}368 + 0{,}135 + 1 = 1{,}553$ — совпало.

Что показывает Colab курса. PyTorch scaled_dot_product_attention с бэкендами MATH и FLASH_ATTENTION: результаты совпадают до ~$10^{-7}$; на A100 при длине 4096 FlashAttention примерно в 9 раз быстрее (в сохранённом выводе 9,49×). На реальной модели (google/gemma-3-4b-it, eager против sdpa) на коротких входах разницы почти нет — много времени уходит на полносвязные слои, — а на длинных выигрыш заметен; если удлинить вход ещё в 10 раз — CUDA out of memory. Переполнился склад, а не верстак: почему — следующая лекция. Обычно PyTorch и так выбирает FlashAttention по умолчанию, то есть вы им, скорее всего, уже пользуетесь.

Лекция «Ускоряем генерацию, часть 2: KV-кэш»

Видео на YouTube · слайды inference.pdf, с. 29–55.

Что такое KV-кэш. После префилла запросы $q$ выбрасывают, а ключи и значения хранят: для нового токена нужны только его собственные $q, k, v$ и старые $k, v$. Ли шутит, что cache и cash звучат одинаково — и в конце лекции окажется, что кэш действительно про деньги.

Почему кэш переполняет склад. Ключи и значения хранятся для каждого токена, каждой головы и каждого слоя. Пример по архитектуре Gemma 2 27B (если считать, что у всех 32 голов свои $k, v$):

$$46\ \text{слоёв} \times 32\ \text{головы} \times 128 \times 2\ \text{байта (FP16)} \times 2\ (k \text{ и } v) = 753\,664\ \text{байта} \approx 0{,}72\ \text{МБ на токен},$$

то есть на 80 ГБ — около 114 тысяч токенов, а агентам нужно больше. Способы «уместить больше»:

1. Делить $k, v$ между запросами. MQA — одна пара $k, v$ на все головы (по словам Ли, качество страдает); GQA — пара на группу голов (Llama, Gemma). Почему делят именно $k, v$, а не $q$? Потому что хранятся только $k, v$.
2. Сжимать $k, v$ — MLA (DeepSeek-V2): хранить маленький латентный вектор $c_i$, а ключи и значения получать умножением на матрицы. Распаковывать не нужно: $q^\top W_K c_i = (W_K^\top q)^\top c_i$ — достаточно один раз преобразовать запрос, а взвешенную сумму считать в сжатом пространстве и распаковать один раз в конце. Это точное тождество, не приближение; модель нужно обучать так с самого начала.
3. Ограничить область внимания. Скользящее окно (например, 4096 токенов) даёт кэшу фиксированный предел, а многослойность позволяет всё равно «видеть» далеко (в GPT-OSS слои с окном чередуются с полными). StreamingLLM: чистое окно ломается на длинных входах, но если всегда держать в окне первые несколько токенов, качество возвращается — веса внимания обязаны суммироваться в 1, и «когда смотреть не на что», модель по привычке смотрит на первый токен (attention sink).
4. Выбрасывать неиспользуемые $k, v$ (Scissorhands, H2O): внимание сосредоточено на немногих токенах; то, на что долго никто не смотрит, удаляют. Scissorhands сжимал кэш в 5 раз почти без потерь, но на трудных задачах такая обрезка вредит — тема ещё исследуется.

Кэш между диалогами — это деньги. «Всем привет, я Дацзинь» и «Всем привет, я Сяоцзинь» совпадают в первых пяти токенах — их $k, v$ можно переиспользовать. А дальше даже одинаковый иероглиф «цзинь» уже не совпадёт: представление токена зависит от всего, что было до него. Переиспользуется только точно совпадающий префикс. Поэтому провайдеры делают скидку: в прайсе OpenAI на слайде кэшированный ввод в 10 раз дешевле обычного. Для агентов это золото: в системном промпте всё неизменное (инструменты, правила, личность, список навыков) ставьте в начало, изменчивое (например, текущую дату) — в конец; и шаблоны пишите так, чтобы переменные шли последними («забронируй рейс из x в y» + значения x, y в конце). В работе «Don’t Break the Cache» на долгих агентных задачах кэширование снизило стоимость на 46–80 % в зависимости от модели.

Итоговая таблица Ли:

Метод	Что делает	Меняет внимание?	Нужно обучать?	Другая цена
FlashAttention	меньше перевозок данных	нет	нет	чуть больше вычислений, сложнее понять
KV-кэш	хранить посчитанные $k, v$	нет	нет	память
MQA	общие $k, v$ для всех запросов	да	да	может заметно ухудшить модель
GQA	общие $k, v$ по группам	да	да
MLA	сжатые $k, v$	да	да
Скользящее окно	ограничить область внимания	да	по-разному
StreamingLLM	окно + первые токены	да	по-разному
Обрезка KV-кэша	выбросить $k, v$	да	нет	может заметно ухудшить модель
Спекулятивное декодирование	маленькая модель «предсказывает»	нет (теоретически)	нет	лишние вычисления малой модели

Связь с курсом и с вашей машиной

На 4× RTX 4080 32 ГБ KV-кэш — главный ограничитель длины контекста в vLLM. Посчитайте его для своей модели по config.json (num_hidden_layers, num_key_value_heads, head_dim) — так же, как в примере выше, но с числом KV-голов (GQA). В vLLM включите --enable-prefix-caching и проверьте выигрыш на агенте с длинным системным промптом (задание HW3 курса делает этот опыт). Уроки: 4.6.4, 5.2.1–5.2.2, 7.1.2.

Лекция «Как Transformer узнаёт порядок токенов: абсолютные, относительные, RoPE и вовсе без позиций»

Видео на YouTube · слайды pos.pdf (64 слайда), Colab.

Зачем позиции. Для токена D внимание к «A B C» и к «C B A» даёт одну и ту же взвешенную сумму, а «ты бьёшь меня» и «я бью тебя» — противоположности.

1. Абсолютные позиции. К вектору токена прибавляется вектор позиции. Синусоидальный вариант «Attention Is All You Need» — набор часовых стрелок: высокочастотные пары координат — «секундные», различают соседей; низкочастотные — «часовые», различают далёкие позиции (на слайде: полный оборот за 6,3, 628,3 и 54 410,1 позиции у разных пар). Но важно относительное расстояние: связь «кошка — рыба» не должна зависеть от того, стоит фраза в начале или на 101-й позиции. Если разложить $(x_m + p_m)\cdot(x_n + p_n)$, в синусоидальном варианте не видно, чтобы «позиционный» член зависел только от $m - n$.

2. Относительные позиции. ALiBi: из оценки внимания вычитается $b\,(m-n)$ — «чем дальше, тем меньше внимания, вот и всё»; $b$ задают вручную, разным головам разный. T5: добавка к оценке зависит от расстояния и обучается.

3. RoPE (Llama, Qwen, Gemma). Каждую пару координат $q$ и $k$ поворачиваем на угол $n\theta_i$, где $n$ — позиция; в скалярное произведение входит только разность углов $(m-n)\theta_i$, поэтому «кошка» и «рыба» на позициях 1 и 3 дают ту же оценку, что на 101 и 103. Не меняет вычисление внимания и совместим с KV-кэшем. Углы $\theta_i = 1/10000^{2i/d}$ — та же «секундная/часовая» структура, только вместо сложения — вращение. Распространённое мнение «RoPE ослабляет внимание с расстоянием» Ли опровергает со ссылкой на «Round and Round We Go!»: затухания нет, и это не обязательно плохо — запрос можно обучить смотреть ровно на два токена назад.

4. Обучение на коротком — применение к длинному. На обучении ключ поворачивался максимум на $N\theta$, на длинном тесте — на $2N\theta$: модель никогда не видела таких направлений. Исправления:

Метод	Идея	Цена
Position Interpolation	сжать все позиции в обучающий диапазон (1, 2, 3, 4 → 0,5, 1, 1,5, 2)	нужно дообучение
частотный подход	быстрые пары и так «прокрутились» много раз — их не трогаем; сжимаем только медленные	решить, какие пары «быстрые»
NTK-aware scaling	плавно: самая быстрая пара не меняется, самая медленная сжимается в $L$ раз	идея из поста на Reddit, не из статьи
YaRN	«статейная» версия частотного подхода
динамическое масштабирование	сжимать только когда вход длиннее обучающего (иначе короткие тексты портятся)	коэффициент зависит от длины
LongRoPE	частотный + динамический подход, коэффициенты ищутся эволюционным поиском	нужен поиск

5. А нужны ли позиции вообще? В каузальном декодере каждый токен видит только себя и прошлое: в «кошка ест рыбу?» и «рыба ест кошку?» последний токен видит одно и то же множество слов, но промежуточные токены видят разные префиксы, и через слои порядок всё равно «просачивается». NoPE: без позиционного кодирования длина обобщается лучше, чем с APE, T5, ALiBi или RoPE (на их задачах рассуждения и математики). DroPE: позиции помогают сойтись при предобучении, но мешают обобщаться на длину; после обучения их можно убрать и после короткой перекалибровки удлинять контекст без обучения.

Связь с курсом

Синусоиды и RoPE — урок 4.5.2 (и наш перевод Qwen-blog с RoPE «на числах»); расширение контекста — 5.5.1. Упражнение из разбора: откройте config.json любимой открытой модели, найдите rope_theta и rope_scaling и определите, какой это метод из таблицы.

---

Что сделать руками

1. Онлайн-softmax. Напишите на NumPy функцию, которая считает softmax-взвешенную сумму блоками по формулам выше, за один проход, и сравните с обычным расчётом на случайных данных (ошибка должна быть около $10^{-15}$ в float64).
2. Калькулятор KV-кэша. Посчитайте байты на токен и максимальный контекст на одной RTX 4080 32 ГБ для Qwen3-8B и Qwen3-32B (возьмите числа слоёв и KV-голов из их config.json; вычтите память под веса).
3. Кэш префикса. Отправьте в vLLM с --enable-prefix-caching 20 запросов с одинаковым длинным системным промптом — сначала с датой в начале промпта, потом в конце; сравните время до первого токена.

Словарик

English	Русский	Пояснение
Prefill / Decode	префилл / декодирование	обработка подсказки / генерация по токену
HBM / SRAM	«склад» / «верстак»	большая медленная / маленькая быстрая память GPU
online softmax	онлайн-softmax	поблочный softmax с поправками
KV Cache	KV-кэш	сохранённые ключи и значения
MQA / GQA / MLA	общие / групповые / латентные $k, v$	способы уменьшить кэш
Sliding Window	скользящее окно	внимание только на последние $W$ токенов
attention sink	«сток внимания»	первые токены, куда уходит «лишнее» внимание
prompt caching	кэширование префикса	переиспользование $k, v$ между запросами
Positional Embedding	позиционный эмбеддинг
RoPE	вращательное позиционное кодирование
Train short, test long	«учим на коротком — применяем к длинному»	расширение контекста
NoPE	без позиционного кодирования	No Positional Embedding