Как работает continuous batching в TGI (Hugging Face Text Generation Inference)?

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

batching|Continuous batching — это техника динамического формирования батчей на каждом шаге генерации токенов, реализованная в TGI (Hugging Face Text Generation Inference). В отличие от статического batching, где все запросы обрабатываются синхронно до завершения самого длинного, TGI использует токен-уровневый scheduler: новые запросы добавляются в очередь, а на каждой итерации scheduler выбирает следующие токены для всех активных запросов, немедленно освобождая ресурсы завершённых. Это позволяет значительно повысить throughput (пропускную способность) и снизить latency (задержку) для сервисов генерации текста.

1. Термин: Continuous batching (непрерывная пакетная обработка)

batching|Continuous batching — это метод организации инференса LLM, при котором батч запросов не фиксируется на весь цикл генерации, а обновляется на каждом шаге декодирования. В каждый момент времени GPU обрабатывает набор запросов, находящихся на разных стадиях генерации. Новые запросы могут быть добавлены в батч сразу после освобождения места завершёнными запросами.

Ключевая идея: GPU эффективен при больших батчах, но в генерации текста длина ответа непредсказуема. Continuous batching позволяет держать батч максимально заполненным, минимизируя простои GPU.

2. Проблема статического batching

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

Недостатки статического batching

Пример: Батч из 4 запросов с длинами генерации [10, 50, 100, 200] токенов. Статический batching обработает все 200 шагов, хотя 3 запроса завершились раньше. GPU будет загружен на 100% только первые 10 шагов, затем загрузка падает.

3. Как работает continuous batching в TGI

TGI использует токен-уровневый scheduler (scheduler на уровне токенов). Архитектура включает три основных компонента:

1. Очередь запросов — все входящие запросы попадают в очередь ожидания.
2. Scheduler — на каждом шаге (iteration) решает, какие запросы войдут в текущий батч.
3. KV cache manager — управляет кэшем ключей и значений для каждого запроса.

Алгоритм работы на каждом шаге

1. Выбор активных запросов scheduler берёт все запросы, которые ещё не завершили генерацию (активные).
2. Добавление новых если в батче есть свободные слоты (после завершения предыдущих запросов), scheduler добавляет новые запросы из очереди.
3. Один шаг декодирования GPU выполняет forward pass для всех запросов в батче, генерируя следующий токен для каждого.
4. Проверка завершения запросы, сгенерировавшие токен <EOS> или достигшие max_new_tokens, помечаются как завершённые.
5. Освобождение ресурсов KV cache завершённых запросов освобождается, слоты становятся доступны для новых запросов на следующем шаге.

Псевдокод цикла инференса TGI

# Упрощённая логика continuous batching
queue = []  # очередь запросов
active = []  # активные запросы (id, tokens, kv_cache)

while True:
    # Добавляем новые запросы, если есть место
    while len(active) < max_batch_size and queue:
        req = queue.pop(0)
        # Инициализируем KV cache для первого токена
        req.kv_cache = init_kv_cache(req.prompt)
        active.append(req)
    
    if not active:
        break  # нет активных запросов
    
    # Формируем батч из активных запросов
    batch = [req.last_token for req in active]
    # Forward pass: получаем логиты для всех запросов
    logits = model.forward(batch, kv_caches=[req.kv_cache for req in active])
    
    # Выбираем следующий токен для каждого запроса
    for i, req in enumerate(active):
        next_token = sample(logits[i])
        req.tokens.append(next_token)
        req.kv_cache.update(next_token)
        if next_token == EOS or len(req.tokens) >= req.max_new_tokens:
            req.finished = True
            # Освобождаем KV cache (в реальности — помечаем слот как свободный)
    
    # Удаляем завершённые запросы из active
    active = [req for req in active if not req.finished]

Ключевые особенности реализации в TGI

• Динамическое управление KV cache TGI использует фиксированный пул памяти под KV cache. Каждый запрос получает слот в этом пуле. При завершении запроса слот освобождается и может быть переиспользован новым запросом.
• Preemption (вытеснение): если очередь переполнена, TGI может вытеснить запрос с низким приоритетом, сохранив его KV cache на CPU и возобновив позже (поддерживается не во всех версиях).
• Оптимизация attention TGI использует PagedAttention (как в vLLM) для эффективного управления KV cache, что позволяет избежать фрагментации памяти.

4. Детали реализации: KV cache management

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

В continuous batching KV cache каждого запроса хранится отдельно. При добавлении нового запроса в батч его KV cache инициализируется (для prompt). При генерации каждого нового токена KV cache расширяется.

Проблема Размер KV cache растёт линейно с длиной генерации. Если запросы имеют разную длину, память фрагментируется.

Решение TGI Используется блочное выделение памяти (block-based allocation) — KV cache разбивается на блоки фиксированного размера (например, 16 или 32 токена). Каждый запрос получает блоки по мере необходимости. Это позволяет эффективно переиспользовать память завершённых запросов.

5. Преимущества и недостатки continuous batching

