Что такое continuous batching? Как реализовано в vLLM?

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

batching|Continuous batching (batching|непрерывная пакетная обработка) — это техника инференса LLM, при которой батч запросов динамически обновляется на каждой итерации генерации токенов. В отличие от статического батчинга, где все запросы начинают и заканчивают генерацию одновременно, batching|continuous batching позволяет удалять завершившиеся запросы и добавлять новые на каждом шаге декодирования. В vLLM это реализовано через iteration-level scheduling и механизм PagedAttention, что позволяет достичь прироста пропускной способности (throughput) в 4-6 раз по сравнению со статическим батчингом.

1. Термины: Batching, Static Batching и Continuous Batching

Batching (пакетная обработка) — объединение нескольких запросов в один батч для параллельной обработки на GPU. Это ключевой приём для повышения утилизации GPU.

Static batching (статическая пакетная обработка) — классический подход, при котором батч формируется до начала генерации. Все запросы в батче обрабатываются синхронно: они начинают генерацию вместе и заканчивают вместе. Проблема в том, что разные запросы генерируют разное количество токенов (например, один запрос требует 10 токенов, а другой — 100). В статическом батчинге все запросы ждут, пока самый длинный не завершится, что приводит к простою GPU (так называемая bubble).

Continuous batching (непрерывная пакетная обработка) — динамический подход, при котором батч обновляется на каждой итерации генерации одного токена. Запросы, которые сгенерировали токен <EOS> (конец последовательности), удаляются из батча, а новые запросы из очереди могут быть добавлены. Это устраняет проблему простоя GPU.

2. Проблема статического батчинга: Bubble и неэффективность

Представьте, что у вас есть GPU, который может обрабатывать батч из 4 запросов. Вы формируете статический батч из 4 запросов:

• Запрос A: генерирует 5 токенов
• Запрос B: генерирует 50 токенов
• Запрос C: генерирует 10 токенов
• Запрос D: генерирует 100 токенов

При статическом батчинге все запросы будут обрабатываться до тех пор, пока не завершится самый длинный (запрос D на 100 токенах). Это означает, что GPU будет простаивать на 95 итерациях для запроса A, на 50 для запроса B и на 90 для запроса C. Эффективность использования GPU резко падает.

Bubble (пузырь) — это время, когда GPU обрабатывает запросы, которые уже сгенерировали свой ответ, но вынуждены ждать завершения других запросов в батче.

3. Continuous Batching: Как это работает

Continuous batching решает проблему bubble за счёт динамического управления батчем на каждой итерации.

Алгоритм работы

1. Очередь запросов (Scheduling Queue): Все входящие запросы помещаются в очередь.
2. Итерация генерации: На каждом шаге декодирования (генерации одного токена) scheduler решает, какие запросы будут в батче.
3. Удаление завершённых: Если запрос сгенерировал токен <EOS> (конец последовательности) или достиг максимальной длины, он удаляется из батча.
4. Добавление новых: Если в батче освободилось место (из-за удаления завершённых запросов), scheduler может добавить новые запросы из очереди.
5. Prefill vs Decode: Новые запросы проходят этап prefill (предварительное заполнение кэша ключей-значений для входного промпта), а затем переключаются на этап decode (генерация токенов).

Пример:

• Итерация 1: Батч = [A, B, C, D]. Все генерируют первый токен.
• Итерация 2: Батч = [A, B, C, D]. A генерирует второй токен, B — второй, C — второй, D — второй.
• ...
• Итерация 5: A генерирует токен <EOS>. A удаляется из батча. Scheduler добавляет новый запрос E из очереди. E проходит prefill. Батч = [B, C, D, E].
• Итерация 6: B, C, D, E генерируют следующий токен.
• Итерация 10: C генерирует <EOS>. C удаляется. Добавляется F. Батч = [B, D, E, F].

Таким образом, GPU постоянно занят обработкой активных запросов, и простои минимизируются.

4. Реализация в vLLM: PagedAttention и Iteration-level Scheduling

vLLM — это высокопроизводительная библиотека для инференса LLM, которая реализует continuous batching через два ключевых механизма: PagedAttention и iteration-level scheduling.

4.1 PagedAttention

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

PagedAttention разбивает KV cache на блоки фиксированного размера (страницы), аналогично виртуальной памяти в операционных системах. Каждый блок может хранить KV-векторы для нескольких токенов. Блоки могут быть не непрерывными в физической памяти, что позволяет:

• Динамически выделять память: Память выделяется по мере необходимости, а не заранее.
• Устранить фрагментацию: Блоки могут быть размещены в любом свободном месте.
• Эффективно управлять continuous batching: При удалении запроса из батча его блоки памяти могут быть немедленно переиспользованы для нового запроса.

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

