Дезагрегированный prefill и декодирование

Для генерации выходных токенов из входного промпта инференс больших языковых моделей (LLM) разделяется на два этапа: предварительное заполнение (prefill) и декодирование (decode). Prefill выполняется для входных токенов, заполняя KV-кеши перед переходом к этапу декодирования, который генерирует токены по одному. Хотя один шаг декодирования обычно занимает десятки миллисекунд, prefill требует значительно больше времени. При запуске на одних и тех же устройствах смешивание prefill с декодированием снижает производительность декодирования. В этой статье мы исследуем проверенное решение в виде дезагрегированного prefill и декодирования, выполняя их на разных устройствах, чтобы максимизировать как пропускную способность prefill, так и задержки декодирования.

Производительность Prefill против Decode

В типичном движке обслуживания LLM планировщик пакетов выбирает запросы для обработки на каждом этапе выполнения модели. При работе на одном устройстве или узле запросы prefill и декодирования объединяются в пакеты. Стоимость механизма внимания (attention), который агрегируется по длине последовательности, растет как для prefill, так и для декодирования, пропорционально длине записей в KV-кеше (kv_len). Запросы декодирования обычно передают один токен (qo_len=1) с минимальными затратами через другие слои, которые работают независимо с токенами последовательности. Запросы prefill передают тысячи или десятки тысяч токенов со значительными затратами через плотные слои (большое значение qo_len).

На задержку прямого прохода сильнее влияет количество независимых токенов, проходящих через плотные слои (qo_len), чем количество токенов, извлеченных из KV-кеша во время выполнения внимания (kv_len). Механизм внимания может распараллеливаться как по количеству запросов, так и по kv_len, пропорционально длине последовательностей, достигая хорошей утилизации. Prefill ограничен вычислительными ресурсами: при высоком значении qo_len ядра GEMM могут распределять достаточно блоков по измерению M, чтобы полностью использовать вычислительные возможности современных графических процессоров (GPU). Декодирование ограничено памятью: из-за обычно небольших размеров пакетов количество входных данных вдоль измерения M мало, что достаточно только для одного блока. Хотя ядра Split-K GEMM могут улучшить утилизацию SM для малых размеров пакетов токенов, кеши и блоки матричного умножения обычно остаются недостаточно загруженными.

Развернуть

При смешивании пакеты, содержащие запросы для prefill, приводят к более высоким задержкам при прямом проходе, негативно влияя на пропускную способность декодирования всего инстанса. Хотя смешивание запросов prefill с запросами декодирования или использование пакетного prefill может немного улучшить производительность декодирования, трудно поддерживать достаточную пропускную способность prefill для обработки достаточного количества запросов на инстансе, чтобы максимизировать пропускную способность декодирования. В случае больших моделей, с типичной длиной вывода, для поддержания большого размера пакета для декодирования prefill должен выполняться достаточно часто, что значительно увеличивает среднюю задержку и вызывает прерывистость вывода.

Развернуть

Развернуть

Эти проблемы можно решить, используя отдельный набор узлов для выполнения prefill и декодирования. Связав узел prefill с несколькими узлами декодирования, можно планировать достаточное количество запросов для prefill, чтобы максимизировать пропускную способность и поддерживать достаточное количество параллельных запросов на узлах декодирования, чтобы также максимизировать пропускную способность декодирования. Узлы prefill заполняют KV-кеши, которые затем передаются на узлы декодирования. Поскольку декодерам больше не нужно прерываться для prefill, задержки становятся гораздо более детерминированными, так как общее влияние растущего значения kv_len активных запросов становится менее выраженным. Цена этого — увеличение времени до первого токена (TTFT), так как передача KV-кешей по сети может занимать от десятков до сотен миллисекунд.

KV Messenger

В Perplexity наша реализация дезагрегированного prefill и декодирования построена вокруг KV messenger, который взаимодействует с движком LLM для организации передачи KV-кешей с узлов prefill на узлы декодирования через сеть. На стороне prefill messenger принимает запросы от узлов декодирования, передает их планировщику пакетов и отслеживает выполнение прямого прохода, чтобы отправлять KV-кеши с минимально возможной задержкой. На стороне декодера, после выделения неизвлекаемых страниц, messenger блокирует запрос от планирования на декодирование до тех пор, пока не поступит уведомление о завершении передачи KV-кеша и контекста декодера.

Развернуть

