PagedAttention

PagedAttention — это алгоритм внимания и менеджер памяти для обслуживания больших языковых моделей (LLM). Он хранит KV-кэш каждого запроса в небольших блоках фиксированного размера, размещённых в любом месте памяти GPU, и отслеживает их в таблице блоков для каждого запроса — подобно тому, как операционная система сопоставляет страницы виртуальной памяти. Вуосук Квон, Чжуохань Ли и их коллеги из UC Berkeley и университетов-партнёров представили этот метод вместе с движком обслуживания моделей vLLM в статье Efficient Memory Management for Large Language Model Serving with PagedAttention, представленной на SOSP 2023. Устранение фрагментации, вызванной последовательным выделением кэша, позволяет одному GPU одновременно хранить гораздо больше последовательностей. В статье сообщается о пропускной способности в 2–4 раза выше, чем у FasterTransformer и Orca при той же задержке.

Проблема с памятью, которую решает метод#

Во время генерации текста трансформер хранит ключевые и векторные значения каждого предыдущего токена, чтобы не пересчитывать их. Кэш увеличивается на одну запись для каждого токена в каждом слое, и одна последовательность для модели LLaMA с 13 миллиардами параметров может занимать до 1,7 ГБ. Длина выходных данных неизвестна в момент поступления запроса, поэтому ранние системы обслуживания моделей резервировали для каждого запроса одну непрерывную область памяти размером, рассчитанным на максимально возможную длину. Большая часть этой области оставалась пустой, а промежутки между областями были слишком малы для повторного использования. Команда vLLM измерила, что такие системы расходовали впустую 60–80% памяти KV-кэша, из-за чего размер пакета ограничивался задолго до исчерпания вычислительных ресурсов.

Как работает страничное распределение#

PagedAttention разбивает кэш на блоки, каждый из которых хранит ключи и значения фиксированного числа токенов; в vLLM по умолчанию это 16 токенов. Новый блок выделяется запросу только после заполнения предыдущего и берётся из любого свободного места, а таблица блоков фиксирует их порядок. Ядро механизма внимания обращается к этой таблице, чтобы собирать ключи и значения блок за блоком, поэтому кэш последовательности не нужно хранить непрерывно. Потери ограничиваются незаполненным остатком последнего блока каждой последовательности и, по измерениям vLLM, составляют менее 4%.

В этом примере на PyTorch ключи одного запроса размещаются в разрозненных блоках общего пула, а затем считываются с помощью таблицы блоков:

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

Выводятся таблица блоков [41, 7, 19] и keys intact: True, а затем — 8 неиспользуемых ячеек для запроса с постраничным распределением против 2 008 для области, зарезервированной на 2 048 токенов.

Совместное использование блоков последовательностями#

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

Авторы vLLM сообщают, что такое совместное использование сокращает расход памяти при параллельной генерации и поиске по лучу до 55%, что позволяет увеличить пропускную способность до 2,2 раза. Хеширование полных блоков распространяет этот подход на разные запросы — именно так vLLM реализует кэширование промптов.

Родственные методы#

PagedAttention определяет место хранения кэшированных ключей и значений и сочетается с методами, которые сокращают объём сохраняемых данных или ускоряют их чтение:

• Flash Attention: разбивает вычисление внимания на блоки, поэтому полная матрица оценок никогда не записывается в основную память GPU. В современных движках используются ядра типа FlashAttention, которые считывают данные из постранично распределённых блоков.
• Непрерывное пакетирование: новый запрос добавляется сразу после завершения другого. Постраничное распределение сразу возвращает блоки завершившейся последовательности в свободный пул, освобождая место для нового запроса.
• Внимание с группировкой запросов: одни и те же головы ключей и значений используются несколькими головами запросов, что уменьшает содержимое каждого блока.
• Квантование кэша FP8: ключи и значения хранятся с точностью 8 бит, что примерно вдвое сокращает объём памяти кэша по сравнению с 16-битными форматами.

Интерфейс LLM Ultralytics передаёт аргументы запроса без изменений. Поэтому вызов Chat Completions с n=3, отправленный через base_url на локальный сервер vLLM, возвращает три варианта ответа, которые совместно используют кэшированные блоки промпта; только различающиеся продолжения занимают новую память.