Глава 3.5 — Варианты внимания для инференса

Содержание

  1. Где остановилась глава 3.4
  2. Сколько на самом деле стоит KV-кэш
  3. Multi-Query Attention: одна голова ключей и значений на всех
  4. Grouped-Query Attention: интерполяция, которая победила
  5. Multi-head Latent Attention: кэшировать сжатие вместо самих векторов
  6. Как выбирать между ними
  7. Взгляд с точки зрения собеседования
  8. Вопросы для самопроверки
  9. Источники

1. Где остановилась глава 3.4

Глава 3.4 представила KV-кэш как приём, который вообще делает авторегрессивную генерацию посильной, и закрыла раздел предупреждением: кэш — это память, эта память растёт линейно по длине последовательности и по размеру батча, и при обслуживании длинного контекста с высокой конкурентностью именно она становится связывающим ограничением на то, сколько запросов система удерживает одновременно. Всё, что та глава предложила в ответ, было исправлением на этапе развёртывания — квантовать кэш, разбить его на страницы, ограничить скользящим окном — применённым к модели, архитектура которой уже зафиксирована.

Эта глава — про вторую половину ответа, ту, которую нужно решить до начала обучения. Между 2019 и 2024 годами область внесла три последовательных изменения в сам механизм внимания, каждое нацеленное прямо в размер KV-кэша и каждое принимающее свой компромисс ради этого: Multi-Query Attention, Grouped-Query Attention и Multi-head Latent Attention. Это не мелкие детали реализации. Grouped-Query Attention используется практически в каждой открытой модели, выпущенной после 2023 года, и знание того, какой из вариантов стоит в модели, скажет вам о её экономике обслуживания больше, чем количество параметров. Эта глава ещё и первое место в книге, где проектное решение диктуется целиком стоимостью инференса, а не качеством или стоимостью обучения, — и это само по себе стоит усвоить, потому что паттерн повторяется.

2. Сколько на самом деле стоит KV-кэш

Начнём с того, чтобы записать стоимость точно. Для одной последовательности на каждом слое кэш хранит по вектору ключа и вектору значения на токен на голову. При $L$ слоях, $h$ головах размерности $d_h$, длине последовательности $n$, размере батча $b$ и $p$ байтах на элемент общий объём равен $$\text{байт кэша} = 2 \cdot L \cdot h \cdot d_h \cdot n \cdot b \cdot p$$ где ведущая двойка — это ключи и значения. Обратите внимание, чего в формуле нет: ничего про число параметров и ничего квадратичного. Кэш растёт линейно по контексту и по батчу, причём с константой, которую целиком задаёт форма модели.

Подставим реальные числа. У LLaMA-2 70B $L = 80$ слоёв, $h = 64$ головы запросов и $d_h = 128$. Обслуживание одной последовательности длиной 4096 токенов в fp16 при обычном многоголовом кэше потребовало бы $2 \cdot 80 \cdot 64 \cdot 128 \cdot 4096 \cdot 2$ байт — около $10{,}7$ ГБ. Для одной последовательности, при длине контекста, которая по меркам 2024 года мала, на модели, чьи fp16-веса занимают $140$ ГБ. Соберите шестнадцать таких запросов в батч или растяните контекст до 32K — и один кэш превзойдёт веса. Именно это число делает остаток главы необходимым: пропускная способность системы обслуживания — это примерно «сколько одновременных последовательностей влезает в память, оставшуюся после весов», а определяет это кэш.

Объём памяти — только половина проблемы и, пожалуй, менее важная. Раздел 2 главы 3.4 установил, что авторегрессивное декодирование ограничено не вычислениями, а пропускной способностью памяти, и KV-кэш — самая наглядная иллюстрация почему. На каждом шаге декодирования вычисление внимания читает весь кэш — каждый ключ и значение каждого предыдущего токена — и делает примерно одно умножение-накопление на прочитанный элемент. Это арифметическая интенсивность около одного FLOP на байт против современных ускорителей, которым нужны сотни FLOP на байт, чтобы загрузить вычислительные блоки. GPU простаивает в ожидании памяти подавляющую часть каждого шага декодирования. Формулировка Шазира в статье, с которой началась вся эта линия работ, стоит того, чтобы привести её в его терминах: узкое место — не арифметика, а отношение обращений к памяти к арифметике, и способ ускорить декодирование — грузить меньше байт на шаг. Каждый приём в этой главе уменьшает одно и то же число и получает экономию памяти и экономию пропускной способности памяти (полосы) от одного и того же изменения.