Дезагрегация prefill требует высокоскоростных соединений с низкой задержкой, поэтому наша реализация оптимизирована для RDMA, поддерживая как EFA, так и сетевые интерфейсные контроллеры (NIC) ConnectX. KV Messenger построен на libfabric, используя наши обертки fabric-lib для предоставления абстракций более высокого уровня с низкой задержкой поверх примитивов удаленного прямого доступа к памяти (RDMA), реализуя эффективную передачу страниц и метаданных, а также сигнализацию с низкой задержкой. В фоновом режиме fabric-lib координирует GPU и его напрямую подключенные NIC для копирования данных с узла prefill на узел декодера.

После получения узел prefill выделяет соответствующий набор исходных страниц KV и планирует запрос для prefill с использованием локального движка. Чтобы минимизировать задержку, передачи не ждут прямого прохода: вместо этого копирование страниц KV инициируется, как только модель завершает добавление записей в KV-кеш для отдельных слоев. Поскольку запросы prefill могут быть разбиты на части, планировщик пакетов уведомляет KV messenger о текущих запланированных частях до начала выполнения. Чтобы поддерживать графы CUDA при возможности отслеживания слоев, messenger поддерживает выделенный тред, опрашивающий счетчик, который увеличивается после выходной проекции внимания. Счетчик поддерживается только на ведущем узле в сегментированной среде: несмотря на то, что записи KV-кеша действительны после добавления и до внимания, выходная проекция сокращается по рангам, неявно синхронизируя их. Как только наблюдается изменение счетчика, messenger уведомляется, и он вызывает fabric-lib для инициализации передачи слоя.

Развернуть

После завершения передачи последней части копируются любые дополнительные метаданные: спекулятивное декодирование или MTP требуют перемещения логитов и скрытых состояний на декодер. Эти копии также выполняются через RDMA, в и из предварительно выделенных буферов.

По завершении всех ожидающих передач последней части узел prefill освобождает страницы KV и завершает запрос. Узел декодера не уведомляется явно: вместо этого он использует немедленные счетчики для отслеживания количества выполненных операций. Количество операций RDMA на стороне prefill пропорционально количеству переданных страниц. По завершении известного количества копий страниц и контекста fabric-lib вызывает KV messenger, чтобы указать, что запрос готов к декодированию. Messenger освобождает любой контекст и передает запрос движку LLM.

Сегментированные передачи KV-кеша

Если prefill и декодер полагаются на тензорный параллелизм (TP) и сегментируют или реплицируют KV-кеши идентично, единый движок передачи координирует несколько устройств для отправки и получения страниц всех реплик. Чтобы иметь возможность использовать один messenger и движок передачи, несмотря на то, что исполнитель модели реплицирован на нескольких устройствах и процессах, cuMem и cuMemImportFromShareableHandle используются для выделения памяти устройства, поддерживающей KV-кеши, и для ее сопоставления в основной процесс. Движок передачи проверяет топологию узла, чтобы найти NIC и CPU в ближайшем узле NUMA для использования при передаче каждого из срезов KV-кеша.

Если источник и место назначения сегментируются идентично, передачи тривиальны, так как существует взаимно однозначное соответствие между устройствами и страницами источника и назначения. В этой ситуации сегментирование неявно помогает уменьшить задержки при передаче: за счет использования большего количества GPU можно задействовать больше соответствующих NIC, приближаясь к полной утилизации пропускной способности. Однако при наличии несоответствия движок передачи должен разбивать или реконструировать страницы в зависимости от соотношения между исходными и целевыми срезами.

Развернуть

Если prefill разделяет KV-кеш между большим количеством устройств, полные страницы реконструируются на декодере путем отправки соответствующих половин с устройств prefill. Если декодер имеет больше сегментов, он получает страницы из нескольких источников. Декодер должен знать схему сегментирования prefill, чтобы иметь возможность вычислить количество операций записи RDMA, которые он ожидает получить. Если задействована репликация, prefill группирует устройства в наборы реплик, которые реплицируют полный KV-кеш внутри себя. Целевым наборам реплик случайным образом назначается один из исходных наборов, чтобы использовать все доступные устройства для инициации записи RDMA.

Развернуть

Сегментированные передачи требуют небольшой корректировки KV-кешей. По умолчанию FlashInfer полагается на макет NHD, который упорядочивает токены внутри страницы внутри голов внимания. Поскольку кеши, скорее всего, сегментированы по количеству голов внимания, это создает разрыв внутри головы. Передачи RDMA не поддерживают неявно чередующиеся записи, требуя одной операции на голову для выполнения передачи. Вместо этого, чтобы уменьшить количество взаимодействий с libfabric, мы организуем KV-кеши с использованием макета HND, который размещает измерение головы перед количеством токенов. Это обеспечивает непрерывность, позволяя копировать страницу за одну операцию записи.

