Как работает PagedAttention в vLLM внутренне?

Краткий тезис

PagedAttention — это техника управления KV-кэшем в LLM, вдохновлённая виртуальной памятью в ОС. Вместо того чтобы выделять непрерывный блок памяти под весь KV-кэш для последовательности, PagedAttention разбивает его на блоки фиксированного размера (обычно 16 токенов) и хранит их непоследовательно через таблицу страниц. Это позволяет снизить фрагментацию памяти с ~70% до ~5%, увеличить throughput и поддерживать эффективную работу с длинными последовательностями и shared prefix'ами.

1. Термин: KV-кэш (Key-Value Cache)

Что это В архитектуре Transformer при генерации каждого нового токена модель вычисляет Key и Value матрицы для всех предыдущих токенов. Чтобы не пересчитывать их заново, они кэшируются в памяти — это и есть KV-кэш.

Проблема KV-кэш растёт линейно с длиной последовательности и batch size. Для больших моделей (например, 70B параметров) он может занимать десятки гигабайт. Традиционные аллокаторы выделяют непрерывный блок памяти под каждый запрос, что приводит к сильной фрагментации.

Термин «Фрагментация памяти» — ситуация, когда память разбита на множество маленьких свободных участков, которые нельзя объединить в один большой блок, достаточный для нового запроса. В контексте KV-кэша это означает, что мы не можем эффективно использовать доступную память.

2. Проблема: Фрагментация и неэффективность традиционного подхода

Традиционный подход: выделить непрерывный блок памяти размером max_seq_len * num_layers * num_heads * head_dim * dtype_size для каждого запроса в batch.

• Внутренняя фрагментация Если реальная длина последовательности меньше max_seq_len, мы выделяем память «про запас», которая не используется.
• Внешняя фрагментация После обработки нескольких запросов разной длины память разбивается на фрагменты, и новый запрос может не поместиться, даже если суммарно свободной памяти достаточно.
• Ограничение batch size Из-за фрагментации мы не можем набрать большой batch, что снижает throughput (пропускную способность).

Статистика В традиционных системах (например, Hugging Face Transformers) фрагментация KV-кэша может достигать 60–70%.

3. Идея PagedAttention: Виртуальная память для KV-кэша

PagedAttention решает проблему, заимствуя концепцию таблицы страниц (page table) из операционных систем.

Аналогия с ОС

• Виртуальная память Процесс видит непрерывное виртуальное адресное пространство, но физические страницы памяти разбросаны.
• PagedAttention Модель «видит» непрерывный KV-кэш для последовательности (логические блоки), но физически блоки хранятся в разных местах (физические блоки).

Ключевые компоненты

1. Логические блоки (Logical Blocks): Последовательность токенов разбивается на блоки фиксированного размера B (обычно 16 токенов). Каждый блок имеет логический индекс.
2. Физические блоки (Physical Blocks): Реальные блоки памяти, выделенные на GPU. Они не обязаны быть непрерывными.
3. Таблица страниц (Block Table / Page Table): Словарь, который отображает логический индекс блока на физический адрес. Для каждого запроса хранится своя таблица страниц.

4. Как работает PagedAttention: Пошагово

Шаг 1: Инициализация

• Приходит новый запрос (промпт).
• Система резервирует место для логических блоков, но не выделяет физическую память сразу.
• Создаётся пустая таблица страниц для этого запроса.

Шаг 2: Prefill (Prefill-фаза)

• Промпт обрабатывается целиком. Для каждого токена вычисляются Key и Value.
• По мере обработки, когда набирается B токенов, система запрашивает новый физический блок из менеджера памяти (memory manager).
• Адрес этого блока записывается в таблицу страниц.

Шаг 3: Decode (Decode-фаза)

• Генерируется один новый токен.
• Его Key и Value добавляются в текущий логический блок.
• Если блок заполнен, запрашивается новый физический блок.

Шаг 4: Attention с таблицей страниц

• При вычислении attention для нового токена, модель должна получить доступ ко всем предыдущим Key и Value.
• Вместо того чтобы обращаться к непрерывному массиву, модель использует таблицу страниц:
  1. Определяет, к какому логическому блоку относится нужный токен.
  2. По таблице страниц находит физический адрес этого блока.
  3. Загружает данные из физического блока.
