Бързи ембединги на GPU
Бързото и точно търсене е от жизненоважно значение за цялата Perplexity – от Search и Computer до нашата API платформа. Зад кулисите тежката работа се извършва от ембединг и ранкинг модели, които помагат на нашите системи да идентифицират най-релевантните резултати за дадена заявка. Постигаме най-съвременното
Бързото и точно търсене е от жизненоважно значение за цялата Perplexity – от Search и Computer до нашата API платформа. Зад кулисите тежката работа се извършва от ембединг и ранкинг модели, които помагат на нашите системи да идентифицират най-релевантните резултати за дадена заявка. Постигаме най-съвременно качество и латентност, като обучаваме и обслужваме наши собствени модели, като например pplx-embed.
Тази статия представя поглед "под капака" към инфраструктурата за обслужване на Perplexity за този специален клас модели. Обсъждаме нашите техники за ефективно посрещане на нуждите от инференция на AI-нативното търсене, което позволява бързо прототипиране и оценка на моделите, като същевременно захранва нашия индекс за търсене в екзабайтов мащаб. Тези техники заедно разширяват границата на Парето за качество и ефективност на търсенето, което ни позволява да обслужваме агенти и потребители с възможно най-добрите резултати при минимални разходи и латентност.
Ембединги за търсене
В типична настройка за търсене индексираните документи се мапират към високомерно векторно пространство с помощта на ембединг модел и се съхраняват във векторна база данни. Чрез ембедингване на заявка със същия модел, сходни документи могат да бъдат открити чрез намиране на векторите, най-близки до тези на заявката. Това води до два различни модела на трафик, които двигателят за инференция трябва да обслужва:
- Партиден ембединг: при изграждане, разширяване или повторно индексиране на базата данни, масивни документи трябва да бъдат ембедингнати във векторното пространство, максимизирайки пропусквателната способност с цел минимизиране на разходите.
След векторното търсене, големи партиди от документи трябва да бъдат оценявани, постигайки баланс между пропусквателна способност и латентност.
- Онлайн ембединг: при заявка към базата данни, кратка заявка трябва да бъде ембедингната за търсения, минимизирайки латентността.
Изградихме нашата инфраструктура за инференция така, че да използва възможно най-много общи компоненти за различните случаи на употреба. Тъй като обикновено използваме малки Transformer модели за създаване на ембединги, ние споделяме голяма част от реализацията с нашия код за LLM инференция: партидните ембединги са сходни с предварителното зареждане (prefill), обвързано с изчисленията, докато онлайн ембедингите, които често се изпълняват върху няколко токена, са изчислително сходни с дезикацията (decode), обвързана с паметта. По този начин ние преизползваме нашите оптимизирани prefill и decode кернели, за да обслужваме ембединг модели. В резултат на това можем да постигнем мащабна пропусквателна способност при партидна инференция с минимални допълнителни инженерни усилия, като същевременно запазваме ниска латентност за работните натоварвания с онлайн ембединги.
Лалета (Tulips), рози (Roses) и малко бръшлян (Ivy)
Предоставяме инференция чрез стандартизирани API, както вътрешно, така и външно чрез нашата API платформа. Зад кулисите множество услуги участват в обработката на ембединг заявка:
- Ivy е Rust HTTP шлюз, който услугите на Perplexity извикват.
Тя обработва работата от страната на CPU за заявки като JSON парсиране, токенизация, темплейтиране на входни данни и разделяне на партиди, като превежда заявките към потребителски gRPC протокол за достъп до сървърите надолу по веригата. Това разделение ни позволява да конфигурираме определени параметри около токенизацията и форматирането на входните данни, без да се налага да пипаме по-тежките инстанции за инференция.
- Tulip е интерфейсът на сървъра за инференция.
Това е gRPC сървър, реализиран с Rust, tokio и `tonic`. Tulip получава gRPC заявки за инференция, като управлява планирането и партидирането. След това тя изпраща партидите към двигателя ROSE, връщайки завършени отговори на клиентите.
- **ROSE** (Runtime-Optimized Serving Engine) реализира инференция на модели.
Тя е дефинирана предимно в Python, като предоставя кернели, слоеве и дефиниции за голямо разнообразие от модели. ROSE реализира правите преминавания (forward passes) през моделите, като същевременно осигурява управление на CUDA графи, специализирано за ембединги. Тя е свързана с Tulip чрез функция step(), която приема партида и връща референция към изчислението, което тя изпълнява върху ускорителя.
Обръщане на внимание отвъд кернела (Paying Attention Beyond the Kernel)
Както базираните на Transformer модели, така и поддържащите ги архитектури Hopper/Blackwell са зрели технологии, така че инференцията на ембединги от страната на GPU се е консолидирала до до голяма степен оптимална реализация в различните двигатели за инференция. Въпреки това открихме допълнителни възможности за подобрение във времената за изпълнение и интерфейсите, които предоставят моделите от край до край на клиента. По-специално установихме, че можем да подобрим латентностите чрез внимателно управление на CUDA графите и изграждане на абстракция LazyTensor за асинхронно проследяване на резултат от страната на GPU в нейния Rust двигател. Реализирахме тези функционалности в Tulip, за да може тя ефективно да взаимодейства с реализациите на модели в ROSE.
Tulip
Проектирахме Tulip така, че да бъде възможно най-лек интерфейс върху обслужването на нашите модели. Тя обработва входящите заявки в асинхронни задачи на Tokio, поддържайки пул от заявки, които проследява, и планира партиди от тях за изпращане към ускорителя. Механизмът за планиране в Tulip е много прост: заявките се акумулират, докато Tulip изпраща работа или чака резултати. От натрупаните заявки се избират последователности на принцип "пръв дошъл, пръв обслужен", които да бъдат изпълнени през модела.
Опростеният механизъм за планиране е мотивиран от наблюдение върху производителността на модела. За малки ембединг модели, при дължините на последователностите, които обслужваме, забележихме, че линейната цена на плътните слоеве доминира над квадратната цена на attention. Следователно латентността е предимно пропорционална на броя токени, а не на броя последователности. Следователно, след като една партида стане достатъчно голяма, за да насите GPU (което е около 512 токена за модел под един милиард параметъра), пакетирането на повече последователности в нея не подобрява ефективността.
За да взаимодейства ефективно с модела, Tulip разчита на CUDA графи и лениво проследяване на резултатите, за да припокрие работата на GPU и CPU и да използва напълно наличните ресурси.
Управление на CUDA графи
Изпълнението на правото преминаване на модела включва работа както от страната на CPU, така и от страната на GPU. CPU отговаря за планирането на партиди и стартирането на кернели с подходящите параметри, докато GPU изпълнява съответните матрични умножения, attention, нормиране или активиращи кернели. За работни натоварвания с висока пропусквателна способност, като обучение и повторно индексиране, овърхедът от страната на CPU е пренебрежим, тъй като и размерите на партидите, и латентността от страната на GPU са големи. Въпреки това, при по-малки размери на партидите, работата от страната на CPU може да надхвърли работата от страната на GPU.
За да се смекчи овърхедът, вместо да се стартират независими кернели, може да се изгради CUDA граф за улавяне на метаданните, необходими за стартиране на всички кернели на право преминаване с едно единствено извикване към CUDA драйвера. Това елиминира необходимостта от повторно изпълнение на скъп Python и PyTorch код за конфигурациите, за които могат да бъдат заснети CUDA графи.
За всеки модел проследяваме точка на пречупване, определяща минималния брой токени, при който изпълнението на GPU е по-скъпо от извикването на кернела от страната на CPU. Тъй като ембединг моделите са малки, наблюдаваме, че тази точка на пречупване настъпва при партиди от хиляди токени и десетки последователности. Някои реализации на attention разчитат на динамични входни данни от страната на хоста за конфигуриране на извикванията на кернели, което възпрепятства пълните CUDA графи за prefill/плътни модели. Ние upstream-нахме промени в съответните кернели, за да ги активираме в нашия двигател за инференция.
За да адресираме овърхеда, ние изграждаме CUDA графи за целия модел за всички ембединг модели и припокриваме работата на CPU с работата на GPU. Тъй като CUDA графите минимизират овърхеда от страната на CPU, след като графът бъде стартиран, разполагаме със свободно време за стартиране и поставяне в опашка на изпълнението на следващата партида, когато тя стане налична. Резултатите от чакащата партида се проследяват с LazyTensor, което позволява на асинхронна задача в Rust да блокира, докато предходната партида завърши изпълнението си. CUDA графите помагат на обслужването с ниска латентност, като гарантират, че не сме възпрепятствани от разходите за стартиране на кернели, и улесняват подобреното планиране в случай на висока пропусквателна способност, тъй като освобождават CPU да работи по следващата партида по-рано.
CUDA графите трябва да бъдат заснети за всяка отделна конфигурация, което за ембединги означава граф за всяка комбинация от брой последователности и брой токени. Тъй като тази мрежа е обширна, ние подлагаме броя на токените на подложки (padding) до кофи (buckets), които са кратни на 64 или 256. Това все пак води до хиляди графи, чието заснемане може да отнеме няколко минути за типичен модел. Цената на заснемането идва от два източника: евристично право преминаване, което трябва да бъде изпълнено за компилиране на кернели и настройване на буфери за различните кернели, които се нуждаят от тях, последвано от заснемащото изпълнение, което преизпълнява Python кода.
Смекчаваме разходите за стартиране чрез лениво заснемане на CUDA графи при обслужването от двигателя. Проследяваме всяка конфигурация и гарантираме, че тя преминава през евристичен тест за загряване преди задействане на заснемането и възпроизвеждането на графа при второто попадение. Всички следващи изпълнения на същата конфигурация на графа след това преминават през възпроизвеждане на CUDA граф. Ленивото заснемане на графи оказва влияние върху p99 латентностите по време на стартиране; въпреки това то е ценно за разпределянето на няколко минути евристична работа в рамките на множество часове. По-бързите времена за стартиране ни позволяват да мащабираме и управляваме по-добре ембединг деплоймънтите.
Лениви тензори (Lazy Tensors)
Чрез CUDA работата с GPU е асинхронна. Тъй като асинхронното стартиране на кернел го поставя в опашка в поток, хост кодът трябва изрично да се синхронизира, за да прочете получените вектори. За да улесним по-висока степен на паралелизъм и да можем да стартираме бъдещи партиди, докато чакаме предходната да завърши на устройството, разчитаме на абстракция LazyTensor за проследяване на стойности.
LazyTensor проследява хост буфер в заключена в паметта страница (page-locked) и операция cudaMemcpyAsync чрез събитие, копиращо данни от устройството. То се стартира след стартирането на правото преминаване върху същия поток. Тъй като операцията по копиране трябва да изчака изпълнението на всички предходни кернели в потока, свързаното събитие проследява както завършването на правото преминаване, така и достъпността на резултата на CPU.
Ние използваме LazyTensor в нашия ембединг двигател ROSE, за да припокрием работата на GPU и CPU. Вместо всяко извикване на step() да изпълнява CUDA графа и да чака завършването му, step() връща LazyTensor, за да проследи асинхронно своя резултат. В комбинация с CUDA графите, това ни помага да постигнем ниски латентности и по-добра пропусквателна способност.
ROSE
Ние адаптирахме нашия двигател ROSE, който първоначално създадохме за LLM обслужване, така че да обработва и изпълнението на ембединг модели. За да минимизираме усилията, необходими за поддръжка на ембединг модели, ROSE агресивно преизползва код между LLM и ембединги. Например, обслужването на pplx-embed и дезикацията на Qwen3.5 LLM минават през едни и същи кернели. Това споделяне ни позволява лесно да обслужваме ембединг модел, който първоначално е бил дообучен от LLM за прототипиране, оценка и производствена инференция.
За плътни слоеве (dense layers) ембедингът и LLM инференцията са идентични, тъй като векторните токени се обработват независимо. В attention слоевете разликите се овладяват чрез добавяне на поддръжка за грапави (ragged) входни данни, наред с настройките за paged prefill и decode, изисквани от LLM. Когато обслужваме ембединг модел, ние не инстанцираме KV кеш и изпращаме заявки към варианти на attention кернели, които поддържат грапавия формат, за да се избегне подложката (padding). Поддържащите процедури за конвертиране и калибриране също се споделят с LLM.
Ivy
Ivy, нашият HTTP прокси слой за инференция, също играе важна роля за производителността. Тъй като полезните товари на заявките варират в производствена среда, насочването на отделни заявки към отделни реплики може да доведе до дисбаланс на натоварването. Ivy разделя заявките с големи партиди на късове (chunks) и балансира натоварването между репликите, подобрявайки използването и изглаждайки латентността. Нашата скорошна работа по вътрешна униграмна токенизация, внедрена напълно в Ivy, драстично подобрява латентностите в сравнение с готови токенизатори.
...но кернелите все още имат значение
ROSE поддържа разнообразни attention бекенди. Различните кернели могат да бъдат подходящи за конкретни размери на задачите. С течение на времето интегрирахме кернели FlashInfer 2, FlashInfer 3 и FlashAttention 4 за реализиране на грапав (ragged) attention.
Като цяло наблюдаваме, че FlashAttention 4 е по-бърз. Въпреки това FlashInfer 3 го превъзхожда при Qwen-базирани модели с много дълги последователности. Тъй като производителността и настройката могат да варират в зависимост от броя и размерността на attention главите, ние поддържаме множество конфигурации и вземаме решение за всеки отделен случай при обслужване.
Бенчмаркове
Извършваме бенчмаркинг спрямо vLLM v0.22.0, изпълнявайки инференция с точност BF16 върху реални тегла на модели и входни данни, извлечени от набори от данни за оценка. На всички тестове за време са предхождани тестове за загряване, които потвърждават, че отклонението в косинусната сходност е в рамките на 0.1%.
Ембединги с ниска латентност (p50 / p90 / p99 / макс. ms)
Отчитаме времена за изпълнение за предварително токенизиран размер на партидата от заявки 1, напълно последователни заявки, дължини на последователностите 128, 512 и 4096 токена.
Оценяване с ниска латентност (p50 / p90 / p99 / макс. ms)
Предварително токенизирани размери на партиди заявки 5, 25 и 50, дължина на последователността 512 токена.
Ембединги с висока пропусквателна способност (emb/s)
Размер на партидата заявки 100, четири едновременни процеса подават заявки, дължини на последователностите 512, 1024 и 4096 токена.
Ембединги с висока едновременност (p50 / p90 / p99 / макс. ms)
Дължина на последователността 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, за да намалим овърхеда.