4.2 Iteration-level Scheduling

Iteration-level scheduling (планирование на уровне итераций) — это механизм, который определяет, какие запросы будут в батче на каждой итерации генерации.

Компоненты scheduler'а в vLLM

1. Waiting Queue: Очередь запросов, ожидающих обработки.
2. Running Queue: Очередь запросов, которые в данный момент генерируют токены.
3. Scheduler Policy: Правила, определяющие, какие запросы из waiting queue добавлять в running queue (например, first-come-first-served, приоритет по длине промпта).

Алгоритм работы scheduler'а на каждой итерации:

1. Проверить, какие запросы в running queue завершили генерацию (сгенерировали <EOS> или достигли max_tokens).
2. Удалить завершённые запросы из running queue и освободить их блоки памяти в PagedAttention.
3. Проверить, есть ли свободное место в батче (с учётом ограничений GPU, таких как max_num_batched_tokens).
4. Если есть свободное место, взять запросы из waiting queue (согласно политике) и добавить их в running queue.
5. Для новых запросов выполнить prefill (заполнить KV cache для входного промпта).
6. Для всех запросов в running queue выполнить decode (сгенерировать один токен).

Пример кода (упрощённая логика scheduler'а):

class Scheduler:
    def __init__(self, max_batch_size, max_tokens_per_batch):
        self.waiting_queue = []
        self.running_queue = []
        self.max_batch_size = max_batch_size
        self.max_tokens_per_batch = max_tokens_per_batch

    def schedule(self):
        # 1. Удаляем завершённые запросы
        self.running_queue = [req for req in self.running_queue if not req.is_finished()]

        # 2. Считаем свободное место
        current_batch_size = len(self.running_queue)
        free_slots = self.max_batch_size - current_batch_size

        # 3. Добавляем новые запросы из очереди ожидания
        while free_slots > 0 and self.waiting_queue:
            new_request = self.waiting_queue.pop(0)
            self.running_queue.append(new_request)
            free_slots -= 1

        # 4. Выполняем prefill для новых запросов (упрощённо)
        for req in self.running_queue:
            if req.is_new():
                req.prefill()

        # 5. Выполняем decode для всех запросов
        for req in self.running_queue:
            req.decode()

        return self.running_queue

5. Преимущества Continuous Batching в vLLM

1. Повышение Throughput: За счёт устранения bubble, throughput (количество сгенерированных токенов в секунду) увеличивается в 4-6 раз по сравнению со статическим батчингом.
2. Снижение Latency: Среднее время ответа (latency) снижается, так как запросы не ждут завершения других.
3. Эффективное использование памяти: PagedAttention позволяет динамически управлять KV cache, что особенно важно при continuous batching, когда запросы постоянно добавляются и удаляются.
4. Масштабируемость: Continuous batching хорошо масштабируется на большое количество запросов и GPU.

6. Недостатки и ограничения

1. Сложность реализации: Требуется сложный scheduler и управление памятью.
2. Overhead на планирование: На каждой итерации scheduler выполняет операции по управлению очередями, что добавляет небольшой overhead.
3. Зависимость от длины промпта: Prefill для новых запросов может быть дорогим, если промпт очень длинный. vLLM решает эту проблему через chunked prefill (разбиение длинного промпта на чанки).
4. Необходимость в PagedAttention: Без эффективного управления KV cache continuous batching может привести к фрагментации памяти.

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

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

Задача: Реализовать симулятор continuous batching для инференса LLM.

Инструменты: Python, NumPy, (опционально) PyTorch для симуляции GPU.

Шаги:

1. Создайте класс Request: с атрибутами prompt_length (длина промпта), max_tokens (максимальное количество генерируемых токенов), generated_tokens (текущее количество сгенерированных токенов), is_finished (флаг завершения).
2. Создайте класс Scheduler: с очередями waiting_queue и running_queue, параметрами max_batch_size и max_tokens_per_batch.
3. Реализуйте метод schedule(): на каждой итерации удаляйте завершённые запросы, добавляйте новые из очереди ожидания, выполняйте prefill для новых и decode для всех.
4. Создайте класс Simulator: генерируйте случайные запросы (с разной длиной промпта и количеством генерируемых токенов), запускайте симуляцию на N итераций.
5. Сравните со статическим батчингом: реализуйте статический батчинг (все запросы начинают и заканчивают вместе) и сравните throughput (количество завершённых запросов за N итераций).

Ожидаемый результат: Вы увидите, что continuous batching обрабатывает значительно больше запросов за то же время, особенно при вариативной длине генерации.

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

10. Навигация

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

Навигация

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