3. Multi-Query Attention: одна голова ключей и значений на всех

Предложение Шазира 2019 года — самый агрессивный доступный ход и самый простой в формулировке: сохранить все $h$ голов запросов, но выдать всему слою одну голову ключей и одну голову значений, общие для всех голов запросов. Запросы остаются полностью многоголовыми — каждый по-прежнему проецируется в собственное подпространство и вычисляет своё распределение внимания, — но все они обращаются к одним и тем же ключам и одним и тем же значениям. Член кэша $h \cdot d_h$ схлопывается до $d_h$, уменьшая кэш в $h$ раз: для формы LLaMA-2 70B выше $10{,}7$ ГБ превращаются в $167$ МБ.

Стоит точно проговорить, чего это стоит, а чего нет, потому что интуиция «вы убрали большую часть способности модели к вниманию» неверна. Количество паттернов внимания не изменилось — по-прежнему есть $h$ различных проекций запросов, порождающих $h$ различных распределений по последовательности. Теряется возможность разным головам смотреть на разные проекции одних и тех же токенов: каждая голова теперь оценивает относительно одного общего взгляда на контекст и извлекает из одного общего пространства значений. Аргумент главы 2.1 в пользу многоголового внимания состоял в том, что головы специализируются — одна отслеживает синтаксическую зависимость, другая кореференцию; MQA позволяет им продолжать специализироваться в том, что они спрашивают, заставляя делить то, что они видят.

Эмпирически это чего-то стоит. Шазир сообщил о небольшой, но измеримой деградации качества на переводе, а более поздняя статья про GQA нашла эффект побольше в масштабе плюс вторую проблему, которая на практике важнее: MQA-модели бывают нестабильны при обучении, особенно на длинных последовательностях, и разрыв в качестве не закрывается простым продлением обучения. Сочетание — реальная потеря качества и реальная сложность обучения в обмен на очень крупное сокращение кэша — объясняет, почему MQA получил настоящее, но ограниченное распространение (PaLM и Falcon среди заметных пользователей), а не стал стандартом. Зато он задал ось, и как только вы начинаете видеть во внимании настраиваемое число голов ключей и значений, напрашивается вопрос: а действительно ли две крайности — единственные варианты?

4. Grouped-Query Attention: интерполяция, которая победила

Ответ Ainslie et al., опубликованный в 2023 году, — ровно та интерполяция, которую вы бы угадали, и его успех хорошо иллюстрирует, как часто полезным вкладом оказывается скучная середина уже существующего спектра. Разбейте $h$ голов запросов на $g$ групп; дайте каждой группе свою голову ключей и свою голову значений, общие для $h/g$ голов запросов внутри неё. При $g = h$ получается стандартное многоголовое внимание, при $g = 1$ — MQA, а всё между ними — Grouped-Query Attention, и кэш уменьшается в $h/g$ раз. Типичный выбор в продакшн-моделях — $g = 8$: LLaMA-2 70B использует восемь голов ключей-значений против шестидесяти четырёх голов запросов, сокращая кэш с $10{,}7$ ГБ до $1{,}34$ ГБ при контексте 4K и удерживая качество в пределах шума относительно полной многоголовой модели.

