Как работает paged attention? (детально)

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

PagedAttention — это техника управления памятью для KV-кэша в больших языковых моделях, вдохновлённая страничной памятью операционных систем. Кэш разбивается на блоки фиксированного размера (например, 16 токенов), которые аллоцируются не последовательно, а через таблицу страниц. Это решает проблему фрагментации памяти (до 70% в наивных реализациях) и позволяет обслуживать больше одновременных запросов, особенно при длинных генерациях и пакетной обработке. PagedAttention является ключевым нововведением движка vLLM.

1. Проблема: фрагментация KV-кэша в традиционных подходах

При авторегрессивной генерации каждый новый токен требует доступа к KV-кэшу всех предыдущих токенов. Для каждого слоя модели и каждой головы внимания хранятся два вектора: Key (K) и Value (V). При пакетной обработке запросы имеют разную длину — как входную, так и генерируемую. Если выделять непрерывную память под каждый запрос (например, по максимуму контекстного окна), то большую часть времени она простаивает или фрагментируется.

Наивное управление памятью:

• Каждому запросу выделяется непрерывный блок размером sequence length|max_seq_len * num_layers * num_heads * d_head * 2 (для K и V).
• Если запрос заканчивается раньше, хвост блока остаётся неиспользованным (внутренняя фрагментация).
• При освобождении одного запроса и аллокации для другого между ними могут образовываться неиспользуемые промежутки (внешняя фрагментация).

Результат: фрагментация памяти может достигать 60–80%, что резко снижает эффективный throughput.

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

2. Основная идея PagedAttention

PagedAttention заимствует концепцию страничной памяти из ОС: вместо одного непрерывного буфера под кэш для каждого запроса память разбивается на физические блоки (pages) одинакового размера (например, 16 токенов). Логическое представление кэша каждого запроса описывается таблицей страниц (block table), где перечислены номера физических блоков, хранящих его KV-кэш. Блоки могут располагаться в памяти не последовательно — точно так же, как виртуальные страницы процесса маппятся на произвольные физические фреймы.

Основные компоненты:

• Логические блоки — виртуальные слоты для каждого запроса (например, 1-й логический блок включает токены 0..15, 2-й — 16..31 и т.д.).
• Физические блоки — реальные участки памяти фиксированного размера. Количество физических блоков ограничено доступной памятью (например, на GPU).
• Block table (таблица страниц) — маппинг от (request_id, logical_block_id) к physical_block_id. Для каждого запроса свой маппинг.
• Block manager — компонент, который выделяет, освобождает и вытесняет физические блоки.

Преимущества:

• Внутренняя фрагментация не превышает размера последнего блока (в среднем полблока). При размере блока 16 токенов и средней длине последовательности 100 токенов фрагментация < 5%.
• Внешняя фрагментация практически отсутствует, так как освобождённый блок сразу может быть повторно использован любым другим запросом.
• Возможен общий доступ к одним и тем же физическим блокам для разных последовательностей (например, beam search или parallel sampling) — copy-on-write.

3. Структуры данных: Block Table и типы блоков

3.1 Физические блоки

Каждый физический блок хранит KV-кэш для всех слоёв и всех голов внимания для набора токенов, попадающих в этот блок. Размер блока конфигурируется (обычно 16, 32 или 64). Например:

• Размер слоя: num_layers * num_heads * d_head * 2 = 80 слоёв * 32 головы * 128 * 2 ≈ 2.5 MB.
• Блок на 16 токенов: ~40 MB на один блок. Вся память GPU делится на такие блоки.

3.2 Block Table (таблица страниц)

Для каждого запроса хранится массив логических идентификаторов блоков. Block manager ведёт глобальную хеш‑таблицу:

physical_block_id -> { 
    ref_count: int, 
    allocated: bool, 
    last_access_time: timestamp 
}

Пример упрощённой структуры на Python:

class PhysicalBlock:
    def __init__(self, block_id, block_size, num_layers, num_heads, d_head):
        self.block_id = block_id
        self.ref_count = 0
        # реальная память (K, V) аллоцируется отдельно, здесь опущена
        self.kv_cache = None  # torch.Tensor shape: (num_layers, 2, num_heads, block_size, d_head)

class BlockManager:
    def __init__(self, num_blocks, block_size, num_layers, num_heads, d_head):
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.blocks = [PhysicalBlock(i, block_size, num_layers, num_heads, d_head) 
                       for i in range(num_blocks)]
        self.free_blocks = list(range(num_blocks))  # очередь свободных ID
        # для каждого request_id храним список физических ID
        self.alloc_table = {}  # request_id -> list[physical_block_id]

    def allocate_blocks(self, request_id, num_blocks):
        # выделяет num_blocks свободных блоков
        allocated = []
        for _ in range(num_blocks):
            if not self.free_blocks:
                raise MemoryError("Out of memory")
            phys_id = self.free_blocks.pop(0)
            self.blocks[phys_id].ref_count = 1
            allocated.append(phys_id)
        self.alloc_table[request_id] = allocated
        return allocated

    def free_request(self, request_id):
        for phys_id in self.alloc_table.get(request_id, []):
            self.blocks[phys_id].ref_count -= 1
            if self.blocks[phys_id].ref_count == 0:
                self.free_blocks.append(phys_id)
        del self.alloc_table[request_id]

3.3 Copy-on-Write (COW) при совместном доступе

