Chunked Prefill

Разбиение предварительной обработки на блоки — это метод обслуживания моделей, при котором длинная входная подсказка делится на небольшие группы токенов и обрабатывается за несколько шагов вперемешку с генерацией токенов для других запросов. Без этого механизм инференса, принявший длинную подсказку, тратит на неё целый шаг, и все ответы, уже передаваемые другим пользователям, приостанавливаются до его завершения. В статье SARATHI 2023 года предложили предварительную обработку блоками, а в последующей работе Sarathi-Serve на её основе реализовали планирование без остановок. В Dynamic SplitFuse от DeepSpeed-FastGen используется та же идея. Сейчас предварительная обработка блоками по умолчанию включена в vLLM и доступна в SGLang и TensorRT-LLM — движках, обслуживающих множество больших языковых моделей (LLM) с открытыми весами.

Предварительная обработка, декодирование и остановки генерации#

Инференс авторегрессионного трансформера состоит из двух этапов:

• Предварительная обработка: Модель параллельно обрабатывает все токены подсказки, вычисляет для них ключи и значения и сохраняет их в KV-кэше. На каждый прочитанный байт весов предварительная обработка выполняет много операций, поэтому обычно её ограничивает вычислительная мощность.
• Декодирование: На каждом шаге модель генерирует по одному токену для каждого запроса, каждый раз считывая весь кэш. На каждый прочитанный байт декодирование выполняет мало арифметических операций, поэтому обычно его ограничивает пропускная способность памяти.

Шаг GPU, включающий предварительную обработку подсказки длиной 20 000 токенов, длится гораздо дольше шага декодирования, поэтому все запросы в пакете ждут его завершения. В Sarathi-Serve такие паузы называют остановками генерации; для пользователей с короткими подсказками они проявляются как всплески задержки между токенами.

Как работает разбиение префилла на фрагменты#

Планировщик выделяет каждому шагу бюджет токенов, задаваемый в vLLM параметром max_num_batched_tokens. Сначала он добавляет по одному токену для каждого выполняемого декодирования, а затем заполняет оставшийся бюджет блоками ожидающих подсказок. Если подсказка не помещается, её обработка продолжается на следующем шаге, постепенно расширяя KV-кэш, пока не обработается вся подсказка и не начнётся её декодирование. Объединение двух этапов в одном шаге также позволяет лучше использовать GPU: декодирование оставляет вычислительные ресурсы свободными, и их могут задействовать блоки предварительной обработки.

Разбиение на блоки меняет только момент обработки токенов подсказки. Каждый блок учитывает ключи и значения из кэша всех предыдущих блоков и собственные токены с помощью причинной маски, поэтому результат совпадает с однопроходной предварительной обработкой. Это показано на одном слое внимания в PyTorch:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
prompt_len, chunk = 1000, 256
q, k, v = torch.randn(3, 1, 8, prompt_len, 64).unbind(0)  # one layer: 8 heads, 64 dims each

# One-shot prefill: all prompt tokens attend causally in a single pass
full = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# Chunked prefill: each chunk attends to the cached keys and values of earlier chunks plus itself
outputs = []
for start in range(0, prompt_len, chunk):
    end = min(start + chunk, prompt_len)
    mask = torch.ones(end - start, end, dtype=torch.bool).tril(diagonal=start)  # causal over absolute positions
    q_chunk, k_cached, v_cached = q[..., start:end, :], k[..., :end, :], v[..., :end, :]
    outputs.append(F.scaled_dot_product_attention(q_chunk, k_cached, v_cached, attn_mask=mask))
chunked = torch.cat(outputs, dim=2)

print(f"{len(outputs)} chunks, same result as one pass: {torch.allclose(full, chunked, atol=1e-5)}")

Скрипт выводит 4 chunks, same result as one pass: True: подсказка длиной 1 000 токенов обрабатывается блоками по 256, 256, 256 и 232 токена, а результаты совпадают с однопроходной обработкой.

Связанные методы обслуживания моделей#

• Непрерывная пакетная обработка: Добавляет и удаляет запросы на каждом шаге, не дожидаясь завершения всего пакета. Разбиение предварительной обработки на блоки позволяет планировать длинные подсказки в рамках этого цикла.
• PagedAttention: Хранит KV-кэш в страницах фиксированного размера, поэтому кэш подсказки, обрабатываемой за несколько шагов, может расти без предварительного резервирования памяти для всей её длины.
• Кэширование подсказок: Повторно использует кэшированные ключи и значения повторяющегося префикса подсказки, полностью исключая предварительную обработку этой части.
• Разделение предварительной обработки и декодирования: Выполняет два этапа в отдельных пулах GPU и передаёт между ними KV-кэш, чтобы каждый пул можно было масштабировать и настраивать под собственное узкое место.
• Спекулятивное декодирование: Ускоряет этап декодирования: модель сначала предлагает несколько токенов, а затем проверяет их за один проход.

Применение#

Разбиение предварительной обработки на блоки особенно полезно, когда один сервер обрабатывает длинные и короткие подсказки. Ассистент с генерацией с дополненной выборкой (RAG), заполняющий своё контекстное окно отчётом на 50 страниц, может обрабатывать его блоками, пока продолжается потоковая передача чатов других пользователей. Модели «зрение — язык» (VLM) добавляют к каждой подсказке сотни или тысячи токенов изображения, поэтому в конвейерах обработки видеопотока, отправляющих кадры на описание, возникает именно такое сочетание запросов. Интерфейс Ultralytics для LLM отправляет текст и изображения на любую совместимую с OpenAI конечную точку, включая собственный сервер vLLM. На нём запросы конвейера YOLO26 на описание кадров с обнаруженными людьми обрабатываются блоками одновременно с трафиком других пользователей.

Настройка бюджета токенов#

Бюджет влияет на два показателя задержки, и между ними приходится искать компромисс. Небольшой бюджет оставляет мало места для предварительной обработки на каждом шаге: это снижает задержку между токенами для пользователей потоковой передачи, но увеличивает время до первого токена для новых подсказок. Большой бюджет даёт обратный эффект. Кроме того, при обработке каждой следующей части длинной подсказки повторно считывается растущий KV-кэш предыдущих блоков, поэтому очень маленькие блоки увеличивают трафик памяти. Операторы настраивают бюджет по медианной и хвостовой задержке инференса, времени до первого токена и пропускной способности на реалистичных длинах подсказок. В руководстве по оптимизации vLLM бюджет в 2 048 токенов рекомендуется для снижения задержки между токенами, а для максимальной пропускной способности на крупных GPU с небольшими моделями — более 8 192 токенов.