Эмпирический факт, который делает GQA стоящим применения, — это резкая нелинейность кривой «качество против кэша» вблизи конца MQA. Переход с $h$ голов ключей-значений до $8$ не стоит почти ничего измеримого; переход с $8$ до $1$ стоит заметно больше. Почти всё сокращение памяти происходит на первом шаге (переход с 64 голов на 8 забирает 87,5% возможной экономии), а почти вся потеря качества — на втором, так что середина диапазона — это не компромисс между двумя хорошими вариантами, а строго лучший выбор, чем одна из крайностей, при любой реалистичной нагрузке. GQA заодно исправляет нестабильность обучения MQA, что для его распространения, пожалуй, столь же важно, как и цифра качества.

Второй вклад статьи — практический, и именно о нём спрашивают на собеседованиях. Если GQA обязательно обучать с нуля, то каждый существующий многоголовый чекпоинт застревает со своим кэшем навсегда — а в 2023 году это означало каждую хорошую открытую модель. Ainslie et al. показали, что существующий чекпоинт можно конвертировать: усреднить матрицы проекций ключей и значений голов внутри каждой группы в одну проекцию на группу, что даёт модель, структурно являющуюся GQA и сразу выдающую осмысленный (хоть и ухудшенный) выход, а затем недолго продолжить предобучение, чтобы остальная сеть адаптировалась. Они называют это uptraining, и поразителен именно объём: примерно 5% исходного вычислительного бюджета предобучения хватает, чтобы восстановить по сути полное качество. Это число и превратило GQA из архитектурного предложения в то, что любой владелец существующей модели может внедрить за вечер кластерного времени, — и объясняет, почему приём разошёлся по экосистеме открытых весов так быстро.

5. Multi-head Latent Attention: кэшировать сжатие вместо самих векторов

И MQA, и GQA работают через уменьшение числа голов ключей-значений. Multi-head Latent Attention из DeepSeek-V2 атакует тот же член иначе: сохранить все головы, но сжать то, что вы храните. Для каждого токена спроецируйте вход вниз в единый низкоранговый латентный вектор $c_t \in \mathbb{R}^{d_c}$, где $d_c$ много меньше $h \cdot d_h$, кэшируйте только этот латент и восстанавливайте полные ключи и значения каждой головы по требованию матрицами обратной проекции $W^{UK}$ и $W^{UV}$: $$c_t = x_t W^{DKV}, \qquad k_t = c_t W^{UK}, \qquad v_t = c_t W^{UV}$$ DeepSeek-V2 использует $d_c = 512$ против 128 голов размерности 128, так что кэшируемый объект на токен на слой — это $512$ чисел, а не $2 \cdot 128 \cdot 128 = 32768$.

В таком изложении кажется, что стоимость просто переехала: вы сэкономили память, но теперь должны разжимать кэш на каждом шаге, а это ровно та арифметика, которой вы пытались избежать. Изящество в том, что не должны. Оценкам внимания нужно $q_t^\top k_s = (x_t W^Q)(c_s W^{UK})^\top = x_t \left(W^Q {W^{UK}}^\top\right) c_s^\top$ — а $W^Q {W^{UK}}^\top$ есть произведение двух фиксированных матриц весов, так что его можно вычислить один раз и вложить в проекцию запросов. То же поглощение работает на выходной стороне, где $W^{UV}$ сливается с $W^O$. На инференсе обратные проекции никогда не выполняются: запросы проецируются прямо в латентное пространство и обращаются напрямую к кэшированным латентам. Вы получаете и маленький кэш, и сниженную полосу вовсе без шага распаковки.