При beam search или parallel sampling несколько последовательностей могут иметь общие префиксы (первые токены). В PagedAttention такие последовательности могут ссылаться на одни и те же физические блоки до тех пор, пока не потребуется модификация (например, при генерации разного следующего токена). При записи в такой блок создаётся копия — это снижает потребление памяти для общих префиксов.

4. Операции: аллокация, деаллокация, preemption

4.1 Allocation (аллокация)

При начале генерации нового токена для запроса block manager выделяет новый физический блок, если текущий логический блок заполнен. Процесс:

1. Текущая длина последовательности (позиция pos).
2. Вычисляется логический номер блока: logical_block = pos // block_size.
3. Если этот логический блок ещё не маппится на физический (т.е. block_table[logical_block] пусто), то запрашивается новый физический блок из списка свободных.
4. Новый физический блок фиксируется в таблице, и ref_count устанавливается в 1.
5. KV-кэш для текущего токена записывается в соответствующую позицию внутри физического блока.

4.2 Deallocation (деаллокация)

Когда запрос завершён (сгенерирован EOS или достигнут лимит длины), все принадлежащие ему физические блоки освобождаются: их ref_count уменьшается на 1. Если ref_count становится 0, блок возвращается в пул свободных.

4.3 Preemption (вытеснение)

При нехватке свободных физических блоков block manager может вытеснить один или несколько запросов. Вытеснение заключается в:

1. Выборе кандидата (обычно по алгоритму FCFS — первый выполненный, или по приоритету).
2. Сохранении его KV-кэша на CPU (с возможной компрессией).
3. Освобождении его физических блоков.
4. Когда запрос снова становится активным, его блоки восстанавливаются из CPU (или пересчитываются заново, если кэш был выгружен).

Термин Preemption — механизм временной приостановки одного запроса для освобождения ресурсов, необходимых для другого.

Виды вытеснения:

• Swap-based: блоки выгружаются на CPU (в paged memory OS‑стиль).
• Recomputation-based: блоки удаляются, и при возобновлении запроса кэш пересчитывается с начала. Последнее проще, но дороже.

5. Pre-computed block table и оптимизация вычислений

5.1 Как attention вычисляется с paged‑блоками?

В стандартном attention для последовательности длины L вычисляется матрица A = Q * K^T. При paged организации K хранится в разрозненных блоках. Attention выполняется по блокам: для каждого физического блока загружается его часть K и V, вычисляются локальные scores, затем маскируются (для каузального внимания — только токены из предыдущих блоков и в пределах текущего блока). Результат агрегируется.

Это напоминает block-sparse attention, но без потери качества, так как все блоки присутствуют, хоть и хранятся непоследовательно.

5.2 Overhead

• Дополнительный расход GPU времени на работу с block table и загрузку блоков из произвольных адресов (непрерывное чтение быстрее). Однако этот overhead компенсируется резким снижением фрагментации и большей утилизацией памяти (больше запросов → выше throughput).
• Современные GPU (с архитектурой Hopper/Ampere) поддерживают asynchronous copy и многопоточность, что смягчает задержки.

6. Сравнение PagedAttention с другими подходами

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

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

• Scheduler — распределяет время GPU между запросами, принимает решения о preemption.
• BlockManager — управляет физическими блоками (free pool, block table, ref count).
• Attention backend — реализует PagedAttention на CUDA ядрах (использует custom kernels от NVIDIA).
• Engine — оркестрирует инференс: prefilling блоками, decode с дозаполнением блоков.

Процесс в vLLM:

1. Запрос приходит, scheduler выделяет начальный физический блок (prefill этап).
2. На каждом decode‑шаге scheduler добавляет новый блок, если необходимо.
3. Если свободных блоков нет → scheduler вытесняет самый старый запрос.
4. После завершения запроса блоки возвращаются.

Поддержка prefix caching: если несколько запросов имеют одинаковый префикс (system prompt), vLLM может повторно использовать физические блоки, соответствующие этому префиксу (global block table).

8. Влияние на throughput и латентность

Throughput (количество токенов в секунду) при PagedAttention может быть в 2–5 раз выше, чем у наивного конвейерного инференса, за счёт:

• большего batch size в памяти;
• эффективной утилизации GPU при длинных запросах;
• поддержки beam search без копирования всего кэша.

Латентность (время до первого токена) незначительно страдает из-за overhead таблицы страниц, но в целом остаётся низкой за счёт оптимизации CUDA ядер.

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

Задача: реализовать упрощённый симулятор PagedAttention на Python (без реального GPU), демонстрирующий выигрыш от страничной организации.

Инструменты: Python 3.9+, numpy, time.

Шаги:

1. Создайте класс SimpleBlockManager с пулом физических блоков (блок = numpy.array для хранения значений K и V).
2. Реализуйте allocate_sequence(seq_len) — возвращает список физических блоков, аллоцированных для этой последовательности.
3. Симулируйте 100 запросов с разными длинами (случайное распределение от 50 до 500 токенов). Подсчитайте общее количество занятых блоков.
4. Сравните с наивным подходом: выделение непрерывного массива под max_len для каждого запроса (max_len = 512). Вычислите фрагментацию.
5. Добавьте preemption: если блоков не хватает, выгрузите самый старый запрос на CPU (в файл).
6. Выведите гистограмму утилизации памяти.

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

• Фрагментация в paged‑версии ≤5% (при block_size=16).
• Количество одновременно обслуживаемых запросов в paged‑версии в 3–5 раз больше.
• Программа выводит таблицу сравнения.

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

Навигация

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