Разделённые предварительное заполнение и декодирование

Чтобы генерировать выходные токены из входного запроса, вывод LLM делится на две стадии: предзаполнение и декодирование. Предзаполнение обрабатывает входные токены, заполняя кеши KV, прежде чем перейти к стадии декодирования, которая генерирует токены по одному.

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

Производительность предзаполнения против декодирования

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

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

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

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

KV Messenger

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

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

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

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

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

Шардированные передачи кеша KV

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

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

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

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

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

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

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

Раздельные развертывания

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

DeepSeek-R1

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

DeepSeek опирается на многоголовое латентное внимание, сжимая кеши KV. Поскольку все головы KV сжаты в один латентный вектор, TP не может шардировать кеши KV, так как вместо этого они должны реплицировать латентные векторы на всех рангах. Шардирование происходит после декомпрессии, так как каждый ранг может извлечь различные головы из одного латентного представления. Следовательно, все шардированные кеши KV идентичны как в узле предзаполнения, так и в узле декодирования.

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

С смешанным предзаполнением-декодированием наше развертывание R1 испытывало трудности с постоянным превышением 50 TPS из-за частых прерываний предзаполнения порядка сотен миллисекунд. В отличие от этого, при разделении предзаполнения мы подвергались наказанию в около 100 мс на TTFT для каждого запроса, но один узел предзаполнения мог поддерживать постоянные размеры пакетов на 3 узлах декодирования, обеспечивая пропускную способность выше 90 TPS при обработке нагрузки около 1 QPS на узел декодера. С развертываниями параллельными по данным TPS была чуть ниже, около 50, однако инстанции могли обрабатывать нагрузку 1 QPS на ранг, с 8 рангами на один узел.

Qwen3-Coder

Эта модель 480B использует внимание с группированным запросом (GQA), так что внимание может быть легко шардировано и может извлекать пользу из тензорного параллелизма, не жертвуя памятью для кешей KV. Следовательно, мы могли шардировать модель на 8 GPUs для обоих, предзаполнения и декодирования, соединяя около 3 узлов декодирования с одним узлом предзаполнения. Поскольку внимание шардировано, мы полагаемся на макет кеша KV HND для шардирования кешей KV узлов предзаполнения и декодирования, соединяя ранги узла предзаполнения с ранги узла декодирования и полностью используя все NICs, чтобы передавать срезы параллельно.