#Быстрый инференс

Сэмплирование выбирает токен по вероятностям модели, и эта часть дешёвая. Дорого получить сами вероятности: на каждый токен нужен полный проход через модель. Обучение модели происходит один раз. Инференс — запуск модели, чтобы ответить на запрос, — происходит миллиарды раз, поэтому почти вся инженерная работа в обслуживании LLM (запуске моделей для множества пользователей сразу) уходит на то, чтобы каждый токен стал дешевле и быстрее.

Глава начинается с железа, потому что почти каждый приём в ней вытекает из одного факта о том, на что GPU тратит время. Дальше идут главные приёмы — примерно в том порядке, в каком с ними встречается запрос.

#Зачем вообще GPU

Трансформер в основном состоит из больших матричных произведений: каждый слой умножает векторы токенов на матрицы весов, а внимание умножает запросы на ключи. GPU (графический процессор) — это чип с тысячами маленьких ядер, созданный, чтобы делать много умножений-сложений одновременно. У CPU, центрального процессора компьютера, всего несколько больших ядер, которые быстро делают одно действие за другим.

Ваш браузер умеет запускать код прямо на GPU через WebGPU. Рисунок ниже дважды перемножает две случайные матрицы: один раз обычным циклом на JavaScript на CPU, другой — крошечным вычислительным шейдером, маленькой программой для GPU, где каждой клетке результата достаётся свой поток (один из тысяч параллельных исполнителей).

Скорости указаны в GFLOP/s — миллиардах операций с плавающей точкой (арифметики с дробными числами) в секунду. На большинстве машин даже наивный шейдер заметно обгоняет CPU. Серверный GPU с настроенными ядрами быстрее ещё в тысячи раз.

#Вычисления и память

Рисунок измерял только арифметику. Но числа сначала должны доехать до ядер, и на GPU эта дорога часто оказывается медленнее самих вычислений. Здесь важны две части чипа:

• Вычисления — ядра. H100, серверный GPU от NVIDIA, выполняет около 1000 триллионов операций в секунду (989 TFLOP/s в 16-битной арифметике).
• Память — там лежат веса и всё остальное. Основная память, HBM (high-bandwidth memory, память с высокой пропускной способностью), вмещает у H100 80 ГБ. Ещё у чипа есть немного SRAM — быстрой памяти прямо рядом с ядрами: всего несколько десятков мегабайт, зато добраться до неё гораздо быстрее.

У памяти два разных числа, и в этой главе нужны оба. Объём — сколько в неё помещается: 80 ГБ. Пропускная способность — как быстро данные из неё поступают в ядра: 3,35 ТБ/с (терабайта в секунду) у HBM.

Поделим одну скорость на другую: 989 TFLOP/s ÷ 3,35 ТБ/с ≈ 300. Пока из HBM приходит один байт, ядра успевают сделать около 300 операций. Поэтому скорость вычисления определяет его арифметическая интенсивность — сколько операций оно делает на каждый прочитанный байт.

• Если операций на байт заметно больше 300, предел ставят ядра. Такая работа упирается в вычисления (compute-bound).
• Если меньше, ядра простаивают в ожидании данных. Работа упирается в память (memory-bound): чип с более быстрой арифметикой не поможет, поможет только читать меньше байт.

Разницу хорошо видно на двух простых случаях. Умножение матрицы на матрицу много раз использует каждое прочитанное число — по разу на каждую строку второй матрицы, так что интенсивность высокая. Умножение матрицы на один вектор использует каждое число один раз: 16-битный вес занимает 2 байта и участвует в одном умножении и одном сложении — 2 операции на 2 байта, интенсивность около 1. Запомните эти два случая: это ровно две фазы инференса.

#Две фазы: префилл и декодирование

Запрос обрабатывается в две очень разные фазы:

1. Префилл (prefill). Весь промпт проходит через модель за один проход. Все токены промпта обрабатываются параллельно, так что каждый прочитанный вес используется по разу для каждого токена промпта. Это умножение матрицы на матрицу: GPU упирается в вычисления и загружен работой. От этой фазы зависит время до первого токена.
2. Декодирование (decode). Новые токены выходят по одному, ведь каждый зависит от предыдущего. Каждый шаг проталкивает через все слои только самый новый токен (почему старым токенам не нужно проходить заново, показано в следующем разделе). Это умножение матрицы на один вектор, с интенсивностью около 1. От этой фазы зависит число токенов в секунду, которое вы видите в потоке.

Значит, медленная часть — декодирование, и причина в памяти, а не в математике. Но прежде чем мерить, насколько оно медленное, нужна структура данных, благодаря которой можно обрабатывать «только самый новый токен».

#KV-кэш

Пусть промпт — «Кошка села на тёплый», и модель только что выдала «мягкий». Чтобы получить следующий токен, она прогоняет «мягкий» через все слои. В каждом слое внимание сравнивает запрос нового токена с ключами всех пяти токенов и смешивает их значения с получившимися весами. Откуда берутся эти пять ключей и значений? Каждый — это обученная матрица, умноженная на вектор токена в этом слое, а этот вектор сам получен на выходе всех слоёв ниже.