Есть одна настоящая сложность, и это та деталь, которая отличает «читал про MLA» от «понял MLA». Трюк с поглощением требует, чтобы между $W^Q$ и $W^{UK}$ не стояло ничего зависящего от позиции, — а RoPE (глава 2.2) работает именно вставкой туда зависящего от позиции поворота, и $R_{t-s}$ зависит от конкретной пары запрос-ключ, так что его нельзя свернуть в фиксированную матрицу. RoPE и поглощение весов напрямую несовместимы. Решение DeepSeek, называемое decoupled RoPE, состоит в том, чтобы разделить размерность головы надвое: бо́льшая часть не несёт позиционного поворота и остаётся поглощаемой, а небольшой дополнительный срез (64 измерения в DeepSeek-V2) несёт RoPE и обрабатывается отдельно, как одна общая голова ключей в стиле MQA, кэшируемая рядом с латентом. Итоговая оценка внимания — это просто сумма двух частей: $q_t^\top k_s = x_t\left(W^Q {W^{UK}}^\top\right)c_s^\top + (q_t^R)^\top R_{t-s}\, k_s^R$, поглощённого слагаемого и слагаемого, несущего RoPE. Кэш на токен становится $d_c + d_h^R = 512 + 64 = 576$ чисел на слой, то есть примерно на 98% меньше, чем $2 \cdot 128 \cdot 128 = 32768$, которые хранил бы многоголовый слой той же формы; сами DeepSeek-V2 сообщают о сокращении на 93,3% относительно их предыдущей плотной DeepSeek 67B. Их более интересное утверждение — по другой оси: они измеряют MLA как сравнимое или слегка превосходящее полное многоголовое внимание по качеству, что — если подтвердится — сделало бы его первой точкой этого спектра, вообще не являющейся компромиссом. Этот результат новее и получил меньше независимых подтверждений, чем результат GQA, и приводить его стоит именно с такой оговоркой.

6. Как выбирать между ними

Если записать кэш всех четырёх схем на токен на слой в одних единицах, прогресс становится очевиден. Многоголовое внимание хранит $2 \cdot h \cdot d_h$ чисел; MQA — $2 \cdot d_h$; GQA — $2 \cdot g \cdot d_h$; MLA — $d_c + d_h^R$. Для модели формы LLaMA-2 70B ($h = 64$, $d_h = 128$) это $16384$, $256$, $2048$ при $g=8$ и — по пропорциям DeepSeek — что-то в районе $576$. Порядок по качеству для первых трёх примерно обратный, а MLA претендует на то, чтобы выбиться из этой закономерности.

На практике решение обычно принято за вас тем, что вы делаете. Если вы обучаете новую модель с нуля, GQA при $g = 8$ — тот выбор по умолчанию, за который по сути никого не критикуют, а MLA — более амбициозный вариант, стоящий рассмотрения, если вы готовы взять на себя сложность decoupled RoPE и целитесь в длинные контексты, где кэш доминирует. Если у вас есть готовый многоголовый чекпоинт и проблема со стоимостью обслуживания, ответ — uptraining до GQA, и он дёшев. Если переобучать нельзя вовсе — вы возвращаетесь в мир главы 3.4: квантуйте кэш, разбивайте на страницы, ограничивайте окном, — и ничего из этой главы вам не поможет, что как раз и объясняет, почему важно, что это архитектурные решения, а не решения развёртывания.

Стоит также назвать, чему всё это ортогонально, а чему нет. Оно свободно сочетается с FlashAttention (меняет, как внимание вычисляется, а не что хранится), с квантованием кэша (меняет число байт на элемент, а не число элементов), с PagedAttention (меняет размещение, а не размер) и со скользящим окном (ограничивает $n$, а не константу на токен). Но эти схемы не сочетаются друг с другом — у слоя одна организация голов ключей-значений, — так что это настоящий выбор, а не ещё один пункт для стека. Часть III к этому моменту разобрала удешевление внимания по длине последовательности (глава 3.1), обобщение позиции дальше обученного диапазона (глава 3.2), покупку параметров без вычислений (глава 3.3), обслуживание результата (глава 3.4) и сжатие того, что обслуживанию приходится помнить (эта глава). Всё это предполагало, что механизм — softmax-внимание, и вопрос лишь в том, как за него заплатить. Глава 3.6 закрывает часть III, разбирая это предположение на части.

7. Взгляд с точки зрения собеседования

