vLLM и PagedAttention: виртуальная память для инференса LLM

Инференс разделяется на обработку входного текста и последовательную генерацию ответа. На первом этапе, prefill, все токены запроса уже известны, поэтому вычисления хорошо укладываются в крупные матричные операции. На этапе decode следующий токен появляется только после предыдущего: приходится снова читать веса модели и обращаться к сохранённым ключам и значениям прошлых токенов. Эти данные образуют KV-кеш. Повторно вычислять всю историю было бы дорого, но хранить её тоже непросто. Чтобы чтение весов окупалось, сервер объединяет несколько запросов в один вычислительный пакет. Размер такого пакета ограничивает память под кеш каждого участника. Поэтому задача ускорения начинается с вопроса о числе одновременно обслуживаемых последовательностей, а не только о скорости умножения матриц. Быстрый одиночный запрос ещё не показывает, как сервер справится с потоком запросов разной длины: каждый удерживает состояние, пока продолжается генерация.

В примере из статьи модель OPT-13B с точностью FP16 занимает около 26 ГБ на GPU с 40 ГБ памяти. После весов и временных данных остаётся ограниченное пространство под растущий кеш. Для этой модели один токен требует примерно 800 КБ KV-данных, а последовательность длиной 2048 токенов — порядка 1,6 ГБ. При непрерывном размещении сервер заранее отводит большой участок под запрос, хотя будущая длина ответа неизвестна. Отсюда три вида потерь: резерв для ещё не сгенерированных токенов, незаполненный остаток после фактического окончания ответа и промежутки между участками разного размера. Ведущий разбирает варианты Orca, включая идеализированный случай, где длина ответа известна заранее: даже такое знание не устраняет все потери. Важна организация размещения целиком. Выделенная память может оставаться пустой и при этом быть недоступной соседнему запросу, который уже ждёт обслуживания.

PagedAttention переносит в инференс идею из операционных систем: логически непрерывная последовательность не обязана занимать непрерывный физический участок. Кеш делится на одинаковые блоки; таблица связывает номера логических блоков запроса с физическими блоками в памяти GPU. Ядро внимания умеет читать данные через это отображение. Для запроса сохраняется правильный порядок токенов, а система может брать свободные блоки из общего пула. Новый блок выделяется по мере заполнения предыдущего. После завершения запроса его блоки возвращаются в пул и могут обслуживать другие последовательности. В примере на слайдах блок содержит четыре токена для наглядности; в обсуждаемой реализации используется размер шестнадцать. Это различие важно при чтении схемы: рисунок объясняет механизм, а не задаёт настройку для любого сервера. Незаполненный остаток теперь ограничен последним блоком, вместо большого личного резерва на весь потенциальный ответ.

Следующий шаг — совместное использование уже рассчитанного префикса. Если кодинговому помощнику нужны несколько вариантов ответа на один промпт, нет смысла заранее копировать одинаковые KV-данные для каждого варианта. Их логические таблицы могут ссылаться на одни физические блоки. Пока данные только читаются, такого разделения достаточно. Когда одна последовательность хочет дописать токен в общий незаполненный блок, система проверяет число ссылок и создаёт отдельную копию для изменяемой ветки: это копирование при записи. Полностью заполненные общие блоки продолжают использоваться совместно. При поиске по нескольким кандидатам, beam search, возникает дерево префиксов: отброшенные ветви освобождают свои блоки, а общая часть сохраняется для оставшихся. Экономия зависит от того, насколько много в запросах общего текста. Само наличие таблицы блоков не означает одинакового выигрыша на любой нагрузке; оно даёт механизмы, которыми планировщик может воспользоваться.

Даже эффективное размещение не делает память бесконечной. Если активным запросам понадобились новые блоки, а свободных больше нет, планировщик должен кого-то приостановить. Здесь знание задачи позволяет упростить решение относительно универсальной операционной системы: очередному шагу внимания нужны все блоки последовательности, поэтому вытесняется целая группа, а не случайная страница. В рассмотренной системе выбирается позднее пришедшая группа, а порядок обслуживания помогает избегать голодания. Освободить память можно двумя способами: перенести KV-данные в память CPU и затем вернуть либо удалить их и восстановить повторной обработкой уже известной последовательности. Первый способ ограничен скоростью передачи, второй платит вычислениями. Ведущий связывает это с архитектурой vLLM: планировщик управляет запросами и таблицами, рабочие процессы на GPU исполняют вычисления. При тензорном параллелизме они используют согласованное отображение блоков для разных частей модели, а результаты вычислений синхронизируют отдельно.

В экспериментах разбираемой статьи авторы получили примерно двукратный–четырёхкратный рост пропускной способности относительно Orca при сопоставимой задержке. Это исторический результат для выбранных моделей и нагрузок. При этом само ядро PagedAttention может работать медленнее обычного ядра внимания: косвенный доступ и таблица блоков требуют дополнительной работы. Размер блока тоже приходится выбирать. Слишком маленький увеличивает накладные расходы, слишком большой возвращает потери в незаполненном остатке. Итоговый выигрыш возникает потому, что экономия памяти позволяет одновременно обслуживать больше последовательностей и лучше использовать дорогое оборудование. Поэтому замер отдельного ядра и замер сервиса могут привести к противоположным выводам. Финальная инженерная мысль выпуска — найти реальное ограничение, взять подходящую идею из другой области и проверить её влияние на всю систему. Изменять саму модель для такого улучшения не понадобилось. Отдельная граница применимости — обучение с более статичной организацией тензоров: там цена косвенного доступа может остаться, а выгода динамического размещения исчезнуть.