Спекулятивное декодирование

Спекулятивное декодирование требует небольших корректировок дезагрегированного prefill-decode. В нашей реализации узлам prefill не разрешается осуществлять выборку токенов. Поскольку модели Sonar от Perplexity поддерживают структурированный вывод, мы не хотим усложнять процесс синхронизации реализаций обработчика схем между prefill и декодерами. В механизмах MTP и спекулятивного декодирования выполнение prefill для draft-модели до последнего токена включает выборку токенов из целевой модели.

Развернуть

Чтобы обойти эти проблемы, prefill не включает последний токен входной последовательности. Вместо этого передаются скрытые состояния или логиты из prefill, предшествующие последнему токену, и этот токен рассматривается как токен декодирования на следующем шаге декодера. Хотя это немного увеличивает задержки, поскольку после prefill должен быть выполнен полный шаг декодирования для выпуска первого токена, сложность реализации значительно снижается.

Дезагрегированные развертывания

Мы развернули или экспериментировали с несколькими дезагрегированными конфигурациями с различными моделями для поддержки либо производственного трафика, либо внутренних рабочих нагрузок оценки. Основываясь на размере и механизме внимания моделей, мы выбрали подходящие схемы сегментирования для узлов prefill и декодера, чтобы наилучшим образом использовать GPU.

DeepSeek-R1

С DeepSeek мы рассматривали развертывания как с тензорным параллелизмом (TP), так и с параллелизмом данных (DP). Как обсуждалось в предыдущих статьях блога, развертывания TP обеспечивают лучшую задержку при более низкой пропускной способности, требуя больше GPU для обслуживания интенсивного трафика. Развертывания DP масштабируются гораздо лучше с нагрузкой, однако их пиковая пропускная способность ниже из-за затрат на связь между устройствами или узлами.

DeepSeek опирается на многоголовое латентное внимание (Multi-Head Latent Attention), сжимая KV-кеши. Поскольку все головы KV сжаты в один латентный вектор, TP не может сегментировать KV-кеши, так как вместо этого он должен реплицировать латентные векторы на всех рангах. Сегментирование происходит после декомпрессии, так как каждый ранг может извлекать разные головы из одного и того же латентного представления. Следовательно, все сегменты KV-кеша идентичны как на сегментах prefill, так и на сегментах декодера.

При настройке TP внутри узла как prefill, так и декодеры сегментированы идентично. Передачи отправляются со всех рангов, чтобы полностью использовать все доступные NIC. Однако при развертывании DP, где размер ранга TP меньше или каждый ранг DP назначен одному GPU, любое устройство prefill, которое содержит реплицированную копию KV-кеша, может отправить его. Чтобы сбалансировать запросы по всем доступным NIC, мы случайным образом выбираем GPU и NIC для отправки KV-кеша от prefill к декодеру.

При смешанном prefill-decode наше развертывание R1 с трудом преодолевало порог в 50 токенов в секунду (TPS) из-за частых прерываний prefill длительностью в сотни миллисекунд. Напротив, при разделении prefill мы получили штраф около 100 мс к TTFT для каждого запроса, но один узел prefill мог поддерживать стабильные размеры пакетов на 3 узлах декодера, обеспечивая пропускную способность более 90 TPS при обработке нагрузки около 1 QPS на узел декодера. При развертывании с параллелизмом данных TPS был немного ниже, около 50, однако инстансы могли обрабатывать нагрузку 1 QPS на ранг, при 8 рангах на один узел.

Qwen3-Coder

Эта модель на 480 млрд параметров использует групповое внимание (GQA), поэтому внимание можно легко сегментировать, и она может извлечь выгоду из тензорного параллелизма без ущерба для памяти под KV-кеши. Следовательно, мы могли бы сегментировать модель на 8 GPU как для prefill, так и для декодирования, связывая около 3 узлов декодера с одним узлом prefill. Поскольку внимание сегментировано, мы полагаемся на макет KV-кеша HND для сегментирования KV-кешей prefill и декодера, связывая ранги prefill с рангами декодера и полностью используя все NIC для передачи срезов параллельно.

Цитата

@online{perplexity2025DisaggregatedPrefillAndDecode,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Kevin Hu and Erin McKnight and Phi Hoang},
  title = {{Disaggregated prefill and decode}},
  date = {2025-08-01},
  url = {https://perplexity.coolkeypoint.com/hub/blog/disaggregated-prefill-and-decode}
}