Push-based KV transfer: как AWS оптимизирует задержки в vLLM

AWS Machine Learning · оригинал

Материал подготовлен автоматизированной редакционной системой. Факты можно сверить по указанному первоисточнику.

Команда Amazon Web Services представила push-based коннектор для vLLM, который переносит передачу кэша KV из критического пути декодирования. Это позволяет сократить время до первого токена (TTFT) до трех раз при высокой нагрузке.

Оптимизация передачи данных между этапами префилла и декода напрямую влияет на отклик LLM-сервисов, что критично для высоконагруженных приложений.

В инференсе больших языковых моделей (LLM) существуют две фазы с разными требованиями к ресурсам: prefill, обрабатывающий входной промпт, и decode, генерирующий ответ по одному токену. При их разделении на отдельные узлы (P-D disaggregation) возникает необходимость передачи кэша ключей и значений (KV cache). Традиционно используется pull-based подход, где узел декодирования инициирует запрос данных после завершения префилла. Однако эта последовательная координация удлиняет критический путь, увеличивая время до первого токена (TTFT).

AWS разработала push-based коннектор для vLLM, который меняет логику взаимодействия. В новой схеме узел декодирования заранее регистрирует свои целевые блоки памяти, пока префилл-узел еще вычисляет кэш. Это позволяет начать передачу данных сразу по завершении вычислений, минуя задержки на обнаружение адресов и настройку соединения со стороны потребителя. Такой подход снимает нагрузку с критического пути генерации токенов.

Архитектура опирается на три ключевых шага: регистрацию, синхронизацию (rendezvous) и передачу. Узел декодирования отправляет метаданные о зарезервированных блоках через NIXL notification. Префилл-узел ожидает совпадения событий — завершения вычислений и получения регистрации — после чего инициирует запись данных. Для обеспечения надежности система учитывает повторные запросы и прерванные операции, корректно сопоставляя их по идентификаторам.

Важным аспектом является модель потоков. Подготовка дескрипторов передачи и отправка данных выполняются в фоновом потоке на стороне префилла. Это предотвращает блокировку основного цикла выполнения модели (forward-pass loop), что особенно важно при высокой конкурентности. В отличие от pull-режима, где настройка соединения происходит синхронно в цикле декодера и «зависает» генерацию токенов, push-режим позволяет декодеру оставаться свободным для обработки уже полученных данных.

Бенчмарки проводились на кластере Amazon EC2 P5en с использованием модели Qwen3-32B и сети AWS EFA. Тесты изолировали производительность передачи, отключив кэширование префиксов. Результаты показали, что при низкой нагрузке оба режима работают сопоставимо. Однако при увеличении количества запросов (до 8 QPS) и длины входа (16k токенов) pull-режим демонстрирует резкое деградирование TTFT до 1744 мс.

В то же время push-режим сохраняет TTFT на уровне около 540 мс, обеспечивая трехкратное ускорение. Аналогичная тенденция наблюдается и для TPOT (время на генерацию одного токена): при высоких нагрузках pull-режим достигает 10.4 мс, тогда как push-режим удерживает показатель вблизи 7 мс. Это подтверждает эффективность выноса операций настройки из критического пути декодера.

Выбор между режимами зависит от требований к маршрутизации. Pull-режим позволяет гибко выбирать узел декодирования после завершения префилла, что полезно для балансировки нагрузки. Push-режим требует раннего выбора узла, но обеспечивает минимальные задержки. AWS продолжает развивать эту функциональность, планируя перенести фоновую модель потоков и в pull-коннектор для дальнейшей оптимизации.