Быстрые эмбеддинги на графических процессорах (GPU)
Быстрый и точный поиск жизненно важен для всей экосистемы Perplexity — от Search и Computer до нашей API-платформы. За кулисами основную тяжелую работу выполняют модели эмбеддингов и ранжирования, которые помогают нашим системам находить наиболее релевантные результаты для заданного запроса. Мы достигаем передового со
Быстрый и точный поиск жизненно важен для всей экосистемы Perplexity — от Search и Computer до нашей API-платформы. За кулисами основную тяжелую работу выполняют модели эмбеддингов и ранжирования, которые помогают нашим системам находить наиболее релевантные результаты для заданного запроса. Мы достигаем передового качества и минимальной задержки, обучая и обслуживая собственные модели, такие как pplx-embed.
В этой статье представлен внутренний обзор инфраструктуры обслуживания Perplexity для этого особого класса моделей. Мы обсуждаем наши методы эффективного удовлетворения потребностей в инференсе для поиска на базе ИИ, обеспечивая быстрое прототипирование и оценку моделей при одновременной поддержке нашего поискового индекса экзабайтного масштаба. В совокупности эти методы расширяют границу Парето качества и эффективности поиска, позволяя нам обслуживать агентов и пользователей с наилучшими результатами при минимальных затратах и минимальной задержке.
Эмбеддинги для поиска
В типичной поисковой установке индексированные документы отображаются в многомерное векторное пространство с помощью модели эмбеддингов и сохраняются в векторной базе данных. Создавая эмбеддинг запроса с помощью той же модели, можно найти похожие документы путем поиска векторов, наиболее близких к вектору запроса. Это порождает два различных паттерна трафика для обслуживания движком инференса:
- Пакетный эмбеддинг: при создании, расширении или реиндексации базы данных большой массив документов должен быть преобразован в эмбеддинги в векторном пространстве для максимизации пропускной способности в целях минимизации затрат.
Сразу после векторного поиска необходимо выполнять оценку больших батчей документов, соблюдая баланс между пропускной способностью и задержкой.
- Онлайн-эмбеддинги: при запросе к базе данных короткий запрос должен быть преобразован в эмбеддинг для поиска, сводя задержку к минимуму.
Мы создали нашу инфраструктуру инференса так, чтобы задействовать как можно больше общих компонентов для разных сценариев использования. Поскольку для получения эмбеддингов мы обычно используем небольшие модели Transformer, мы разделяем большую часть реализации с нашим кодом инференса LLM: батчевые эмбеддинги похожи на префилл, ограниченный вычислениями (compute-bound prefill), в то время как онлайн-эмбеддинги, которые часто работают на нескольких токенах, вычислительно похожи на декодирование, ограниченное памятью (memory-bound decode). Таким образом, мы повторно используем наши оптимизированные кернелы префилла и декодирования для обслуживания моделей эмбеддингов. В результате мы можем достичь огромной пропускной способности пакетного инференса с минимальными дополнительными инженерными затратами, сохраняя при этом низкую задержку для рабочих нагрузок онлайн-эмбеддингов.
Тюльпаны, розы и немного плюща
Мы предоставляем инференс через стандартизированные API как внутри компании, так и снаружи через нашу API-платформу. Под капотом в обработке запроса эмбеддинга задействовано несколько сервисов:
- Ivy — это HTTP-шлюз на Rust, к которому обращаются сервисы Perplexity.
Он обрабатывает работу на стороне CPU для таких запросов, как разбор JSON, токенизация, шаблонизация входных данных и разбиение батчей, транслируя запросы в кастомный протокол gRPC для нижестоящих серверов. Это разделение позволяет нам настраивать определенные параметры, связанные с токенизацией и форматированием входных данных, без необходимости затрагивать более тяжелые инстансы инференса.
- Tulip — это интерфейс сервера инференса.
Это gRPC-сервер, реализованный на Rust, tokio и `tonic`. Tulip получает gRPC-запросы инференса, управляя планированием и батчингом. Затем он отправляет батчи в движок ROSE, возвращая клиентам готовые ответы.
- **ROSE** (Runtime-Optimized Serving Engine) реализует инференс моделей.
Он в основном написан на Python и предоставляет кернелы, слои и определения для самого широкого спектра моделей. ROSE реализует прямые проходы (forward passes) через модели, а также обеспечивает управление графами CUDA, специализированное для эмбеддингов. Он связан с Tulip через функцию step(), которая принимает батч и возвращает ссылку на вычисление, выполняемое на ускорителе.
Уделяем внимание тому, что за пределами кернела
Как модели на базе Transformer, так и лежащие в их основе архитектуры Hopper/Blackwell являются зрелыми технологиями, поэтому встраивание инференса на стороне GPU свелось к в целом оптимальной реализации в различных движках инференса. Тем не менее, мы обнаружили дополнительные возможности для оптимизации в средах выполнения и обертках, которые предоставляют модели клиенту «от конца до конца» (end-to-end). В частности, мы выяснили, что можем улучшить задержки за счет тщательного управления графами CUDA и создания абстракции LazyTensor для асинхронного отслеживания результата на стороне GPU в нативном движке на Rust. Мы реализовали эти функции в Tulip, чтобы она могла эффективно взаимодействовать с реализациями моделей в ROSE.
Tulip
Мы спроектировали Tulip максимально легковесным интерфейсом поверх нашего обслуживания моделей. Он обрабатывает входящие запросы в асинхронных задачах Tokio, поддерживая пул запросов, который он отслеживает, и планируя из него батчи для отправки на ускоритель. Механизм планирования в Tulip очень прост: запросы накапливаются, пока Tulip отправляет работу или ожидает результатов. Из накопленных запросов последовательности выбираются в порядке очереди (first-come, first-served) для запуска через модель.
Простой механизм планирования мотивирован наблюдением за производительностью моделей. Для небольших моделей эмбеддингов при длинах последовательностей, которые мы обслуживаем, мы заметили, что линейная стоимость плотных слоев доминирует над квадратичной стоимостью внимания. Таким образом, задержка в основном пропорциональна количеству токенов, а не количеству последовательностей. Следовательно, как только батч становится достаточно большим для насыщения GPU (что составляет около 512 токенов на модели с количеством параметров менее одного миллиарда), добавление в него новых последовательностей не повышает эффективность.
Для эффективного взаимодействия с моделью Tulip полагается на графы CUDA и ленивое отслеживание результатов, чтобы совместить работу GPU и CPU и полностью задействовать доступные ресурсы.
Управление графами CUDA
Выполнение прямого прохода модели включает в себя работу как на стороне CPU, так и на стороне GPU. Центральный процессор отвечает за планирование батчей и запуск кернелов с соответствующими параметрами, в то время как графический процессор выполняет соответствующие умножения матриц, кернелы внимания, нормализации или активации. Для высокопроизводительных рабочих нагрузок, таких как обучение и реиндексация, накладные расходы на стороне CPU ничтожно малы, поскольку размеры батчей и задержка на стороне GPU велики. Однако при меньших размерах батчей работа CPU может перевешивать работу GPU.
Для снижения накладных расходов вместо запуска независимых кернелов можно построить граф CUDA, чтобы захватить метаданные, необходимые для запуска всех кернелов прямого прохода с помощью всего одного вызова драйвера CUDA. Это устраняет необходимость повторного выполнения дорогостоящего кода Python и PyTorch для конфигураций, для которых могут быть захвачены графы CUDA.
Для каждой модели мы отслеживаем точку перегиба, определяя минимальное количество токенов, при котором выполнение на GPU обходится дороже запуска кернела на стороне CPU. Поскольку модели эмбеддингов малы, мы замечаем, что эта точка перегиба наступает при батчах в тысячи токенов и десятки последовательностей. Некоторые реализации внимания зависят от динамических входных данных на стороне хоста для настройки запуска кернелов, что препятствует использованию полных графов CUDA для префилла/плотных слоев модели. Мы добавили в апстрим изменения для соответствующих кернелов, чтобы задействовать их в нашем движке инференса.
Для устранения накладных расходов мы создаем графы CUDA для всей модели для всех моделей эмбеддингов и совмещаем работу CPU с работой GPU. Поскольку графы CUDA минимизируют накладные расходы на стороне CPU, после запуска графа у нас есть свободное время для инициирования и постановки в очередь выполнения следующего батча, как только он становится доступным. Результаты ожидающего батча отслеживаются с помощью LazyTensor, что позволяет асинхронной задаче в Rust блокироваться до завершения выполнения предыдущего батча. Графы CUDA помогают обслуживанию с низкими задержками, гарантируя, что нас не сдерживают затраты на запуск кернелов, и способствуют улучшению планирования в сценарии с высокой пропускной способностью, так как они высвобождают CPU для работы над следующим батчом раньше.
Графы CUDA должны быть захвачены для каждой отдельной конфигурации, что для эмбеддингов означает граф на каждую комбинацию количества последовательностей и токенов. Поскольку эта сетка обширна, мы дополняем количество токенов до бакетов, кратных 64 или 256. Это по-прежнему приводит к тысячам графов, на захват которых для типичной модели может уйти несколько минут. Стоимость захвата складывается из двух источников: жадный прямой проход, который должен быть выполнен для компиляции кернелов и настройки буферов для различных нуждающихся в них кернелов, за которым следует проход захвата, повторно исполняющий код Python.
Мы смягчаем затраты на старте, захватывая графы CUDA лениво по мере работы движка. Мы отслеживаем каждую конфигурацию и гарантируем, что она проходит прогревочный запуск (eager warmup run) перед триггером захвата графа и его воспроизведения при втором обращении. Все последующие выполнения той же конфигурации графа затем проходят через воспроизведение графа CUDA. Ленивый захват графов влияет на задержки p99 во время запуска; однако он ценен тем, что распределяет несколько минут жадной работы на несколько часов. Более быстрое время запуска позволяет нам лучше масштабировать и управлять развертываниями эмбеддингов.
Ленивые тензоры (Lazy Tensors)
Благодаря CUDA работа GPU является асинхронной. Поскольку асинхронный запуск кернела ставит его в очередь в поток, код хоста должен явно выполнить синхронизацию для считывания результирующих векторов. Чтобы обеспечить более высокую степень параллелизма и иметь возможность запускать будущие батчи в ожидании завершения предыдущего на устройстве, мы полагаемся на абстракцию LazyTensor для отслеживания значений.
LazyTensor отслеживает буфер хоста в закрепленной памяти (page-locked memory) и операцию cudaMemcpyAsync с помощью события, копирующего данные с устройства. Он запускается после старта прямого прохода на том же потоке. Поскольку операция копирования должна дождаться выполнения всех предыдущих кернелов в потоке, связанное событие отслеживает как завершение прямого прохода, так и доступность результата на CPU.
Мы используем LazyTensor в нашем движке кодировщика ROSE, чтобы совместить работу GPU и CPU. Вместо того чтобы каждый вызов step() запускал граф CUDA и ожидал его завершения, step() возвращает LazyTensor для асинхронного отслеживания своего результата. В сочетании с графами CUDA это помогает нам достичь низких задержек и более высокой пропускной способности.
ROSE
Мы адаптировали наш движок ROSE, изначально созданный для обслуживания LLM, для выполнения моделей эмбеддингов. Чтобы минимизировать усилия, необходимые для поддержки моделей эмбеддингов, ROSE активно переиспользует код между LLM и эмбеддингами. Например, обслуживание pplx-embed и декодирование LLM Qwen3.5 проходят через одни и те же кернелы. Такое совместное использование позволяет нам с легкостью обслуживать модель эмбеддингов, которая изначально была дообучена (fine-tuned) на базе LLM для прототипирования, оценки и инференса в продакшене.
Для плотных слоев инференс эмбеддингов и LLM идентичен, поскольку векторы токенов обрабатываются независимо. В слоях внимания различия обрабатываются за счет добавления поддержки рваных (ragged) входных данных, а также постраничных настроек префилла и декодирования, требуемых для LLM. При обслуживании модели эмбеддингов мы не создаем кэш KV (KV cache) и перенаправляем запросы к вариантам кернелов внимания, которые поддерживают рваный формат, во избежание дополнения паддингами. Сопутствующие подпрограммы конвертации и калибровки также разделяются с LLM.
Ivy
Ivy, наш прокси-слой HTTP для инференса, также играет важную роль в производительности. Поскольку полезная нагрузка запросов различается в продакшене, маршрутизация отдельных запросов к отдельным репликам может привести к дисбалансу нагрузки. Ivy разбивает запросы с большими батчами на фрагменты и балансирует нагрузку между репликами, улучшая утилизацию и сглаживая задержки. Наша недавняя работа над собственной униграмной токенизацией, полностью развернутой в Ivy, кардинально улучшает задержки по сравнению со стандартными токенизаторами.
...но кернелы все еще имеют значение
ROSE поддерживает различные бэкенды внимания. Разные кернелы могут подходить для определенных размеров задач. Со временем мы интегрировали кернелы FlashInfer 2, FlashInfer 3 и FlashAttention 4 для реализации рваного внимания.
В целом мы замечаем, что FlashAttention 4 работает быстрее. Однако FlashInfer 3 превосходит его на моделях на базе Qwen при очень большой длине последовательностей. Поскольку производительность и настройка могут варьироваться в зависимости от количества и размерности голов внимания, мы сохраняем поддержку нескольких конфигураций и принимаем решение в каждом конкретном случае при обслуживании.
Бенчмарки
Мы проводим бенчмаркинг по сравнению с vLLM v0.22.0, выполняя инференс в точности BF16 на реальных весах моделей и входных данных, полученных из оценочных датасетов. Перед всеми тестовыми прогонами выполнялись прогревочные запуски, которые подтвердили, что расхождение в косинусном расстоянии находится в пределах 0.1%.
Эмбеддинги с низкой задержкой (p50 / p90 / p99 / макс. мс)
Мы приводим время выполнения для предварительно токенизированного размера батча запросов 1, полностью последовательных запросов и длин последовательностей 128, 512 и 4096 токенов.
Скоринг с низкой задержкой (p50 / p90 / p99 / макс. мс)
Предварительно токенизированные размеры батчей запросов: 5, 25 и 50, длина последовательности: 512 токенов.
Эмбеддинги с высокой пропускной способностью (эмб/с)
Размер батча запросов: 100, четыре параллельных процесса отправляют запросы, длина последовательностей: 512, 1024 и 4096 токенов.
Эмбеддинги с высокой степенью параллелизма (p50 / p90 / p99 / макс. мс)
Длина последовательности: 512, размер батча: 1, но мы отправляем 1, 2, 4, 8 и 16 параллельных запросов. Этот бенчмарк также включает затраты на токенизацию через Ivy, а также накладные расходы сети между Ivy и Tulip.
Заключение и будущие работы
Инфраструктура обслуживания, состоящая из Ivy, Tulip и ROSE, позволяет нам обслуживать эмбеддинги для Perplexity с меньшей задержкой и более высокой пропускной способностью, обеспечивая более точный поиск при сниженных затратах по сравнению с готовыми решениями.
Сосредоточившись на конкретных моделях и взяв под контроль весь стек, мы получаем свободу, необходимую для поиска эффективного баланса между производительностью и гибкостью, сочетая высоко переиспользуемые и производительные примитивы Rust с более общим кодом моделирования на Python. Многие движки инференса с открытым исходным кодом, такие как vLLM, SGLang и TokenSpeed, интегрируют в свой стек такие языки, как Rust и C++. Мы инвестировали в Rust в течение последних двух лет и добились отличных результатов как в производительности, так и в удобстве поддержки. Разделяя большую часть реализации эмбеддингов с нашим стеком обслуживания LLM, мы также получаем прирост пропускной способности, не требуя значительных инженерных затрат на обслуживание моделей эмбеддингов.
По мере эволюции моделей мы продолжим совершенствовать каждый слой нашего стека для снижения задержек, ограниченных как CPU, так и GPU. Наши кастомные протоколы на базе gRPC в Ivy и Tulip позволяют нам настраивать коммуникацию для сокращения сетевых задержек, в то время как ROSE обеспечивает основу для повышения вычислительной пропускной способности. Кроме того, по мере роста поддержки свободопоточного (free-threaded) Python в экосистеме мы сможем еще больше улучшить взаимодействие Python и Rust для снижения накладных расходов.