Как работает paged attention в vLLM? Чем это отличается от стандартного attention механизма?

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

Paged attention — это техника оптимизации инференса LLM, реализованная в библиотеке vLLM. Она решает проблему фрагментации и неэффективного использования памяти при хранении KV-кэша (Key-Value cache) в стандартном attention. Вместо выделения непрерывного блока памяти под весь кэш для каждого запроса, Attention|paged attention разбивает KV-кэш на блоки фиксированного размера (pages) и управляет ими как виртуальной памятью в операционной системе. Это позволяет выделять память по требованию, уменьшает фрагментацию на 70–80% и значительно увеличивает пропускную способность (throughput) инференса.

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

В стандартном attention-механизме (например, в трансформерах) при генерации каждого нового токена модель вычисляет Key и Value для всех предыдущих токенов в последовательности. Эти значения сохраняются в KV-кэше, чтобы избежать повторных вычислений. Размер KV-кэша растёт линейно с длиной последовательности и количеством запросов.

Основные проблемы стандартного подхода

• Непрерывное выделение памяти (contiguous allocation): Для каждого запроса выделяется непрерывный блок памяти под весь возможный KV-кэш (до максимальной длины последовательности). Это приводит к внутренней фрагментации, так как реальная длина ответа часто меньше максимальной.
• Внешняя фрагментация При параллельной обработке нескольких запросов (batching) пулы памяти разных запросов могут фрагментироваться, делая невозможным выделение большого непрерывного блока, даже если суммарно свободной памяти достаточно.
• Неэффективное использование GPU-памяти Из-за фрагментации до 60–80% памяти может быть занято, но не использовано.

Термин: KV-кэш — это структура данных, хранящая вычисленные Key и Value матрицы для всех предыдущих токенов в последовательности. Позволяет избежать пересчёта attention для уже обработанных токенов.

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

2. Как работает paged attention: идея и механизм

Paged attention заимствует концепцию страничной организации памяти (paging) из операционных систем. KV-кэш разбивается на блоки фиксированного размера — pages (обычно 16 или 32 токена). Каждый page — это непрерывный блок памяти, но сами pages могут располагаться в памяти непоследовательно.

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

• Logical KV-blocks (логические блоки): Виртуальное представление KV-кэша для каждого запроса. Логические блоки нумеруются последовательно (0, 1, 2, ...).
• Physical KV-blocks (физические блоки): Реальные блоки памяти на GPU. Они выделяются по мере необходимости и могут быть разбросаны по всей доступной памяти.
• Block table (таблица блоков): Словарь (или массив), который отображает логический номер блока на физический адрес. Для каждого запроса ведётся своя block table.

Процесс работы

1. Инициализация При старте генерации для запроса выделяется один логический блок (page 0). Физическая память ещё не выделена.
2. Генерация первого токена Модель вычисляет Key и Value для первого токена. Они записываются в физический блок, который выделяется по требованию. Block table обновляется: logical page 0 → physical page N.
3. Генерация следующих токенов По мере заполнения текущего физического блока (например, после 16 токенов) выделяется новый физический блок. Block table добавляет новую запись: logical page 1 → physical page M.
4. Attention с paging При вычислении attention для нового токена модель обращается к block table, чтобы получить физические адреса всех предыдущих блоков. Затем она выполняет attention, читая данные из разрозненных физических блоков, как если бы они были непрерывными.

Пример (упрощённый):

• Запрос A: logical pages [0, 1, 2] → physical pages [5, 2, 8]
• Запрос B: logical pages [0, 1] → physical pages [3, 7]

Физические блоки 5, 2, 8, 3, 7 могут быть разбросаны в памяти, но block table обеспечивает логическую непрерывность.

3. Отличия от стандартного attention механизма

Термин: Copy-on-write (COW) — техника, при которой несколько запросов могут разделять одни и те же физические блоки, пока один из них не попытается изменить данные. Это позволяет экономить память при обработке общих префиксов (например, системного промпта).

4. Преимущества paged attention

1. Уменьшение фрагментации памяти на 70-80% За счёт выделения блоков по требованию и непоследовательного хранения.
2. Увеличение пропускной способности (throughput): vLLM может обрабатывать больше запросов одновременно (больший batch size) при том же объёме GPU-памяти.
3. Поддержка copy-on-write Позволяет эффективно разделять KV-кэш для запросов с общим префиксом (например, в чат-ботах с системным промптом).
4. Гибкое управление памятью Блоки могут быть освобождены и переиспользованы сразу после завершения запроса, без ожидания освобождения всего непрерывного блока.
5. Поддержка переменной длины последовательностей Нет необходимости резервировать память под максимальную длину для каждого запроса.