Собеседование на senior-позицию по LLM очень часто спрашивает про KV-кэш, и именно на этом материале кандидат либо пересказывает определение, либо показывает, что ему приходилось за кэш платить:

  • «Ваш KV-кэш съедает память GPU при высокой конкурентности. Разберите варианты.» Сильный ответ сортирует варианты по тому, требуют ли они переобучения. На этапе развёртывания: квантовать кэш, использовать постраничное несмежное размещение, чтобы убрать потери на фрагментацию, ограничить окно. Архитектурно: сократить головы ключей-значений через GQA или перейти на латентный кэш через MLA. Самые сильные ответы добавляют, что если модель — существующий многоголовый чекпоинт, GQA достижим через uptraining, а не полное переобучение, что делает его живым вариантом, а не пунктом «в следующей модели».
  • «Чем именно Multi-Query Attention жертвует относительно многоголового и почему GQA его вытеснил?» Ответ должен точно указать, что MQA сохраняет все головы запросов и все паттерны внимания и делит только ключи и значения — так что головы по-прежнему специализируются в том, что спрашивают, но больше не видят разных проекций контекста. GQA победил, потому что кривая «качество против кэша» резко нелинейна: восемь голов ключей-значений забирают бо́льшую часть экономии памяти почти без потери качества, а переход к одной стоит непропорционально дороже, плюс у MQA есть проблемы со стабильностью обучения, которых у GQA нет.
  • «Как бы вы сконвертировали обученную многоголовую модель в GQA?» Усреднить проекции ключей и значений внутри каждой группы, чтобы инициализировать групповые проекции, затем провести uptraining на малой доле — Ainslie et al. использовали около 5% — исходного вычислительного бюджета предобучения. Кандидат, который знает про инициализацию именно усреднением, а не просто «дообучить», статью читал.
  • «Объясните, как MLA может кэшировать сжатый латент, не платя за распаковку, и что ломается при добавлении RoPE.» Это самая острая версия вопроса. Ожидаемый ответ: $W^{UK}$ можно поглотить в $W^Q$ (а $W^{UV}$ — в $W^O$), потому что обе матрицы фиксированы, так что запросы обращаются напрямую в латентном пространстве; а RoPE вставляет между ними зависящий от позиции поворот, который нельзя свернуть в фиксированное произведение, — поэтому DeepSeek отщепляет небольшую отдельную размерность под RoPE, обрабатываемую в стиле MQA.

Сквозная линия здесь — понимаете ли вы, что стоимость LLM состоит из двух во многом независимых половин. Почти всё в стоимости обучения — функция параметров и FLOPs; почти всё в стоимости обслуживания на тех уровнях конкурентности, на которых работают реальные продукты, — функция числа, которое в подсчёте параметров вообще не появляется. Кандидаты, которые только обучали модели, тянутся к квантованию по любому поводу; кандидаты, которые их обслуживали, первым делом смотрят на кэш.

8. Вопросы для самопроверки

  1. Выпишите формулу размера KV-кэша и определите, какие члены задаются архитектурой, какие нагрузкой, а какие конфигурацией обслуживания. Какие из них можно изменить без переобучения?
  2. Объясните, почему авторегрессивное декодирование ограничено пропускной способностью памяти, через арифметическую интенсивность шага внимания по кэшу, и почему из этого следует, что сжатие кэша покупает вам и задержку, а не только память.
  3. Что именно в MQA становится общим, а что остаётся на голову? Используйте это, чтобы объяснить, почему «MQA убирает бо́льшую часть способности модели к вниманию» — неверное описание.
  4. Почему $g = 8$ — частый выбор для GQA, а не $g = 2$ или $g = 32$? Сформулируйте ответ через форму кривой «качество против памяти», а не как произвольное соглашение.
  5. Опишите процедуру uptraining для конвертации многоголового чекпоинта в GQA, включая то, как инициализируются групповые проекции и во сколько это примерно обходится.
  6. Объясните аргумент о поглощении весов, позволяющий MLA не распаковывать кэш, а затем точно объясните, почему добавление RoPE его ломает и что с этим делает decoupled RoPE.
  7. Для каждого из FlashAttention, квантования кэша, PagedAttention и скользящего окна скажите, сочетается ли оно с GQA и почему, — и объясните, почему GQA и MLA не сочетаются друг с другом.

