Что такое PagedAttention: размещение KV cache

Что такое PagedAttention: размещение KV cache блоками с обращением к ним через таблицу. кэш делится на блоки фиксированного размера, а таблица блоков связывает логические позиции с физическими блоками. Так описывает механизм блог vLLM. Дальше — как устроено размещение, где оно помогает и где упирается в предел.

Зачем KV cache вообще нужен

В рассматриваемом полном причинном внимании модель генерирует токены по одному, используя ключи и значения предыдущих позиций. Пересчитывать их на каждом шаге дорого, поэтому движок вывода складывает их в KV cache. В документации Hugging Face для Text Generation Inference сказано, что при декодировании все ключи и значения предыдущих токенов хранятся в памяти GPU для повторного использования, и для больших моделей и длинных последовательностей кэш занимает много памяти.

В справочнике Modular описано, как это выглядит в ранних движках: KV cache выделялся одним сплошным тензором под худший случай, размером под максимальную длину последовательности. Выделение шло на каждый активный запрос. Трафик переменчив: один запрос даёт короткий ответ, другой тянет долгую беседу, третий обрывается рано. Если каждый запрос резервирует память под максимальную длину, значительная часть зарезервированной памяти GPU простаивает.

В статье arXiv 2309.06180 авторы фиксируют две причины неэффективности: кэш динамически растёт и сокращается по мере генерации, а его время жизни и длина заранее неизвестны; кроме того, сплошные раздельные области не дают делиться памятью между последовательностями одного запроса. В исходной публикации блог vLLM оценивал потери рассмотренных тогда систем в 60%–80% памяти на фрагментацию и избыточное резервирование.

Размер кэша зависит от длины последовательности, а она непредсказуема. Это и есть задача, которую решает блочное размещение.

Механизм: блоки, логические позиции, таблица

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

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

Документация Hugging Face формулирует то же короче: кэш делится на блоки, доступ к ним идёт через таблицу поиска, блоки выделяются по необходимости. Справочник Modular добавляет, что модель загружает только нужные блоки.

Разберём учебную карту. Блок вмещает 4 токена. Запрос А длиной 6: логический блок 0 → физический 7 [А0, А1, А2, А3]; логический блок 1 → физический 2 [А4, А5, пусто, пусто]. Запрос Б длиной 3: логический блок 0 → физический 9 [Б0, Б1, Б2, пусто]. В выделенных блоках занято 9 мест из 12, три места — внутренний остаток. Это условная учебная схема. Вся память GPU здесь не рассчитывается.

Учебная трасса: рост последовательности

Продолжим пример. Токены А6 и А7 заполняют физический блок 2, для А8 выделяется блок 4. Таблица А становится [7, 2, 4]. После завершения запроса Б освобождается блок 9. Читатель восстанавливает порядок А0…А8 по логической таблице, а физические номера 7, 2, 4 идут вразнобой. Запуска GPU в этом разборе не было.

Проверка результата строится на трёх вопросах. Первый: совпадает ли число логических блоков с числом занятых физических. Второй: покрывает ли таблица все токены последовательности без пропусков. Третий: освободился ли блок завершённого запроса и вернулся ли он в пул.

Отдельно разведём структурную корректность и фактическую правильность. Структурная корректность — таблица ссылается на существующие блоки, счётчики ссылок сходятся, освобождение не оставляет висячих записей. Фактическая правильность — значения ключей и значений в блоках соответствуют тем токенам, которые в них записаны. Таблица может быть безупречной, а содержимое блока — от другого запроса. Это разные проверки, и провал второй не виден по первой.

Ещё один слой — разделение неизменного префикса. При совместимых модели и состоянии одинаковый неизменный префикс может храниться в общих блоках для нескольких последовательностей. Как только одна из них меняет содержимое страницы, нужна корректная обработка изменения: блог vLLM называет это механизмом Copy-on-Write со счётчиками ссылок на физические блоки.

Что даёт блочное размещение на практике

Блог vLLM сообщает, что при PagedAttention память тратится впустую только в последнем блоке последовательности, и на практике потери составляют менее 4%. Справочник Modular приводит цифры из оригинальной работы: без PagedAttention на хранение реальных состояний токенов шло 20,4%–38,2% выделенной памяти кэша, остальное терялось на фрагментацию; с PagedAttention потери близки к нулю.

Экономия памяти позволяет держать больше последовательностей в одном батче и поднимать загрузку GPU. В блоге vLLM приведены замеры на LLaMA-7B с NVIDIA A10G и LLaMA-13B с NVIDIA A100 (40GB), длины входов и выходов брались из датасета ShareGPT: до 24x выше пропускной способности относительно HuggingFace Transformers и до 3,5x относительно Text Generation Inference. Это результаты конкретных конфигураций, и переносить их на другой сервер не стоит.

Второй эффект — совместное использование памяти. В параллельной выборке несколько выходных последовательностей растут из одного промпта, и вычисления с памятью для промпта делятся между ними. По данным блога vLLM, такое разделение снижает расход памяти сложных схем сэмплирования до 55% и даёт до 2,2x прироста пропускной способности.

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