5. Реализация в vLLM: ключевые компоненты

vLLM — это библиотека для высокопроизводительного инференса LLM. Paged attention — её ключевая инновация.

Основные компоненты реализации

• Scheduler (планировщик): Управляет выделением и освобождением физических блоков. Решает, какие запросы включить в текущий batch, основываясь на доступной памяти.
• Block manager (менеджер блоков): Ведёт учёт всех физических блоков, их состояния (свободен/занят) и block tables для каждого запроса.
• Custom CUDA kernel (пользовательское ядро CUDA): Реализует операцию attention, которая работает с разрозненными физическими блоками. Это ядро оптимизировано для чтения из непоследовательной памяти.

Пример упрощённого кода на Python (концептуально):

class PagedAttention:
    def __init__(self, block_size=16, num_physical_blocks=1000):
        self.block_size = block_size
        self.physical_blocks = [None] * num_physical_blocks  # None = свободен
        self.block_tables = {}  # request_id -> {logical_block: physical_block}

    def allocate_block(self):
        for i, block in enumerate(self.physical_blocks):
            if block is None:
                self.physical_blocks[i] = []  # инициализируем блок
                return i
        raise MemoryError("No free physical blocks")

    def add_token(self, request_id, key, value):
        if request_id not in self.block_tables:
            self.block_tables[request_id] = {}
            logical_block = 0
            physical_block = self.allocate_block()
            self.block_tables[request_id][logical_block] = physical_block
        else:
            # Найти последний логический блок
            logical_block = max(self.block_tables[request_id].keys())
            physical_block = self.block_tables[request_id][logical_block]
            if len(self.physical_blocks[physical_block]) >= self.block_size:
                # Блок заполнен, выделяем новый
                logical_block += 1
                physical_block = self.allocate_block()
                self.block_tables[request_id][logical_block] = physical_block
        # Добавляем key, value в физический блок
        self.physical_blocks[physical_block].append((key, value))

    def get_kv_cache(self, request_id):
        # Возвращает список (key, value) для всех токенов запроса
        cache = []
        for logical_block in sorted(self.block_tables[request_id].keys()):
            physical_block = self.block_tables[request_id][logical_block]
            cache.extend(self.physical_blocks[physical_block])
        return cache

Примечание: Реальная реализация в vLLM использует CUDA и работает с тензорами, а не списками Python.

6. Влияние на latency и throughput

Latency (задержка): Paged attention может незначительно увеличить latency для одного запроса из-за накладных расходов на управление block table и непоследовательное чтение памяти. Однако это увеличение обычно составляет менее 5%.

Throughput (пропускная способность): Главное преимущество. За счёт более эффективного использования памяти vLLM может обрабатывать в 2-4 раза больше запросов в секунду по сравнению с традиционными системами (например, Hugging Face Transformers) при том же оборудовании.

Пример сравнения (гипотетический):

7. Ограничения и недостатки

1. Сложность реализации Требует написания custom CUDA kernels, что усложняет поддержку и портирование на другие архитектуры.
2. Накладные расходы на управление Block table и планировщик добавляют вычислительные затраты, особенно при очень малых размерах блоков.
3. Неоптимально для очень коротких последовательностей Если длина ответа меньше размера одного блока, выигрыш в памяти минимален.
4. Зависимость от размера блока Слишком маленький блок увеличивает накладные расходы, слишком большой — снижает эффективность борьбы с фрагментацией.

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

Задача Реализовать упрощённую симуляцию paged attention на Python (без CUDA) и сравнить использование памяти со стандартным подходом.

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

Шаги:

1. Реализуйте класс StandardAttention Выделяет непрерывный массив под KV-кэш для каждого запроса (максимальная длина — 1024 токена). Симулируйте обработку 100 запросов со случайной длиной ответа (от 50 до 500 токенов). Замерьте общее использование памяти.
2. Реализуйте класс PagedAttention Используйте блоки по 16 токенов. Выделяйте физические блоки по мере необходимости. Замерьте общее использование памяти.
3. Сравните результаты Постройте график зависимости использования памяти от количества запросов. Рассчитайте процент фрагментации для стандартного подхода.
4. Добавьте copy-on-write Реализуйте разделение блоков для запросов с общим префиксом (первые 50 токенов одинаковы). Замерьте экономию памяти.

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

• Paged attention покажет на 60-80% меньшее использование памяти.
• Copy-on-write даст дополнительную экономию при наличии общих префиксов.
• Вы получите практическое понимание механизма и его преимуществ.

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

10. Навигация

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

Навигация

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