• Это называется непрямым доступом (indirect access).

Шаг 5: Освобождение памяти

• Когда запрос завершён, все его физические блоки возвращаются в пул свободной памяти.
• Таблица страниц удаляется.

5. Преимущества PagedAttention

6. Дополнительные возможности: Shared Prefix и Copy-on-Write

Shared Prefix (Общий префикс)

• Если несколько запросов начинаются с одинакового текста (например, системный промпт), их первые логические блоки могут ссылаться на одни и те же физические блоки.
• Это экономит память, особенно в чат-ботах с длинными системными промптами.

Copy-on-Write (Копирование при записи)

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

7. Реализация в vLLM

vLLM — это библиотека для высокопроизводительного инференса LLM, которая первой реализовала PagedAttention.

Ключевые компоненты vLLM

• Scheduler (Планировщик): Управляет очередью запросов, решает, какие запросы добавить в batch.
• Block Manager (Менеджер блоков): Отвечает за выделение и освобождение физических блоков.
• Worker (Воркер): Выполняет вычисления attention на GPU, используя таблицу страниц.

Псевдокод (упрощённо):

class PagedAttention:
    def __init__(self, block_size=16):
        self.block_size = block_size
        self.physical_blocks = {}  # {block_id: tensor}
        self.free_blocks = set()   # пул свободных блоков

    def allocate_block(self):
        # Берём блок из пула или создаём новый
        block_id = self.free_blocks.pop()
        self.physical_blocks[block_id] = torch.zeros(...)
        return block_id

    def process_request(self, request):
        page_table = {}  # {logical_block_id: physical_block_id}
        logical_block_id = 0
        for token in request.tokens:
            # Добавляем токен в текущий логический блок
            if len(current_logical_block) == self.block_size:
                # Блок заполнен, выделяем новый физический
                phys_id = self.allocate_block()
                page_table[logical_block_id] = phys_id
                logical_block_id += 1
                current_logical_block = []
            current_logical_block.append(token)
        # ... вычисление attention с page_table

8. Влияние на производительность

• Throughput (пропускная способность): vLLM показывает прирост в 2–4x по сравнению с Hugging Face Transformers для batch-инференса.
• Latency (задержка): Для одиночных запросов latency может быть немного выше из-за накладных расходов на таблицу страниц, но для batch-запросов она снижается.
• Memory footprint Значительно меньше, особенно для длинных последовательностей.

9. Ограничения и компромиссы

• Накладные расходы на таблицу страниц Для очень коротких последовательностей (1–2 блока) overhead может быть заметен.
• Block size Выбор размера блока — компромисс. Маленький блок (8 токенов) → меньше фрагментации, но больше overhead на таблицу. Большой блок (32 токена) → меньше overhead, но больше фрагментации. Оптимальный размер обычно 16 токенов.
• Сложность реализации Требует кастомного CUDA-кода для attention с непрямым доступом.

Пет-проект для закрепления

Задача Реализовать упрощённую версию PagedAttention на Python с использованием PyTorch.

Инструменты Python, PyTorch, NumPy.

Шаги:

1. Создайте класс PagedAttentionManager, который управляет пулом физических блоков.
2. Реализуйте allocate_block() и free_block().
3. Напишите функцию simulate_inference(prompt_tokens, model, manager), которая:
   • Разбивает промпт на логические блоки.
   • Выделяет физические блоки по мере необходимости.
   • Вычисляет attention, используя таблицу страниц (для простоты можно использовать обычный dot-product attention, но с непрямым доступом к памяти).
4. Сравните использование памяти с традиционным подходом (выделение непрерывного блока под max_seq_len).
5. Добавьте поддержку shared prefix: два запроса с общим началом должны использовать одни и те же физические блоки для первых токенов.

Ожидаемый результат

• Вы увидите, что при традиционном подходе фрагментация растёт с увеличением batch size и разбросом длин последовательностей.
• PagedAttention покажет гораздо более эффективное использование памяти, особенно для длинных последовательностей.
• Shared prefix позволит обрабатывать несколько запросов с общим началом с минимальным потреблением памяти.

Связь с другими вопросами

Навигация

• Предыдущий: 438
• Следующий: 440
• Индекс: 00. Индекс разборов