6. Сравнение со static batching

Экспериментальные данные (из бенчмарков TGI): При нагрузке 100 запросов с длинами от 50 до 500 токенов continuous batching даёт прирост throughput в 2–4 раза по сравнению со static batching при том же batch size.

7. Влияние на метрики (latency, throughput)

• Throughput (токенов/сек): Растёт за счёт более плотного заполнения батча. В TGI continuous batching позволяет достичь throughput, близкого к теоретическому максимуму для данного GPU.
• Latency (время ответа): Для отдельных запросов latency снижается, особенно для коротких. Однако при высокой нагрузке latency может увеличиться из-за ожидания в очереди. TGI использует batching timeout — максимальное время ожидания перед формированием батча, чтобы балансировать между latency и throughput.
• Time-to-first-token (TTFT): Незначительно увеличивается, так как запрос может ждать завершения текущего шага, но обычно это миллисекунды.

Формула для оценки throughput при continuous batching:

Throughput ≈ (batch_size_avg * tokens_per_step) / step_time

Где batch_size_avg — средний размер батча за всё время генерации. При static batching batch_size_avg падает к концу, при continuous — остаётся высоким.

8. Связь с другими техниками оптимизации инференса

9. Практический пример: использование TGI с continuous batching

TGI включает continuous batching по умолчанию. Пример запуска сервера:

# Запуск TGI с моделью Llama-2-7B
text-generation-launcher \
  --model-id meta-llama/Llama-2-7b-chat-hf \
  --num-shard 1 \
  --max-batch-prefill-tokens 4096 \
  --max-batch-total-tokens 40960 \
  --max-input-length 2048 \
  --max-total-tokens 4096 \
  --waiting-served-ratio 1.2

Параметры, влияющие на continuous batching:

• --max-batch-prefill-tokens — максимальное количество токенов для префилла (обработки промпта) в одном батче.
• --max-batch-total-tokens — максимальное общее количество токенов (промпт + генерация) в батче.
• --waiting-served-ratio — соотношение между ожидающими и обслуживаемыми запросами, влияет на агрессивность добавления новых запросов.

Пример запроса через curl

curl -X POST http://localhost:8080/generate \
  -H "Content-Type: application/json" \
  -d '{"inputs": "Расскажи про continuous batching", "parameters": {"max_new_tokens": 100}}'

TGI автоматически поместит этот запрос в очередь и scheduler решит, когда добавить его в батч.

10. Ограничения и альтернативы

Ограничения continuous batching в TGI

• Требует поддержки dynamic batching на уровне модели (не все кастомные модели легко адаптировать).
• При очень малом batch size (1–2 запроса) выигрыш незначителен.
• Сложность отладки — неочевидное поведение при переполнении очереди.

Альтернативы

Выбор между TGI и vLLM TGI лучше интегрирован с экосистемой Hugging Face, vLLM часто даёт больший throughput. Для production-систем с высокой нагрузкой рекомендуется сравнивать оба решения на своих данных.

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

Задача Написать симулятор continuous batching на Python, который моделирует обработку потока запросов с разной длиной генерации.

Инструменты Python, библиотеки numpy, time, random.

Шаги:

1. Создайте класс Request с полями: id, prompt_length (фиксированная, например 10 токенов), generation_length (случайная от 10 до 200), status (waiting, active, finished).
2. Реализуйте класс Scheduler:
   • Параметры: max_batch_size (например, 4), queue (список запросов), active (список активных).
   • Метод step(): добавляет новые запросы из очереди, если есть место; выполняет один шаг генерации для всех активных (уменьшает оставшиеся токены); удаляет завершённые.
3. Создайте класс Simulator:
   • Генерирует пуассоновский поток запросов (среднее время между запросами — 0.1 секунды).
   • Запускает цикл, на каждой итерации вызывая scheduler.step().
   • Собирает статистику: среднее время ожидания в очереди, throughput (завершённые запросы в секунду), средний размер батча.
4. Сравните с симулятором статического batching (батч формируется один раз, все запросы обрабатываются до конца самого длинного).

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

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

import random
import time
from collections import deque

class Request:
    def __init__(self, req_id, gen_length):
        self.id = req_id
        self.remaining = gen_length
        self.finished = False

class Scheduler:
    def __init__(self, max_batch_size):
        self.max_batch_size = max_batch_size
        self.queue = deque()
        self.active = []
    
    def add_request(self, req):
        self.queue.append(req)
    
    def step(self):
        # Добавляем новые запросы, пока есть место
        while len(self.active) < self.max_batch_size and self.queue:
            self.active.append(self.queue.popleft())
        
        # Обрабатываем один шаг для всех активных
        for req in self.active:
            req.remaining -= 1
            if req.remaining <= 0:
                req.finished = True
        
        # Удаляем завершённые
        self.active = [req for req in self.active if not req.finished]
        return len(self.active)

# Симуляция (фрагмент)
scheduler = Scheduler(max_batch_size=4)
# ... генерация запросов и цикл

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

Навигация

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