9. Источники

  • Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150. https://arxiv.org/abs/1911.02150
  • Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. arXiv:2305.13245. https://arxiv.org/abs/2305.13245
  • DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434. https://arxiv.org/abs/2405.04434
  • DeepSeek-AI (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437. https://arxiv.org/abs/2412.19437
  • Touvron, H., Martin, L., Stone, K., et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288. https://arxiv.org/abs/2307.09288
  • Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, S., & Dean, J. (2022). Efficiently Scaling Transformer Inference. arXiv:2211.05102. https://arxiv.org/abs/2211.05102

Тест для самопроверки

Проверьте понимание главы с помощью короткого теста — вопросы по теории и небольшие расчёты.

У модели $L = 30$ слоёв, $h = 40$ голов, $d_h = 128$. При обслуживании одной последовательности в 2048 токенов в fp16 с полным многоголовым KV-кэшем — сколько гибибайт занимает кэш?

Объяснение: $2 \cdot 30 \cdot 40 \cdot 128 \cdot 2048 \cdot 2 = 1{,}258{,}291{,}200$ байт $= 1.17$ ГиБ — для одной последовательности.

Почему сжатие KV-кэша выигрывает не только память, но и задержку?

Объяснение: Каждый шаг декодирования прогоняет весь кэш из памяти и делает примерно одно умножение-накопление на прочитанный элемент. При ~1 FLOP/байт против ускорителей, которым нужны сотни, GPU простаивает в ожидании памяти, так что загрузка меньшего числа байт напрямую укорачивает шаг.

Что в Multi-Query Attention становится общим для голов, а что остаётся на голову?

Объяснение: MQA сохраняет все $h$ проекций запросов, так что различных распределений внимания по-прежнему $h$. Теряется то, что головы больше не видят разных проекций контекста — все они оценивают относительно одного общего взгляда.

У модели 64 головы запросов. Во сколько раз конвертация в GQA с 8 группами ключей-значений уменьшает KV-кэш?

Объяснение: Кэш масштабируется числом голов ключей-значений, поэтому множитель равен $h/g = 64/8 = 8$ — это забирает 87,5% той экономии, которую дал бы MQA.

Почему $g = 8$ — частая настройка GQA, а не $g = 1$?

Объяснение: Переход с 64 голов ключей-значений на 8 не стоит почти ничего измеримого; переход с 8 до 1 стоит заметно больше, а у MQA вдобавок есть проблемы со стабильностью обучения, которых GQA избегает.

Как uptraining у Ainslie et al. инициализирует групповые проекции ключей и значений при конвертации многоголового чекпоинта?

Объяснение: Усреднение внутри каждой группы даёт структурно GQA-модель, которая уже выдаёт нечто осмысленное; затем примерно 5% исходного бюджета предобучения восстанавливают по сути полное качество.

Благодаря чему сжатый кэш MLA не требует распаковки во время инференса?

Объяснение: $q_t^\top k_s = x_t (W^Q {W^{UK}}^\top) c_s^\top$, а $W^Q {W^{UK}}^\top$ — произведение двух фиксированных матриц весов, вычисляемое один раз. Обратные проекции на инференсе не выполняются никогда.

Почему RoPE ломает трюк с поглощением весов в MLA и в чём состоит решение DeepSeek?

Объяснение: $R_{t-s}$ зависит от конкретной пары запрос-ключ, поэтому его нельзя домножить на фиксированные веса заранее. Decoupled RoPE оставляет бо́льшую часть размерности головы поглощаемой, а RoPE проводит через небольшой дополнительный срез (64 измерения в DeepSeek-V2), кэшируемый как одна общая голова ключей.

Что из перечисленного НЕ сочетается с GQA?

Объяснение: MLA и GQA — альтернативные ответы на один и тот же вопрос: у слоя ровно одна организация голов ключей-значений. Остальные три меняют то, как внимание вычисляется, сколько байт занимает элемент и как размещается память, — ничто из этого с GQA не конфликтует.