[1/2] DistServe: зачем разделять prefill и decode (Рубрика #Research)

Этот пост посвящен статье DistServe от команды из Peking University, UC San Diego и StepFun, опубликованной на OSDI 2024. Причем в разборе llm-d я уже рассказывал про разделение prefill и decode. А в работе DistServe очень хорошо разобрана инженерная сторона этого решения: что выигрываем, когда разносим стадии по разным GPU, и чем за это платим. Помимо самой статьи есть и открытый код который позволяет не только посмотреть на результаты авторов, но и почитать при желании код.

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

Когда обе стадии живут вместе, новый тяжёлый prefill задерживает уже идущую генерацию. Даём приоритет decode — новые пользователи дольше ждут начала ответа. Плюс обеим стадиям достаётся одна конфигурация параллелизма, хотя потребности у них разные.

В DistServe процесс выглядит так: Запрос → очередь prefill → обработка промпта и первый токен → передача KV-кэша → очередь decode → генерация остальных токенов.

У prefill- и decode-экземпляров свои копии весов модели. KV-кэш содержит промежуточные ключи и значения attention: с ними другая GPU продолжает вычисление, не обрабатывая промпт заново. Prefill и decode теперь можно отдельно масштабировать и настраивать.

Дальше начинается знакомая жизнь распределённой системы: появились сеть, дополнительная память и ещё одно место, где может вырасти очередь :)

Для OPT-66B авторы оценивают кэш промпта из 512 токенов примерно в 1,13 ГБ. При 10 запросах/с это около 90 Гбит/с только на кэш. Причём совпадения средней потребности с полосой мало: всплескам нагрузки нужен запас. Если decode принимает работу медленнее prefill, готовые кэши тоже начинают ждать и занимать память.

Поэтому DistServe учитывает топологию. При медленной сети между серверами соответствующие стадии prefill и decode размещаются на разных GPU внутри одного сервера, чтобы передавать кэш по NVLink. В их измерениях более 95% запросов передавали кэш менее чем за 30 мс.

Цель авторов — goodput на GPU: предельный поток запросов при заданной доле соблюдения требований к задержке (SLO). Проверяются TTFT (time to first token) — время до первого токена — и TPOT (time per output token), среднее время на следующий токен. Обычно целевая доля — 90% запросов. Средний TPOT, кстати, ещё не гарантирует отсутствие отдельных пауз в генерации.

Они проверяли свой подход на OPT-13B/66B/175B, кластере из 32 A100 и данных для чата, дополнения кода и суммаризации. В чатовых тестах получили в 2–4,6 раза больший поток на GPU относительно тогдашнего vLLM; максимум 7,4 раза — относительно DeepSpeed-MII. Это выигрыш всей конфигурации при выбранных SLO. Времена поступления запросов синтезировали, а пороги задержки задали сами; переносить множители на произвольный сервис нельзя.

Остаётся самая интересная часть: сколько GPU выделить каждой стадии и как разложить на них модель. Во второй части поста будет рассказ про intra-op, inter-op и теорию очередей, которая помогает понять, почему самый быстрый отдельный запрос ещё не определяет лучшую конфигурацию сервиса.

#Research #AI #Architecture #Engineering #DistributedSystems