Snelle embeddings op GPU's
Snel en nauwkeurig zoeken is van vitaal belang voor heel Perplexity, van Search en Computer tot ons API Platform. Achter de schermen wordt het zware werk verricht door embedding- en ranking-modellen, die onze systemen helpen om de meest relevante resultaten te identificeren voor een bepaalde query. We bereiken state-of-the
Snel en nauwkeurig zoeken is van vitaal belang voor heel Perplexity, van Search en Computer tot ons API Platform. Achter de schermen wordt het zware werk verricht door embedding- en ranking-modellen, die onze systemen helpen om de meest relevante resultaten te identificeren voor een bepaalde query. We bereiken state-of-the-art kwaliteit en latentie door onze eigen modellen te trainen en te serveren, zoals pplx-embed.
Dit artikel geeft een kijkje achter de schermen van de serving-infrastructuur van Perplexity voor deze speciale klasse van modellen. We bespreken onze technieken om efficiënt in te spelen op de inferentiebehoeften van AI-native zoeken, wat snelle prototyping en evaluatie van modellen mogelijk maakt terwijl het onze exabyte-scale search index aandrijft. Deze technieken breiden gezamenlijk het Pareto-grensvlak van zoekkwaliteit en -efficiëntie uit, waardoor we agents en gebruikers kunnen bedienen met de best mogelijke resultaten tegen de laagste kosten en latentie.
Embeddings voor Zoeken
In een typische zoekopstelling worden geïndexeerde documenten afgebeeld op een hoog-dimensionale vectorruimte met behulp van een embedding-model en opgeslagen in een vector database. Door een query te embedden met behulp van hetzelfde model, kunnen vergelijkbare documenten worden gevonden door de vectoren te vinden die het dichtst bij die van de query liggen. Dit leidt tot twee verschillende verkeerspatronen die een inferentie-engine moet bedienen:
- Batch Embedding: bij het bouwen, uitbreiden of opnieuw indexeren van de database moeten bulkdocumenten worden ingebed in de vectorruimte, waarbij de doorvoer wordt gemaximaliseerd om de kosten te minimaliseren.
Na vectorzoeken moeten grote batches documenten worden gescoord, waarbij een balans wordt gevonden tussen doorvoer en latentie.
- Online Embedding: bij het bevragen van de database moet een korte query worden ingebed voor lookups, om latentie te minimaliseren.
We hebben onze inferentie-infrastructuur zo opgebouwd dat we zoveel mogelijk gemeenschappelijke componenten hergebruiken over verschillende use cases heen. Aangezien we doorgaans kleine Transformer-modellen gebruiken om embeddings te produceren, delen we het grootste deel van de implementatie met onze LLM-inferentiecode: batch-embeddings zijn vergelijkbaar met compute-bound prefill, terwijl online embeddings, die vaak op enkele tokens draaien, computationeel vergelijkbaar zijn met memory-bound decode. We hergebruiken dus onze geoptimaliseerde prefill- en decode-kernels om embedding-modellen te serveren. Hierdoor kunnen we een enorme doorvoer voor batch-inferentie bereiken met minimale extra engineering-inspanning, met behoud van een lage latentie voor online embedding-workloads.
Tulips, Roses en wat Ivy
We stellen inferentie bloot via gestandaardiseerde API's, zowel intern als extern via ons API Platform. Achter de schermen zijn meerdere services betrokken bij de verwerking van een embedding-request:
- Ivy is een Rust HTTP-gateway die wordt aangeroepen door Perplexity-services.
Het handelt het werk aan de CPU-kant af voor requests zoals JSON-parsing, tokenisatie, input-templating en batch-splitsing, waarbij requests worden vertaald naar een aangepast gRPC-protocol voor downstream-servers. Deze scheiding stelt ons in staat om bepaalde parameters rondom tokenisatie en input-opmaak te configureren zonder dat we de zwaardere inferentie-instanties hoeven aan te raken.
- Tulip is de inferentieserverinterface.
Het is een gRPC-server geïmplementeerd met Rust, tokio en `tonic`. Tulip ontvangt gRPC-inferentierequests en handelt planning en batching af. Vervolgens stuurt het de batches naar de ROSE-engine en retourneert voltooide responses aan clients.
- **ROSE** (Runtime-Optimized Serving Engine) implementeert model-inferentie.
Het is voornamelijk gedefinieerd in Python en biedt kernels, lagen en definities voor een breed scala aan modellen. ROSE implementeert de forward passes door modellen en biedt tevens CUDA-grafiekbeheer dat is gespecialiseerd voor embeddings. Het is gekoppeld aan Tulip via een step()-functie, die een batch ontvangt en een referentie retourneert naar de berekening die het uitvoert op de accelerator.
Aandacht besteden voorbij de kernel
Zowel op Transformer gebaseerde modellen als de onderliggende Hopper/Blackwell-architecturen zijn volwassen technologieën, waardoor inferentie inbedden aan de GPU-kant is geconvergeerd naar een grotendeels optimale implementatie over verschillende inferentie-engines heen. Desondanks ontdekten we extra mogelijkheden voor verbetering in runtimes en harnesses die de modellen end-to-end aan een client blootstellen. In het bijzonder merkten we dat we latenties kunnen verbeteren door zorgvuldig CUDA-grafieken te beheren en een LazyTensor-abstractie te bouwen om asynchroon een resultaat aan de GPU-kant bij te houden in de native Rust-engine. We hebben deze functies geïmplementeerd in Tulip, zodat deze effectief kon koppelen met de modelimplementaties van ROSE.
Tulip
We hebben Tulip ontworpen om een zo licht mogelijke interface te zijn over onze model-serving heen. Het verwerkt binnenkomende requests in Tokio async-taken, waarbij het een pool van requests bijhoudt en batches inplant om te dispatchen naar de accelerator. Het planningsmechanisme in Tulip is heel eenvoudig: requests hopen zich op terwijl Tulip werk dispatcht of wacht op resultaten. Uit de opgehoopte requests worden sequenties gekozen op basis van wie het eerst komt, het eerst maalt om door het model te worden uitgevoerd.
Het eenvoudige planningsmechanisme wordt gemotiveerd door een observatie over modelprestaties. Voor kleine embedding-modellen merkten we bij de sequentiengtes die we serveren dat de lineaire kosten van dense lagen zwaarder wegen dan de kwadratische kosten van attention. Hierdoor is de latentie grotendeels evenredig met het aantal tokens, en niet met het aantal sequenties. Zodra een batch groot genoeg is om de GPU te verzadigen (wat neerkomt op ongeveer 512 tokens op een model met minder dan een miljard parameters), verbetert het inpakken van meer sequenties in de batch de efficiëntie niet.
Om effectief te koppelen met het model vertrouwt Tulip op CUDA-grafieken en lui resultaten bijhouden om GPU- en CPU-werk te overlappen en de beschikbare resources ten volle te benutten.
CUDA-grafiekbeheer
Het uitvoeren van de forward pass van een model omvat zowel werk aan de CPU-kant als aan de GPU-kant. De CPU is verantwoordelijk voor het plannen van batches en het starten van kernels met de juiste parameters, terwijl de GPU de relevante matrixvermenigvuldigings-, attention-, norm- of activeringskernels uitvoert. Voor workloads met een hoge doorvoer, zoals training en re-indexering, zijn de overheads aan de CPU-kant verwaarloosbaar omdat de batchgroottes en de latentie aan de GPU-kant beide groot zijn. Bij kleinere batchgroottes kan het werk aan de CPU-kant echter zwaarder wegen dan het werk aan de GPU-kant.
Om overheads te verminderen, kan in plaats van het starten van onafhankelijke kernels een CUDA-grafiek worden gebouwd om de metadata vast te leggen die nodig is om alle kernels van een forward pass te starten met een enkele aanroep naar de CUDA-driver. Dit elimineert de noodzaak om dure Python- en PyTorch-code opnieuw uit te voeren voor de configuraties waarvoor CUDA-grafieken kunnen worden vastgelegd.
Voor elk model traceren we een buigpunt dat het minimumaantal tokens bepaalt waarbij GPU-uitvoering duurder is dan het starten van kernels aan de CPU-kant. Omdat embedding-modellen klein zijn, merken we dat dit buigpunt optreedt bij batches van duizenden tokens en tientallen sequenties. Sommige attention-implementaties vertrouwen op dynamische inputs aan de host-kant om kernel-lanceringen te configureren, wat full-model prefill/dense CUDA-grafieken verhindert. We hebben wijzigingen upstream gezet in relevante kernels om deze mogelijk te maken in onze inferentie-engine.
Om overheads aan te pakken, bouwen we CUDA-grafieken voor het gehele model voor alle embedding-modellen en overlappen we CPU-werk met GPU-werk. Aangezien CUDA-grafieken overhead aan de CPU-kant minimaliseren, hebben we, zodra een grafiek is gestart, vrije tijd om de uitvoering van de volgende batch te starten en in de wachtrij te plaatsen zodra deze beschikbaar is. De resultaten van de in behandeling zijnde batch worden bijgehouden met een LazyTensor, waardoor een asynchrone taak in Rust kan blokkeren totdat de vorige batch de uitvoering heeft voltooid. CUDA-grafieken helpen bij serving met lage latentie doordat ze ervoor zorgen dat we niet worden afgeremd door de kosten van kernel-lanceringen, en ze vergemakkelijken een verbeterde planning in het geval van een hoge doorvoer doordat ze de CPU vrijmaken om eerder aan de volgende batch te werken.
CUDA-grafieken moeten worden vastgelegd voor elke distincte configuratie, wat voor embeddings betekent: een grafiek per combinatie van aantal sequenties en aantal tokens. Omdat dit raster expansief is, vullen we aantallen tokens aan tot buckets die een veelvoud zijn van 64 of 256. Dit resulteert nog steeds in duizenden grafieken waarvan het vastleggen bij een typisch model enkele minuten kan duren. De kosten van het vastleggen komen uit twee bronnen: een eager forward pass die moet worden uitgevoerd om kernels te compileren en buffers in te stellen voor diverse kernels die deze nodig hebben, gevolgd door de capture-run die Python-code opnieuw uitvoert.
We beperken opstartkosten door CUDA-grafieken lui vast te leggen naarmate de engine taken bedient. We houden elke configuratie bij en zorgen ervoor dat deze een eager warmuprun doorloopt voordat het vastleggen en afspelen van de grafiek bij de tweede hit wordt getriggerd. Alle volgende executies van dezelfde grafiekconfiguratie doorlopen vervolgens het afspelen van de CUDA-grafiek. Lui vastleggen van grafieken heeft invloed op p99-latenties tijdens het opstarten; het is echter waardevol om meerdere minuten aan eager werk te spreiden over meerdere uur. Snellere opstarttijden stellen ons in staat om embedding-implementaties beter te schalen en te beheren.
Lazy Tensors
Via CUDA is GPU-werk asynchroon. Aangezien het asynchroon starten van een kernel deze in een queue op een stream plaatst, moet hostcode expliciet synchroniseren om de resulterende vectoren uittenezen. Om een hogere graad van theoretiese parallelle uitvoering te vergemakkelijken en toekomstige batches te kunnen starten terwijl wordt gewacht tot de vorige is voltooid op het device, vertrouwen we op een LazyTensor-abstractie om waarden bij te houden.
De LazyTensor volgt een hostbuffer in page-locked geheugen en een cudaMemcpyAsync-operatie via een event dat data van het device kopieert. Dit wordt gestart na de lancering van de forward pass op dezelfde stream. Omdat de kopieeroperatie moet wachten tot alle voorgaande kernels op de stream zijn uitgevoerd, volgt het bijbehorende event zowel de voltooiing van de forward pass als de beschikbaarheid van het resultaat op de CPU.
We benutten LazyTensors in onze ROSE-encoder-engine om GPU- en CPU-werk te overlappen. In plaats van dat elke step()-aanroep de CUDA-grafiek uitvoert en wacht tot deze klaar is, retourneert step() een LazyTensor om het resultaat asynchroon bij te houden. Gekoppeld aan CUDA-grafieken helpt dit ons om lage latenties en een betere doorvoer te bereiken.
ROSE
We hebben onze ROSE-engine, die we oorspronkelijk hebben gebouwd voor LLM-serving, aangepast om ook de uitvoering van embedding-modellen af te handelen. Om de benodigde inspanning voor het ondersteunen van embedding-modellen te minimaliseren, hergebruikt ROSE agressief code tussen LLMs en embeddings. Zo gaan bijvoorbeeld pplx-embed-serving en Qwen3.5 LLM-decoding allemaal door dezelfde kernels. Deze deling stelt ons in staat om eenvoudig een embedding-model te serveren dat oorspronkelijk is gefinetuned vanuit een LLM voor prototyping, evaluatie en productie-inferentie.
Voor dense lagen zijn embedding- en LLM-inferentie identiek omdat tokenvectoren onafhankelijk worden verwerkt. In attention-lagen worden verschillen opgevangen door ondersteuning toe te voegen voor ragged inputs, samen met de paged prefill- en decode-instellingen die vereist zijn voor LLMs. Bij het serveren van een embedding-model instantiëren we geen KV-cache en sturen we aan op varianten van attention-kernels die het ragged formaat ondersteunen om padding te vermijden. De ondersteunende conversie- en kalibratieroutines worden eveneens gedeeld met de LLMs.
Ivy
Ivy, onze inferentie-HTTP-proxylaag, speelt ook een belangrijke rol in de prestaties. Omdat request-payloads variëren in productie, kan het routeren van individuele requests naar individuele replica's leiden tot laadonbalans. Ivy deelt grote batch-requests op in chunks en balanceert de load daartussen tussen replica's, wat het gebruik verbetert en de latentie afvlakt. Ons recente werk aan eigen unigram-tokenisatie, dat volledig is uitgerold in Ivy, verbetert de latenties drastisch ten opzichte van kant-en-klare tokenizers.
...maar de kernels doen er nog steeds toe
ROSE ondersteunt verschillende attention-backends. Verschillende kernels zijn mogelijk geschikt voor specifieke probleemgroottes. Na verloop van tijd hebben we FlashInfer 2-, FlashInfer 3- en FlashAttention 4-kernels geïntegreerd om ragged attention te implementeren.
Over het algemeen merken we dat FlashAttention 4 sneller is. FlashInfer 3 presteert echter beter op Qwen-gebaseerde modellen bij zeer lange sequentiengtes. Omdat prestaties en tuning kunnen variëren met het aantal en de dimensie van attention-heads, behouden we ondersteuning voor meerdere configuraties en maken we een weloverwogen beslissing per geval bij het serveren.
Benchmarks
We benchmarken tegen vLLM v0.22.0, waarbij we inferentie uitvoeren op BF16-precisie op werkelijke modelgewichten en inputs afkomstig van evaluatiedatasets. Aan alle timingruns gingen warmupruns vooraf die verifieerden dat de divergentie in cosinussimilariteit binnen 0,1% valt.
Embeddings met lage latentie (p50 / p90 / p99 / max ms)
We rapporteren runtimes voor voorgetokeniseerde request-batchgrootte 1, volledig sequentiële requests, sequentiengtes van 128, 512 en 4096 tokens.
Scoring met lage latentie (p50 / p90 / p99 / max ms)
Voorgetokeniseerde request-batchgroottes 5, 25 en 50, sequentiengte van 512 tokens.
Embeddings met hoge doorvoer (emb/s)
Request batchgrootte 100, vier gelijktijdige processen die requests indienen, sequentiengtes van 512, 1024 en 4096 tokens.
Embeddings met hoge gelijktijdigheid (p50 / p90 / p99 / max ms)
Sequentiengte 512, batchgrootte 1, maar we sturen 1, 2, 4, 8 en 16 gelijktijdige requests. Deze benchmark omvat ook de kosten voor tokenisatie via Ivy, samen met de netwerkoverhead tussen Ivy en Tulip.
Conclusie en Toekomstig Werk
De serving-infrastructuur bestaande uit Ivy, Tulip en ROSE stelt ons in staat om embeddings te serveren voor Perplexity met een lagere latentie en een betere doorvoer, wat resulteert in nauwkeuriger zoeken tegen lagere kosten in vergelijking met kant-en-klare oplossingen.
Door te focussen op specifieke modellen en de volledige stack in eigen hand te nemen, verkrijgen we de nodige vrijheid om een effectieve balans te vinden tussen prestaties en flexibiliteit, waarbij we sterk herbruikbare en performante Rust-primitieven combineren met meer generieke Python-modelleringscode. Veel open-source inferentie-engines, zoals vLLM, SGLang en TokenSpeed, integreren talen zoals Rust en C++ in hun stack. We hebben de afgelopen twee jaar geïnvesteerd in Rust en hebben daar grote vruchten van geplukt op het gebied van zowel prestaties als onderhoudbaarheid. Door het grootste deel van de embedding-implementatie te delen met onze LLM-serving-stack, behalen we ook winst in doorvoer, zonder dat er aanzienlijke technische inspanningen besteed hoeven te worden aan het onderhoud van embedding-modellen.
Naarmate modellen evolueren, blijven we elke laag van onze stack verbeteren om zowel CPU-gebonden als GPU-gebonden latenties te verminderen. Onze aangepaste gRPC-gebaseerde protocollen binnen Ivy en Tulip stellen ons in staat om communicatie te tweaken om netwerklatenties te verminderen, terwijl ROSE een fundament biedt om de computationele doorvoer te verbeteren. Bovendien, naarmate de ondersteuning voor free-threaded Python groeit binnen het ecosysteem, zullen we de interoperabiliteit tussen Python en Rust verder kunnen verbeteren om overheads te verminderen.