Без кэша всё считается заново. Чтобы восстановить ключи и значения «Кошка села на тёплый», модели пришлось бы снова прогнать эти четыре токена через все слои. На следующем шаге — пять токенов, потом шесть. По сути, она повторяла бы весь префилл ради каждого нового токена, и общий объём работы рос бы квадратично по длине.

Хранить их безопасно. Внимание причинное: маска разрешает токену смотреть только на себя и на более ранние токены. Поэтому, когда приходит «мягкий», ничего из посчитанного для «Кошка села на тёплый» не меняется ни в одном слое. Их ключи и значения ровно такие же, как шагом раньше. KV-кэш хранит ключи и значения каждого токена в каждом слое, а каждый шаг декодирования добавляет всего одну новую строку на слой.

Ключи и значения, но не запросы. Запрос — это вопрос, который токен задаёт, когда считает свой собственный выход: «какие токены важны для меня?». Старые токены свои вопросы уже задали. Их выходы готовы и, благодаря маске, больше никогда не изменятся, так что их запросы больше не понадобятся. С ключами и значениями наоборот: их читают более поздние токены, так что они нужны каждому будущему токену. То же со всем остальным, что считает слой, например с выходами MLP: повторно смотрят только на ключи и значения, поэтому хранят только их.

Кэш превращает квадратичный пересчёт в линейную работу, но стоит памяти, и эта память растёт с каждым токеном каждого разговора. Для модели с $L$ слоями, $h_{kv}$ головами ключей и значений по $d_h$ чисел, контекстом в $n$ токенов и $b$ байтами на число:

Множитель 2 — за ключи и значения. У Llama 3 8B $L = 32$ слоя и $h_{kv} = 8$ голов ключей и значений по $d_h = 128$ чисел, по $b = 2$ байта на число: $2 \times 32 \times 8 \times 128 \times 2 = 131\,072$ байта, то есть 128 КБ на токен. Разговор на 8000 токенов держит около 1 ГБ кэша поверх 16 ГБ весов. Обслуживайте сотню разговоров — и 80 ГБ GPU заполнит кэш, а не модель.

#Общие ключи и значения: GQA и MQA

Один множитель в этой формуле можно уменьшить, не сокращая контекст, — число голов ключей и значений. Обычное многоголовое внимание (MHA) даёт каждой голове запросов свои головы ключей и значений. Внимание с общими ключами (MQA) отдаёт одну голову ключей и значений всем головам запросов. Групповое внимание (GQA) — посередине: головы запросов делятся на группы, и каждая группа делит одну голову ключей и значений. У Llama 3, например, 32 или 64 головы запросов, но только 8 голов ключей и значений. Запросы не кэшируются, так что много голов запросов памяти не стоит.

KV-кэш настолько ценен, что ему посвящена отдельная глава: его переиспользование между запросами (кэширование промптов и префиксов), хранение страницами (PagedAttention), сжатие и вытеснение разобраны в главе 11 «Кэширование». А в главе 14 показано многоголовое латентное внимание, которое ещё сильнее уменьшает кэш, сжимая ключи и значения в маленький латентный вектор.

#Почему декодирование упирается в память

Вернёмся к скорости. Чтобы получить один токен, каждый вес модели должен доехать из HBM до ядер, и для одного токена каждый вес используется ровно в одном умножении-сложении: интенсивность около 1. Модель на 8 миллиардов параметров в 16-битных числах занимает 16 ГБ, значит каждый шаг декодирования читает 16 ГБ. Насколько быстро это возможно, определяет пропускная способность памяти:

Это верхняя граница для модели на 8B на H100 при размере батча 1 (одна последовательность за раз). На каждом шаге нужно прочитать ещё и KV-кэш, так что на деле получается меньше. При этом ядра почти простаивают: они могли бы делать около 300 операций на байт, а декодирование даёт им примерно одну.

Это хорошо видно на картинке roofline («крыша»). По горизонтали — арифметическая интенсивность, по вертикали — достижимая скорость. Слева от излома (около 300 операций на байт) предел ставит пропускная способность памяти, и линия идёт вверх; справа — сами вычисления, и линия плоская.

Почти каждый приём ниже — это способ читать меньше байт на токен или получать больше токенов с каждого прочитанного байта:

• Батчинг — объединять многих пользователей: веса читаются один раз и используются для каждой последовательности в батче.
• Квантизация — хранить веса в 8 или 4 битах вместо 16 (глава 13).
• Меньший KV-кэш — GQA выше и приёмы из главы 11.
• Спекулятивное декодирование — проверять несколько токенов за одно чтение весов (ниже).

#Батчинг: статический и непрерывный

Батчинг — самый мощный рычаг, но запросы неудобно объединять: они приходят в разное время и хотят разное число токенов. Наивная схема, статический батч, собирает несколько запросов, гоняет их вместе, пока не закончатся все, и только потом берёт следующую группу. Ответ из трёх слов вынужден ждать в своём слоте, пока рядом пишется сочинение.