Где подход не решает задачу

Блочное размещение отвечает за управление памятью. Его основная задача — организация доступа к KV-блокам и управление их размещением. Это не уменьшает математическое число связей плотного внимания. Оптимизация вычислительного ядра, например FlashAttention, решает смежную задачу и может сочетаться с блочным хранением.

Оно не устраняет предел, заданный объёмом памяти GPU. Когда суммарный кэш активных запросов упирается в потолок, рост числа одновременных последовательностей останавливается, и блочная раскладка этот потолок не поднимает. Учебная карта с 9 занятыми местами из 12 показывает остаток внутри блоков; в реальной работе к нему добавляются накладные расходы на таблицы и счётчики ссылок, которые в схеме не отражены.

Оно не даёт выигрыша там, где кэш мал. При коротких последовательностях и небольшом числе одновременных запросов экономия от блочного размещения мала, а таблица добавляет косвенность при доступе к блокам.

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

Как читать упоминание PagedAttention в описании сервера

Первое: уточните, о чём именно речь — о ядре внимания или о слое управления памятью. Справочник Modular разводит эти вещи.

Второе: посмотрите, как в продукте назван размер блока и чем он задаётся. Справочник Modular отмечает, что vLLM реализовал подход первым, а затем его переняли другие проекты, включая MAX, и выставляют размер блока флагом сервера под разными именами. Документация TensorRT LLM говорит, что число токенов в блоке задаёт пользователь при создании движка модели, и оно должно быть степенью двойки больше единицы.

Третье: проверьте, какие механизмы переиспользования включены. В TensorRT LLM блоки с посчитанным состоянием складываются в дерево поиска по мере заполнения, и при добавлении нового запроса выполняется поиск: совпавшие блоки переиспользуются вместо вычисления. Вытеснение идёт по приоритетам от 0 до 100, и при равных приоритетах вытесняется наименее недавно использованный блок. При вытеснении из основной памяти состояние копируется в блок вторичной памяти.

Четвёртое: держите в голове, что цифры ускорения привязаны к конфигурации. Замеры блога vLLM сделаны на конкретных моделях и GPU с длинами из ShareGPT.

Промпт, чтобы разобрать чужую схему размещения кэша по шагам:

Ты помогаешь разобрать схему размещения KV cache в LLM-сервере.
Дано: [вставьте описание или фрагмент документации].
Сделай по шагам:
1. Назови единицу хранения и её размер, если он указан.
2. Опиши, как логические позиции последовательности связаны с физическими блоками.
3. Укажи, когда выделяется новый блок и когда блок освобождается.
4. Перечисли механизмы переиспользования: общий префикс, счётчики ссылок, вытеснение.
5. Отдели управление памятью от вычисления внимания.
6. Перечисли ограничения, названные для этой реализации.
Формат ответа: таблица «сущность — смысл — где в тексте сказано».
Если данных для пункта нет, напиши «в приведённом тексте не указано» и переходи дальше.

Меняйте под себя фрагмент в квадратных скобках и добавляйте пункты про интересующий продукт. Вариации той же задачи: разбор таблицы блоков из статьи arXiv; сравнение двух движков по названию флага размера блока; проверка, как реализовано освобождение блока после завершения запроса; разбор схемы Copy-on-Write при общем префиксе; составление вопросов к документации перед выбором движка.

Частые вопросы

Что такое PagedAttention простыми словами?

Это способ хранить кэш ключей и значений не одним сплошным куском, а наборами блоков фиксированного размера. Таблица блоков связывает позиции внутри последовательности с местами в памяти. Так устроен механизм по описанию блога vLLM и документации Hugging Face.

Чем PagedAttention отличается от обычного внимания?

В исходном сравнении использовалось сплошное резервирование KV-кэша. Это способ хранения, а сама операция внимания допускает разные реализации размещения. PagedAttention допускает несплошное размещение и подтягивает нужные блоки через таблицу. Справочник Modular относит основной выигрыш к слою управления памятью движка.

PagedAttention ускоряет саму операцию внимания?

Основной эффект рассматриваемого подхода связан с размещением и совместным использованием памяти. Справочник Modular указывает, что за эффективность внимания на уровне ядра отвечает FlashAttention. Ускорение вывода растёт за счёт того, что в память помещается больше одновременных последовательностей.

Какой размер блока используется?

Единого значения нет. Документация TensorRT LLM говорит, что число токенов в блоке задаёт пользователь при создании движка модели, и оно должно быть степенью двойки больше единицы. Справочник Modular отмечает, что разные проекты выставляют размер блока флагом сервера под своими именами.

Где взять исходный код и статью?

Блог vLLM ссылается на репозиторий проекта на GitHub, а технические детали вынесены в отдельную работу. Оригинальная статья доступна на arXiv по адресу arxiv.org/html/2309.06180v1. Документация Hugging Face отмечает, что реализация в TGI использует ядра CUDA, разработанные проектом vLLM.