Szybkie osadzenia na procesorach GPU
Szybkie i dokładne wyszukiwanie ma kluczowe znaczenie dla całego Perplexity, od Search and Computer do naszej platformy API. Za kulisami ciężką pracę wykonują modele osadzania i rankingu, które pomagają naszym systemom zidentyfikować najbardziej trafne wyniki dla danego zapytania. Osiągamy najwyższą jakość
Szybkie i dokładne wyszukiwanie ma kluczowe znaczenie dla całego Perplexity, od Search and Computer po naszą platformę API. Za kulisami ciężką pracę wykonują modele osadzania i rankingu, które pomagają naszym systemom zidentyfikować najbardziej trafne wyniki dla danego zapytania. Osiągamy najwyższą jakość i opóźnienia dzięki trenowaniu i obsłudze własnych modeli, takich jak pplx-embed.
Ten artykuł przedstawia kulisy infrastruktury obsługi Perplexity dla tej specjalnej klasy modeli. Omawiamy nasze techniki mające na celu wydajne zaspokojenie potrzeb inferencyjnych wyszukiwania opartego na sztucznej inteligencji, umożliwiając szybkie prototypowanie i ocenę modeli przy Jednoczesnym zasilaniu naszego indeksu wyszukiwania na skalę eksabajtową. Techniki te łącznie poszerzają granicę Pareto w zakresie jakości i wydajności wyszukiwania, umożliwiając nam obsługę agentów i użytkowników z najlepszymi możliwymi wynikami przy najniższych kosztach i opóźnieniach.
Osadzenia do wyszukiwania
W typowej konfiguracji wyszukiwania zindeksowane dokumenty są odwzorowywane na wysokowymiarową przestrzeń wektorową przy użyciu modelu osadzającego i przechowywane w bazie danych wektorowych. Poprzez osadzenie zapytania przy użyciu tego samego modelu można zlokalizować podobne dokumenty, znajdując wektory najbliższe wektorowi zapytania. Powoduje to pojawienie się dwóch różnych wzorców ruchu dla silnika inferencji do obsługi:
- Osadzanie wsadowe: podczas budowania, rozszerzania lub ponownego indeksowania bazy danych duże partie dokumentów muszą zostać osadzone w przestrzeni wektorowej, maksymalizując przepustowość w celu zminimalizowania kosztów.
Po wyszukiwaniu wektorowym duże wsady dokumentów muszą zostać ocenione, co pozwala zachować równowagę między przepustowością a opóźnieniem.
- Osadzanie online: podczas odpytywania bazy danych krótkie zapytanie musi zostać osadzone w celu wyszukiwania, co minimalizuje opóźnienia.
Zbudowaliśmy naszą infrastrukturę inferencji tak, aby wykorzystywać jak najwięcej wspólnych komponentów w różnych przypadkach użycia. Ponieważ zazwyczaj używamy małych modeli Transformer do generowania osadzeń, dzielimy większość implementacji z naszym kodem inferencji LLM: osadzenia wsadowe są podobne do fazy wstępnego wypełniania (prefill) ograniczonej obliczeniowo, podczas gdy osadzenia online, które często działają na kilku tokenach, są obliczeniowo podobne do dekodowania ograniczonego pamięciowo. W związku z tym powtórnie wykorzystujemy nasze zoptymalizowane jądra prefill i dekodowania do obsługi modeli osadzających. W rezultacie możemy osiągnąć ogromną przepustowość inferencji wsadowej przy minimalnym dodatkowym nakładzie pracy inżynieryjnej, zachowując jednocześnie niskie opóźnienie dla obciążeń związanych z osadzeniami online.
Tulipy, róże i bluszcz
Udostępniamy inferencję poprzez standaryzowane interfejsy API, zarówno wewnętrznie, jak i zewnętrznie za pośrednictwem naszej platformy API. W tle w przetwarzaniu żądania osadzania bierze udział wiele usług:
- Ivy to brama HTTP w języku Rust, którą wywołują usługi Perplexity.
Obsługuje on pracę po stronie procesora CPU dla żądań, takich jak parsowanie JSON, tokenizacja, szablonowanie danych wejściowych i podział wsadów, tłumacząc żądania na niestandardowy protokół gRPC dla serwerów podrzędnych. To oddzielenie pozwala nam konfigurować niektóre parametry dotyczące tokenizacji i formatowania danych wejściowych bez konieczności dotykania cięższych instancji inferencji.
- Tulip to interfejs serwera inferencji.
Jest to serwer gRPC zaimplementowany w języku Rust przy użyciu tokio i `tonic`. Tulip odbiera żądania inferencji gRPC, obsługując planowanie i tworzenie wsadów. Następnie wysyła wsady do silnika ROSE, zwracając ukończone odpowiedzi klientom.
- **ROSE** (Runtime-Optimized Serving Engine) implementuje inferencję modeli.
Jest on zdefiniowany głównie w Python, zapewniając jądra, warstwy i definicje dla szerokiej gamy modeli. ROSE implementuje przejścia w przód (forward passes) przez modele, zapewniając również zarządzanie grafikami CUDA wyspecjalizowane dla osadzeń. Jest on połączony z Tulip za pomocą funkcji step(), która pobiera wsad i zwraca odniesienie do obliczeń wykonywanych na akceleratorze.
Zwracanie uwagi na to, co poza jądrem (Paying Attention Beyond the Kernel)
Zarówno modele oparte na Transformerach, jak i leżące u ich podstaw architektury Hopper/Blackwell to dojrzałe technologie, więc inferencja osadzeń po stronie GPU zbiegła się w kierunku w dużej mierze optymalnej implementacji dla różnych silników inferencji. Mimo to odkryliśmy dodatkowe możliwości poprawy w środowiskach wykonawczych i szkieletach, które udostępniają modele klientowi od końca do końca. W szczególności odkryliśmy, że możemy poprawić opóźnienia poprzez staranne zarządzanie grafikami CUDA oraz budowę abstrakcji LazyTensor do asynchronicznego śledzenia wyników po stronie GPU w rodzimym silniku Rust. Wdrożyliśmy te funkcje w Tulip, dzięki czemu może on skutecznie współpracować z implementacjami modeli ROSE.
Tulip
Zaprojektowaliśmy Tulip tak, aby był jak najbardziej lekkim interfejsem nad obsługą naszych modeli. Obsługuje on przychodzące żądania w zadaniach asynchronicznych Tokio, utrzymując pulę żądań, które śledzi i z której harmonogramuje wsady w celu wysłania ich do akceleratora. Mechanizm planowania w Tulip jest bardzo prosty: żądania kumulują się, gdy Tulip wysyła pracę lub czeka na wyniki. Spośród skumulowanych żądań sekwencje są wybierane na zasadzie „kto pierwszy, ten lepszy” do uruchomienia przez model.
Ten prosty mechanizm planowania jest motywowany obserwacją wydajności modelu. W przypadku małych modeli osadzających, przy długościach sekwencji, które obsługujemy, zauważyliśmy, że liniowy koszt warstw gęstych dominuje nad kwadratowym kosztem atencji. W związku z tym opóźnienie jest w dużej mierze proporcjonalne do liczby tokenów, a nie liczby sekwencji. W rezultacie, gdy wsad jest wystarczający duży, aby wysycić GPU (co wynosi około 512 tokenów w modelu o liczbie parametrów poniżej miliarda), upakowanie w nim większej liczby sekwencji nie poprawia wydajności.
Aby skutecznie współpracować z modelem, Tulip opiera się na grafach CUDA i leniwym śledzeniu wyników, aby nakładać na siebie pracę GPU i CPU oraz w pełni wykorzystać dostępne zasoby.
Zarządzanie grafikami CUDA
Uruchomienie przejścia w przód modelu wiąże się z pracą zarówno po stronie procesora (CPU), jak i karty graficznej (GPU). Procesor jest odpowiedzialny za harmonogramowanie wsadów i uruchamianie jąder z odpowiednimi parametrami, podczas gdy karta graficzna wykonuje odpowiednie jądra mnożenia macierzy, atencji, normy lub aktywacji. W przypadku obciążeń o wysokiej przepustowości, takich jak trenowanie i ponowne indeksowanie, narzuty po stronie CPU są pomijalne, ponieważ rozmiary wsadów i opóźnienia po stronie GPU są duże. Jednak przy mniejszych rozmiarach wsadów praca po stronie CPU może przeważać nad pracą po stronie GPU.
Aby złagodzić narzuty, zamiast uruchamiać niezależne jądra, można zbudować graf CUDA w celu przechwycenia metadanych wymaganych do uruchomienia wszystkich jąder przejścia w przód za pomocą jednego wywołania sterownika CUDA. Eliminuje to potrzebę ponownego uruchamiania kosztownego kodu Python i PyTorch dla konfiguracji, dla których można przechwycić grafy CUDA.
Dla każdego modelu śledzimy punkt przegięcia, określający minimalną liczbę tokenów, przy której wykonanie na GPU jest droższe niż uruchomienie jądra po stronie CPU. Ponieważ modele osadzające są małe, zauważamy, że ten punkt przegięcia występuje przy wsadach liczących tysiące tokenów i dziesiątki sekwencji. Niektóre implementacje mechanizmu atencji opierają się na dynamicznych danych wejściowych po stronie hosta w celu skonfigurowania uruchomień jądra, co uniemożliwia pełne grafy CUDA prefill/gęste dla modelu. Wdrożyliśmy upstreamowe zmiany w odpowiednich jądrach, aby umożliwić ich użycie w naszym silniku inferencji.
Aby zaradzić narzutom, budujemy pełne grafy CUDA dla wszystkich modeli osadzających i nakładamy pracę procesora CPU na pracę karty graficznej (GPU). Ponieważ grafy CUDA minimalizują narzuty po stronie procesora CPU, po uruchomieniu grafu mamy wolny czas na zainicjowanie i dodanie do kolejki wykonania następnego wsadu, gdy tylko będzie dostępny. Wyniki oczekującego wsadu są śledzone za pomocą obiektu LazyTensor, co pozwala zadaniu asynchronicznemu w języku Rust blokować się do momentu zakończenia wykonywania poprzedniego wsadu. Grafy CUDA pomagają w obsłudze o niskim opóźnieniu, zapewniając, że nie jesteśmy powstrzymywani przez koszt uruchamiania jąder, i ułatwiają lepsze planowanie w przypadku wysokiej przepustowości, ponieważ zwalniają procesor CPU do szybszej pracy nad następnym wwadem.
Grafy CUDA muszą zostać przechwycone dla każdej odrębnej konfiguracji, co w przypadku osadzeń oznacza jeden graf na kombinację liczby sekwencji i liczby tokenów. Ponieważ ta siatka jest rozległa, dopełniamy liczbę tokenów do kubełków będących wielokrotnością 64 lub 256. Wciąż skutkuje to tysiącami grafów, których przechwycenie dla typowego modelu może zająć kilka minut. Koszt przechwycenia pochodzi z dwóch źródeł: zachłannego przejścia w przód, które musi zostać wykonane w celu skompilowania jąder i skonfigurowania buforów dla różnych ich rodzajów, a następnie przebiegu przechwytywania, który ponownie wykonuje kod Python.
Łagodzimy koszty uruchomienia, leniwie (lazily) przechwjując grafy CUDA w miarę obsługi żądań przez silnik. Śledzimy każdą konfigurację i dbamy o to, aby przeszła ona przez zachłanny przebieg rozgrzewkowy przed wyzwoleniem przechwytywania i odtwarzania grafu przy drugim trafieniu. Wszystkie kolejne wykonania tej samej konfiguracji grafu przechodzą następnie przez odtwarzanie grafu CUDA. Leniwe przechwytywanie grafu ma wpływ na opóźnienia p99 podczas uruchamiania; jest jednak cenne w rozłożeniu kilku minut pracy zachłannej na wiele godzin. Krótszy czas uruchamiania pozwala nam lepiej skalować i zarządzać wdrożeniami osadzeń.
Leniwe tensory (Lazy Tensors)
Poprzez CUDA praca GPU jest asynchroniczna. Ponieważ uruchomienie jądra asynchronicznie umieszcza je w kolejce w strumieniu, kod hosta musi jawnie się zsynchronizować, aby odczytać wynikowe wektory. Aby ułatwić wyższy stopień zrównoleglenia i umożliwić inicjowanie przyszłych wsadów podczas oczekiwania na zakończenie poprzedniego na urządzeniu, polegamy na abstrakcji LazyTensor do śledzenia wartości.
LazyTensor śledzi bufor hosta w pamięci zablokowanej na stronie (page-locked) oraz operację cudaMemcpyAsync za pomocą zdarzenia kopiującego dane z urządzenia. Jest on uruchamiany po uruchomieniu przejścia w przód w tym samym strumieniu. Ponieważ operacja kopiowania musi czekać na wykonanie wszystkich poprzednich jąder w strumieniu, powiązane zdarzenie śledzi zarówno zakończenie przejścia w przód, jak i dostępność wyniku na procesorze CPU.
Wykorzystujemy obiekty LazyTensor w naszym silniku kodującym ROSE, aby nakładać na siebie pracę GPU i CPU. Zamiast każdego wywołania step() uruchamiającego graf CUDA i oczekiwania na jego zakończenie, step() zwraca LazyTensor w celu asynchronicznego śledzenia jego wyniku. W połączeniu z grafikami CUDA pomaga nam to osiągnąć niskie opóźnienia i lepszą przepustowość.
ROSE
Zaadaptowaliśmy nasz silnik ROSE, pierwotnie zbudowany do obsługi LLM, tak aby obsługiwał również wykonywanie modeli osadzających. Aby zminimalizować wysiłek potrzebny do obsługi modeli osadzających, ROSE agresywnie wielokrotnie wykorzystuje kod między LLM a osadzeniami. Na przykład obsługa pplx-embed i dekodowanie LLM Qwen3.5 przechodzą przez te same jądra. To współdzielenie pozwala nam łatwo obsługiwać model osadzający, który został pierwotnie dostrojony z LLM do celów prototypowania, oceny i inferencji produkcyjnej.
W przypadku warstw gęstych inferencja osadzania i LLM jest identyczna, ponieważ wektory tokenów są przetwarzane niezależnie. W warstwach atencji różnice są obsługiwane poprzez dodanie obsługi danych wejściowych o poszarpanej strukturze (ragged inputs), wraz z ustawieniami paginated prefill i dekodowania wymaganymi przez modele LLM. Podczas obsługi modelu osadzającego nie instancjonujemy pamięci podręcznej KV (KV cache) i wysyłamy żądania do wariantów jąder atencji obsługujących format ragged, aby uniknąć dopełniania (padding). Powiązane procedury konwersji i kalibracji są również współdzielone z LLM.
Ivy
Ivy, nasza warstwa proxy HTTP inferencji, również odgrywa ważną rolę w wydajności. Ponieważ ładunki żądań różnią się w produkcji, routowanie poszczególnych żądań do poszczególnych replik może powodować nierównowagę obciążenia. Ivy dzieli żądania o dużych wsadach na fragmenty i równoważy ich obciążenie między replikami, poprawiając wykorzystanie i wygładzając opóźnienia. Nasza niedawna praca nad własną tokenizacją unigramową, w pełni wdrożoną w Ivy, drastycznie poprawia opóźnienia w porównaniu z gotowymi tokenizatorami.
...ale jądra nadal mają znaczenie
ROSE obsługuje różne backendy atencji. Różne jądra mogą być dopasowane do konkretnych rozmiarów problemów. Z czasem zintegrowaliśmy jądra FlashInfer 2, FlashInfer 3 i FlashAttention 4 w celu zaimplementowania poszarpanej atencji (ragged attention).
Ogólnie rzecz biorąc, zauważamy, że FlashAttention 4 jest szybszy. Jednak FlashInfer 3 przewyższa go pod względem wydajności na modelach opartych na Qwen przy bardzo długich długościach sekwencji. Ponieważ wydajność i strojenie mogą się różnić w zależności od liczby i wymiaru głów atencji, zachowujemy wsparcie dla wielu konfiguracji i podejmujemy decyzję dla każdego przypadku z osobna podczas obsługi.
Testy porównawcze (Benchmarks)
Przeprowadzamy testy porównawcze z vLLM v0.22.0, uruchamiając inferencję na precyzji BF16 na rzeczywistych wagach modeli i danych wejściowych pochodzących z zestawów danych ewaluacyjnych. Wszystkie przebiegi czasowe były poprzedzone przebiegami rozgrzewkowymi, które weryfikowały, czy rozbieżność w podobieństwie cosinusowym mieści się w granicach 0,1%.
Osadzenia o niskim opóźnieniu (p50 / p90 / p99 / max ms)
Raportujemy czasy wykonywania dla wstępnie zetokienizowanego wsadu żądań o rozmiarze 1, w pełni sekwencyjnych żądań, o długościach sekwencji 128, 512 i 4096 tokenów.
Ocenianie o niskim opóźnieniu (p50 / p90 / p99 / max ms)
Rozmiary wstępnie zetokienizowanych wsadów żądań 5, 25 i 50, długość sekwencji 512 tokenów.
Osadzenia o wysokiej przepustowości (emb/s)
Rozmiar wsadu żądania 100, cztery współbieżne procesy przesyłające żądania, długości sekwencji 512, 1024 i 4096 tokenów.
Osadzenia o wysokiej współbieżności (p50 / p90 / p99 / max ms)
Długość sekwencji 512, rozmiar wsadu 1, ale wysyłamy 1, 2, 4, 8 i 16 współbieżnych żądań. Ten benchmark uwzględnia również koszty tokenizacji przez Ivy oraz narzut sieciowy między Ivy a Tulip.
Podsumowanie i przyszłe prace
Infrastruktura obsługi złożona z Ivy, Tulip i ROSE pozwala nam obsługiwać osadzenia dla Perplexity z niższym opóźnieniem i lepszą przepustowością, co skutkuje dokładniejszym wyszukiwaniem przy obniżonych kosztach w porównaniu z gotowymi rozwiązaniami.
Skupiając się na konkretnych modelach i przejmując pełną odpowiedzialność za cały stos, uzyskujemy swobodę niezbędną do osiągnięcia skutecznej równowagi między wydajnością a elastycznością, łącząc wysoko wysoce wielokrotnego użytku i wydajne prymitywy w języku Rust z bardziej ogólnym kodem modelowania w Python. Wiele silników inferencji open-source, takich jak vLLM, SGLang i TokenSpeed, integruje ze swoim stosem języki takie jak Rust i C++. W ciągu ostatnich dwóch lat zainwestowaliśmy w Rust i zebraliśmy wielkie owoce zarówno pod względem wydajności, jak i łatwości utrzymania. Dzieląc większość implementacji osadzeń z naszym stosem obsługi LLM, uzyskujemy również zyski w przepustowości, bez konieczności ponoszenia znacznych nakładów inżynieryjnych na utrzymanie modeli osadzających.
W miarę ewolucji modeli będziemy nadal ulepszać każdą warstwę naszego stosu, aby zmniejszyć opóźnienia zarówno procesora (CPU), jak i karty graficznej (GPU). Nasze niestandardowe protokoły oparte na gRPC w ramach Ivy i Tulip pozwalają nam dostosować komunikację w celu zmniejszenia opóźnień sieciowych, podczas gdy ROSE stanowi podstawę do poprawy przepustowości obliczeniowej. Ponadto wraz ze wzrostem wsparcia dla Pythona z wolnymi wątkami (free-threaded Python) w całym ekosystemie będziemy w stanie dalszej poprawy interoperacyjności Python-Rust w celu ograniczenia narzutów.