Непрерывный батчинг (continuous batching, его ещё называют in-flight или iteration-level batching) планирует работу на уровне отдельных шагов декодирования. После каждого шага готовые последовательности уходят, а ожидающие встают на их место, так что слот не простаивает, пока в очереди есть работа.

Все современные движки обслуживания (vLLM, SGLang, TensorRT-LLM, TGI) используют непрерывный батчинг. Новым запросам нужен ещё и префилл, который гораздо тяжелее шага декодирования; движки режут длинные промпты на куски (chunked prefill) и подмешивают их в батчи декодирования, чтобы пользователи с потоковым ответом не замирали, пока кто-то вставляет в чат целую книгу.

#Спекулятивное декодирование

Большая модель читает все свои веса, чтобы выдать один токен. Поскольку декодирование упирается в память, проверить несколько предложенных токенов ей стоит примерно столько же, сколько выдать один: веса читаются один раз и используются для каждой позиции, как при префилле. Спекулятивное декодирование этим пользуется:

1. Маленькая быстрая черновая модель угадывает несколько следующих токенов.
2. Большая модель проходит один раз по всем догадкам и считает, что выдала бы сама в каждой позиции.
3. Догадки сохраняются до первого расхождения, а на месте расхождения берётся собственный токен большой модели.

Когда черновик обычно угадывает, большая модель выдаёт несколько токенов за проход. При аккуратном правиле принятия и отклонения вывод в точности следует распределению вероятностей большой модели, так что спекуляция — чистое ускорение без изменения качества. Типичный выигрыш — в 2–3 раза. Есть варианты без отдельной черновой модели: к большой модели добавляют дополнительные головы предсказания — небольшие выходные слои, не путать с головами внимания (Medusa, EAGLE, предсказание нескольких токенов сразу).

#FlashAttention: внимание без большой матрицы

Внимание сравнивает каждый запрос с каждым ключом, получая матрицу оценок $n \times n$. Для контекста в 32 000 токенов в ней миллиард элементов на голову — слишком много, чтобы поместиться в SRAM. Прямолинейная реализация записывает её в HBM и читает обратно, и время работы определяет этот трафик памяти, а не арифметика.

FlashAttention никогда не строит полную матрицу. Он режет запросы и ключи на плитки, которые помещаются в SRAM, и для каждого блока запросов прогоняет через неё плитки ключей и значений, храня только текущий результат. Сложность — в softmax: ему нужны максимум и сумма по всей строке, прежде чем что-то можно нормировать. Решение — онлайн-softmax, который хранит текущий максимум $m$ и текущую сумму $\ell$ и пересчитывает накопленное, когда появляется больший максимум:

Результат — точное внимание, а не приближение, но с памятью, которая растёт линейно по длине контекста, и с трафиком памяти в несколько раз меньше. Больше всего это важно при префилле и обучении, где много запросов встречаются с множеством ключей; при декодировании запрос на шаге один, и вариант Flash-Decoding вместо этого делит его длинную строку ключей между ядрами. Все крупные стеки для обучения и инференса используют FlashAttention или его потомков.

#Когда одного GPU мало

Модели на 70B в 16 битах нужно 140 ГБ только под веса — больше, чем 80 ГБ одного GPU. Движки обслуживания делят модель между несколькими GPU:

• Тензорный параллелизм режет каждую матрицу весов на части, по одной на GPU. Каждый GPU считает свою часть каждого слоя, а после каждого слоя они обмениваются частичными результатами по быстрому каналу вроде NVLink. Каждый GPU читает только свою долю весов, так что тензорный параллелизм заодно умножает пропускную способность памяти, работающую на каждый токен.
• Конвейерный параллелизм раскладывает разные слои по разным GPU и передаёт активации дальше, как на сборочной линии.
• Экспертный параллелизм размещает разных экспертов модели со смесью экспертов (глава 13) на разных GPU.

Большие развёртывания сочетают все три подхода и часто выполняют префилл и декодирование на разных пулах GPU (дезагрегированное обслуживание): одна фаза упирается в вычисления, другая в память, и им нужны разные компромиссы железа.

#Итоги

• GPU успевает сделать около 300 операций, пока из основной памяти приходит один байт. Работа, где операций на байт меньше, упирается в память.
• Префилл обрабатывает промпт параллельно и упирается в вычисления; декодирование выдаёт один токен за проход и ограничено пропускной способностью памяти.
• KV-кэш хранит ключи и значения каждого токена — единственное, что читают более поздние токены, — так что каждый шаг декодирования обрабатывает только новый токен. Его память растёт с длиной контекста.
• GQA и MQA уменьшают этот кэш, деля головы ключей и значений между головами запросов.
• Батчинг использует каждый байт весов для многих пользователей; непрерывный батчинг держит батч полным шаг за шагом.
• Спекулятивное декодирование позволяет дешёвой черновой модели предлагать токены, которые большая модель проверяет за один проход.
• FlashAttention считает точное внимание плитка за плиткой в быстрой памяти на кристалле с помощью онлайн-softmax.
• Модели, слишком большие для одного GPU, делят тензорным, конвейерным и экспертным параллелизмом.

KV-кэш оказался главной структурой данных всего обслуживания. Следующая глава — о том, как выжать из него максимум: кэширование.