Швидкі ембеддинги на GPU
Швидкий і точний пошук є життєво важливим для всього Perplexity: від Search та Computer до нашої платформи API. За лаштунками важку роботу виконують моделі ембеддингів та ранжування, які допомагають нашим системам визначати найрелевантніші результати для заданого запиту. Ми досягаємо стану найсучаснішої
Швидкий і точний пошук є життєво важливим для всього Perplexity: від Search та Computer до нашої платформи API. За лаштунками важку роботу виконують моделі ембеддингів та ранжування, які допомагають нашим системам визначати найрелевантніші результати для заданого запиту. Ми досягаємо найсучаснішої якості та затримки завдяки навчанню та обслуговуванню власних моделей, таких як pplx-embed.
У цій статті представлено погляд зсередини на інфраструктуру обслуговування Perplexity для цього спеціального класу моделей. Ми обговорюємо наші методи ефективного задоволення потреб пошуку на базі ШІ у висновках, що дозволяє швидко створювати прототипи та оцінювати моделі, одночасно забезпечуючи роботу нашого пошукового індексу екзабайтного масштабу. Ці методи спільно розширюють межу Парето якості та ефективності пошуку, дозволяючи нам обслуговувати агентів і користувачів із найкращими можливими результатами за найнижчої вартості та затримки.
Ембеддинги для пошуку
У типовій пошуковій конфігурації проіндексовані документи зіставляються з багатовимірним векторним простором за допомогою моделі ембеддингів і зберігаються у векторній базі даних. Зробивши ембеддинг запиту за допомогою тієї ж моделі, можна знайти схожі документи, відшукавши вектори, найближчі до вектора запиту. Це породжує два різні шаблони трафіку для обслуговування механізмом висновків:
- Пакетний ембеддинг: під час створення, розширення або повторного індексування бази даних масиви документів повинні бути переведені в ембеддинги у векторний простір, максимізуючи пропускну здатність для мінімізації витрат.
Після векторного пошуку великі пакети документів мають бути оцінені з дотриманням балансу між пропускною здатністю та затримкою.
- Онлайн-ембеддинг: під час запиту до бази даних короткий запит має бути перетворений на ембеддинг для пошуку, що мінімізує затримку.
Ми створили нашу інфраструктуру висновків так, щоб задіяти якомога більше спільних компонентів для різних сценаріїв використання. Оскільки ми зазвичай використовуємо невеликі моделі Transformer для створення ембеддингів, ми ділимо більшу частину реалізації з нашим кодом висновків LLM: пакетні ембеддинги подібні до попереднього заповнення (prefill), обмеженого обчисленнями, тоді як онлайн-ембеддинги, які часто працюють на кількох токенах, обчислювально подібні до декодування, обмеженого пам'яттю. Таким чином, ми повторно використовуємо наші оптимізовані ядра попереднього заповнення та декодування для обслуговування моделей ембеддингів. У результаті ми можемо досягти величезної пропускної здатності пакетних висновків із мінімальними додатковими інженерними зусиллями, зберігаючи при цьому низьку затримку для робочих навантажень онлайн-ембеддингів.
Тюльпани, троянди та трохи плюща
Ми надаємо висновки через стандартизовані API, як внутрішньо, так і зовнішньо через нашу платформу API. Під капотом у процесі обробки запиту на ембеддинг беруть участь кілька служб:
- Ivy — це шлюз HTTP на Rust, до якого звертаються служби Perplexity.
Він обробляє роботу на боці CPU для таких запитів, як синтаксичний аналіз JSON, токенізація, створення шаблонів вхідних даних та розбиття пакетів, перетворюючи запити на власний протокол gRPC для серверів нижнього рівня. Це розділення дозволяє нам налаштовувати певні параметри щодо токенізації та форматування вхідних даних без необхідності торкатися більш важких екземплярів висновків.
- Tulip — це інтерфейс сервера висновків.
Це сервер gRPC, реалізований за допомогою Rust, tokio та `tonic`. Tulip отримує запити висновків gRPC, керуючи плануванням та пакетуванням. Потім він надсилає пакети до рушія ROSE, повертаючи клієнтам завершені відповіді.
- **ROSE** (Runtime-Optimized Serving Engine) реалізує виконання моделей.
Він переважно визначений у Python і надає ядра, шари та визначення для найрізноманітніших моделей. ROSE реалізує прямі прогони через моделі, а також забезпечує керування графами CUDA, спеціалізоване для ембеддингів. Він з'єднаний із Tulip за допомогою функції step(), яка приймає пакет і повертає посилання на обчислення, що виконуються на прискорювачі.
Звертаємо увагу на те, що за межами ядра
Як моделі на базі Transformer, так і базові архітектури Hopper/Blackwell є зрілими технологіями, тому вбудовування висновків на боці GPU здебільшого зійшлося до оптимальної реалізації в різних механізмах висновків. Попри це, ми виявили додаткові можливості для покращення в середовищах виконання та обгортках, які забезпечують взаємодію моделей із клієнтом від початку до кінця. Зокрема, ми виявили, що можемо покращити затримки шляхом ретельного керування графами CUDA та створення абстракції LazyTensor для асинхронного відстеження результату на боці GPU в нативному рушії Rust. Ми реалізували ці функції в Tulip, щоб забезпечити ефективну взаємодію з реалізаціями моделей ROSE.
Tulip
Ми розробили Tulip так, щоб він був настільки легким інтерфейсом поверх нашого обслуговування моделей, наскільки це можливо. Він обробляє вхідні запити в асинхронних завданнях Tokio, підтримуючи пул запитів, які він відстежує, і планує пакети для передачі на прискорювач. Механізм планування в Tulip дуже простий: запити накопичуються, поки Tulip надсилає роботу або чекає на результати. З накопичених запитів послідовності вибираються за принципом «першим прийшов — першим обслуговується» для прогону через модель.
Простий механізм планування мотивується спостереженням за продуктивністю моделі. Для невеликих моделей ембеддингів при довжинах послідовностей, які ми обслуговуємо, ми помітили, що лінійна вартість щільних шарів домінує над квадратичною вартістю уваги. Таким чином, затримка здебільшого пропорційна кількості токенів, а не кількості послідовностей. Отже, щойно пакет стає достатньо великим, щоб наситити GPU (що становить близько 512 токенів на моделі з кількістю параметрів менше мільярда), упаковка в нього додаткових послідовностей не підвищує ефективність.
Для ефективної взаємодії з моделлю Tulip покладається на графи CUDA та лінійне відстеження результатів, щоб перекривати роботу GPU та CPU і повністю використовувати доступні ресурси.
Керування графами CUDA
Виконання прямого прогону моделі передбачає роботу як на боці CPU, так і на боці GPU. CPU відповідає за планування пакетів та запуск ядер із відповідними параметрами, тоді як GPU виконує відповідні ядра множення матриць, уваги (attention), нормування або активації. Для високопродуктивних робочих навантажень, таких як навчання та повторне індексування, накладні витрати на боці CPU є незначними, оскільки розміри пакетів і затримка на боці GPU є великими. Однак при менших розмірах пакетів робота на боці CPU може перевищувати роботу на боці GPU.
Для зменшення накладних витрат замість запуску незалежних ядер можна побудувати граф CUDA для захоплення метаданих, необхідних для запуску всіх ядер прямого прогону за допомогою одного виклику драйвера CUDA. Це усуває необхідність повторного запуску дорогого коду Python та PyTorch для конфігурацій, для яких можуть бути захоплені графи CUDA.
Для кожної моделі ми відстежуємо точку перегину, визначаючи мінімальну кількість токенів, за якої виконання на GPU є дорожчим, ніж запуск ядра на боці CPU. Оскільки моделі ембеддингів невеликі, ми спостерігаємо, що ця точка перегину настає при пакетах у тисячі токенів і десятки послідовностей. Деякі реалізації уваги покладаються на динамічні вхідні дані на боці хоста для конфігурації запусків ядер, що перешкоджає повному попередньому заповненню моделі / щільним графам CUDA. Ми перенесли зміни до відповідних ядер у вихідний репозиторій (upstream), щоб увімкнути їх у нашому рушії висновків.
Для усунення накладних витрат ми створюємо графи CUDA для всієї моделі для всіх моделей ембеддингів і перекриваємо роботу CPU роботою GPU. Оскільки графи CUDA мінімізують накладні витрати на боці CPU, після запуску графіка ми маємо вільний час для ініціювання та постановки в чергу виконання наступного пакета щоразу, коли він стає доступним. Результати виконання пакета, що очікує на розгляд, відстежуються за допомогою LazyTensor, що дозволяє асинхронному завданню в Rust блокувати виконання до завершення попереднього пакета. Графи CUDA допомагають обслуговуванню з низькою затримкою, гарантуючи, що ми не гальмуємо через вартість запуску ядер, і сприяють покращенню планування у випадку високої пропускної здатності, оскільки вони звільняють CPU для виконання роботи над наступним пакетом раніше.
Графи CUDA повинні захоплюватися для кожної окремої конфігурації, що для ембеддингів означає графік для кожної комбінації кількості послідовностей та кількості токенів. Оскільки ця сітка є величезною, ми доповнюємо кількість токенів до бакетів, що є кратними 64 або 256. Це все одно призводить до тисяч графів, захоплення яких для типової моделі може зайняти кілька хвилин. Вартість захоплення складається з двох джерел: жадібного прямого прогону, який повинен бути виконаний для компіляції ядер і налаштування буферів для різних ядер, які їх потребують, за яким слідує запуск захоплення, що повторно виконує код Python.
Ми зменшуємо витрати на запуск, ледаче (lazily) захоплюючи графи 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 і декодування LLM Qwen3.5 проходять через одні й ті самі ядра. Такий спільний доступ дозволяє нам легко обслуговувати модель ембеддингів, спочатку доопрацьовану (fine-tuned) з LLM для створення прототипів, оцінки та виробничих висновків.
Для щільних шарів (dense layers) висновки ембеддингів та LLM ідентичні, оскільки вектори токенів обробляються незалежно. У шарах уваги відмінності обробляються шляхом додавання підтримки нерівних (ragged) вхідних даних разом із налаштуваннями попереднього заповнення та декодування з розподілом на сторінки, які вимагаються LLM. Під час обслуговування моделі ембеддингів ми не створюємо кеш KV і надсилаємо запити до варіацій ядер уваги, які підтримують нерівний формат, щоб уникнути заповнення (padding). Допоміжні процедури перетворення та калібрування також є спільними з 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. Багато open-source механізмів висновків, таких як vLLM, SGLang і TokenSpeed, інтегрують такі мови, як Rust та C++, у свій стек. Останні два роки ми інвестували в Rust і досягли великих успіхів як у продуктивності, так і в зручності підтримки. Поділившись більшою частиною реалізації ембеддингів із нашим стеком обслуговування LLM, ми також отримуємо приріст пропускної здатності, не витрачаючи значних інженерних зусиль на підтримку моделей ембеддингів.
У міру еволюції моделей ми продовжуватимемо вдосконалювати кожен шар нашого стека, щоб зменшити як затримки, обмежені процесором (CPU), так і затримки, обмежені графічним процесором (GPU). Наші власні протоколи на основі gRPC в Ivy та Tulip дозволяють нам налаштовувати зв'язок для зменшення мережевих затримок, тоді як ROSE забезпечує основу для підвищення обчислювальної пропускної здатності. Крім того, у міру зростання підтримки вільного від потоків Python (free-threaded Python) в екосистемі ми зможемо ще більше покращити взаємодію між Python та Rust, щоб зменшити накладні витрати.