Embeddings rapides sur GPU
Une recherche rapide et précise est essentielle pour l'ensemble de Perplexity, de Search and Computer à notre plateforme API. Dans les coulisses, le travail de fond est effectué par des modèles d'embedding et de classement, qui aident nos systèmes à identifier les résultats les plus pertinents pour une requête donnée. Nous atteignons un niveau de
Une recherche rapide et précise est essentielle pour l'ensemble de Perplexity, de Search and Computer à notre plateforme API. Dans les coulisses, le travail de fond est effectué par des modèles d'embedding et de classement, qui aident nos systèmes à identifier les résultats les plus pertinents pour une requête donnée. Nous atteignons une qualité et une latence de pointe en entraînant et en servant nos propres modèles, tels que pplx-embed.
Cet article présente une vue interne de l'infrastructure de service de Perplexity pour cette classe spéciale de modèles. Nous discutons de nos techniques pour répondre efficacement aux besoins d'inférence de la recherche native par IA, permettant un prototypage et une évaluation rapides des modèles tout en alimentant notre index de recherche à l'échelle de l'exaoctet. Ces techniques élargissent collectivement la frontière de Pareto de la qualité et de l'efficacité de la recherche, nous permettant de servir les agents et les utilisateurs avec les meilleurs résultats possibles au coût et à la latence les plus bas.
Embeddings pour la recherche
Dans une configuration de recherche typique, les documents indexés sont mappés à un espace vectoriel de grande dimension à l'aide d'un modèle d'embedding et stockés dans une base de données vectorielle. En transformant une requête en embedding à l'aide du même modèle, des documents similaires peuvent être localisés en trouvant les vecteurs les plus proches de celui de la requête. Cela donne lieu à deux schémas de trafic différents à desservir pour un moteur d'inférence :
- Embedding par lots : lors de la construction, de l'expansion ou de la réindexation de la base de données, des documents en vrac doivent être transformés en embeddings dans l'espace vectoriel, maximisant le débit pour minimiser le coût.
Après la recherche vectorielle, de grands lots de documents doivent être évalués, trouvant un équilibre entre le débit et la latence.
- Embedding en ligne : lors de l'interrogation de la base de données, une requête courte doit être transformée en embedding pour les recherches, ce qui minimise la latence.
Nous avons conçu notre infrastructure d'inférence pour exploiter autant de composants communs que possible entre les cas d'utilisation. Étant donné que nous utilisons généralement de petits modèles Transformer pour produire des embeddings, nous partageons la majeure partie de l'implémentation avec notre code d'inférence LLM : les embeddings par lots sont similaires au préremplissage limité par le calcul, tandis que les embeddings en ligne, qui s'exécutent souvent sur quelques tokens, sont computationnellement similaires au décodage limité par la mémoire. Nous réutilisons ainsi nos noyaux de préremplissage et de décodage optimisés pour servir les modèles d'embedding. Par conséquent, nous pouvons atteindre un débit d'inférence par lots massif avec un travail d'ingénierie supplémentaire minimal, tout en préservant une faible latence pour les charges de travail d'embedding en ligne.
Tulips, Roses, and some Ivy
Nous exposons l'inférence via des API standardisées, à la fois en interne et en externe via notre plateforme API. Sous le capot, plusieurs services sont impliqués dans le traitement d'une requête d'embedding :
- Ivy est une passerelle HTTP Rust que les services Perplexity appellent.
Il gère le travail côté CPU pour les requêtes telles que l'analyse JSON, la tokenisation, le modèle d'entrée et la division des lots, traduisant les requêtes vers un protocole gRPC personnalisé pour les serveurs en aval. Cette séparation nous permet de configurer certains paramètres concernant la tokenisation et le formatage des entrées sans avoir à toucher aux instances d'inférence plus lourdes.
- Tulip est l'interface du serveur d'inférence.
Il s'agit d'un serveur gRPC implémenté avec Rust, tokio et `tonic`. Tulip reçoit les requêtes d'inférence gRPC, gérant la planification et la mise en lots. Il envoie ensuite les lots au moteur ROSE, renvoyant les réponses terminées aux clients.
- **ROSE** (Runtime-Optimized Serving Engine) implémente l'inférence de modèles.
Il est principalement défini en Python, fournissant des noyaux, des couches et des définitions pour une grande variété de modèles. ROSE implémente les passes avant à travers les modèles, fournissant également une gestion des graphes CUDA spécialisée pour les embeddings. Il est relié à Tulip via une fonction step(), qui prend un lot et renvoie une référence au calcul qu'il effectue sur l'accélérateur.
Prêter attention au-delà du noyau
Les modèles basés sur Transformer et les architectures Hopper/Blackwell sous-jacentes sont des technologies matures. Par conséquent, l'inférence par embedding côté GPU a convergé vers une implémentation largement optimale à travers différents moteurs d'inférence. Malgré cela, nous avons découvert des opportunités d'amélioration supplémentaires dans les moteurs d'exécution et les interfaces qui exposent les modèles de bout en bout à un client. En particulier, nous avons constaté que nous pouvions améliorer les latences en gérant avec soin les graphes CUDA et en construisant une abstraction LazyTensor pour suivre de manière asynchrone un résultat côté GPU dans le moteur Rust natif. Nous avons implémenté ces fonctionnalités dans Tulip afin qu'il puisse s'interfacer efficacement avec les implémentations de modèles de ROSE.
Tulip
Nous avons conçu Tulip pour qu'il soit une interface aussi légère que possible sur le service de nos modèles. Il gère les requêtes entrantes dans des tâches asynchrones Tokio, maintenant un pool de requêtes qu'il suit et à partir duquel il planifie les lots à envoyer à l'accélérateur. Le mécanisme de planification de Tulip est très simple : les requêtes s'accumulent pendant que Tulip distribue du travail ou attend des résultats. Parmi les requêtes accumulées, les séquences sont choisies selon le principe du premier arrivé, premier servi pour être exécutées par le modèle.
Le mécanisme de planification simple est motivé par une observation sur les performances du modèle. Pour les petits modèles d'embedding, aux longueurs de séquence que nous servons, nous avons remarqué que le coût linéaire des couches denses domine sur le coût quadratique de l'attention. Par conséquent, la latence est principalement proportionnelle au nombre de tokens, et non au nombre de séquences. Par conséquent, dès qu'un lot est suffisamment grand pour saturer le GPU (ce qui représente environ 512 tokens sur un modèle de moins d'un milliard de paramètres), y regrouper davantage de séquences n'améliore pas l'efficacité.
Pour s'interfacer efficacement avec le modèle, Tulip s'appuie sur les graphes CUDA et le suivi paresseux des résultats pour superposer le travail du GPU et du CPU et utiliser pleinement les ressources disponibles.
Gestion des graphes CUDA
L'exécution de la passe avant d'un modèle implique à la fois un travail côté CPU et côté GPU. Le CPU est responsable de la planification des lots et du lancement des noyaux avec les paramètres appropriés, tandis que le GPU exécute les noyaux de multiplication matricielle, d'attention, de normalisation ou d'activation correspondants. Pour les charges de travail à haut débit telles que l'entraînement et la réindexation, les surcharges côté CPU sont négligeables car les tailles de lots et la latence côté GPU sont toutes deux élevées. Cependant, pour des tailles de lots plus petites, le travail côté CPU peut l'emporter sur le travail côté GPU.
Pour atténuer les surcharges, au lieu de lancer des noyaux indépendants, un graphe CUDA peut être construit pour capturer les métadonnées requises pour lancer tous les noyaux d'une passe avant avec un seul appel au pilote CUDA. Cela élimine le besoin de réexécuter du code Python et PyTorch coûteux pour les configurations pour lesquelles les graphes CUDA peuvent être capturés.
Pour chaque modèle, nous suivons un point d'inflexion, déterminant le nombre minimal de tokens pour lequel l'exécution sur GPU est plus coûteuse que le lancement de noyaux côté CPU. Les modèles d'embedding étant petits, nous observons que ce point d'inflexion intervient à des lots de milliers de tokens et de dizaines de séquences. Certaines implémentations d'attention reposent sur des entrées dynamiques côté hôte pour configurer les lancements de noyaux, empêchant les graphes CUDA denses de préremplissage de modèle complet. Nous avons intégré en amont des modifications aux noyaux concernés pour les activer dans notre moteur d'inférence.
Pour remédier aux surcharges, nous construisons des graphes CUDA pour modèle entier pour tous les modèles d'embedding et superposons le travail du CPU avec le travail du GPU. Les graphes CUDA minimisant les surcharges côté CPU, une fois qu'un graphe est lancé, nous disposons d'un temps libre pour lancer et mettre en file d'attente l'exécution du lot suivant dès qu'il est disponible. Les résultats du lot en attente sont suivis à l'aide d'un LazyTensor, ce qui permet à une tâche asynchrone en Rust de se bloquer jusqu'à ce que le lot précédent termine son exécution. Les graphes CUDA aident au service à faible latence en garantissant que nous ne sommes pas freinés par le coût des lancements de noyaux et facilitent une meilleure planification dans le cas d'un haut débit en libérant le CPU pour qu'il puisse travailler plus tôt sur le lot suivant.
Les graphes CUDA doivent être capturés pour chaque configuration distincte, ce qui, pour les embeddings, signifie un graphe par combinaison de nombre de séquences et de nombre de tokens. Cette grille étant vaste, nous complétons (pad) les nombres de tokens à des buckets qui sont des multiples de 64 ou 256. Cela se traduit toujours par des milliers de graphes dont la capture peut prendre plusieurs minutes pour un modèle typique. Le coût de la capture provient de deux sources : une passe avant avide qui doit être exécutée pour compiler les noyaux et configurer les tampons pour les divers noyaux qui en ont besoin, suivie de la série de captures qui réexécute le code Python.
Nous atténuons les coûts de démarrage en capturant paresseusement les graphes CUDA au fur et à mesure que le moteur sert. Nous gardons une trace de chaque configuration et veillons à ce qu'elle passe par une exécution d'échauffement avide avant de déclencher la capture et la relecture du graphe lors de la deuxième occurrence. Toutes les exécutions ultérieures de la même configuration de graphe passent ensuite par la relecture du graphe CUDA. La capture de graphes paresseuse a un impact sur les latences p99 au démarrage ; cependant, elle est précieuse pour répartir plusieurs minutes de travail avide sur plusieurs heures. Des temps de démarrage plus rapides nous permettent de mieux dimensionner et gérer les déploiements d'embedding.
Tenseurs paresseux
Grâce à CUDA, le travail du GPU est asynchrone. Le lancement d'un noyau le met en file d'attente de manière asynchrone sur un flux, le code de l'hôte doit donc se synchroniser explicitement pour lire les vecteurs résultants. Pour faciliter un degré de parallélisme plus élevé et pouvoir lancer de futurs lots tout en attendant que le précédent se termine sur le périphérique, nous nous appuyons sur une abstraction LazyTensor pour suivre les valeurs.
Le LazyTensor suit un tampon hôte dans une mémoire verrouillée en page et une opération cudaMemcpyAsync via un événement copiant les données du périphérique. Il est déclenché après le lancement de la passe avant sur le même flux. L'opération de copie devant attendre l'exécution de tous les noyaux précédents sur le flux, l'événement associé suit à la fois la fin de la passe avant et la disponibilité du résultat sur le CPU.
Nous exploitons les LazyTensors dans notre moteur d'encodeur ROSE pour superposer le travail du GPU et du CPU. Au lieu que chaque appel de step() exécute le graphe CUDA et attende qu'il se termine, step() renvoie un LazyTensor pour suivre son résultat de manière asynchrone. Couplé aux graphes CUDA, cela nous aide à obtenir de faibles latences et un meilleur débit.
ROSE
Nous avons adapté notre moteur ROSE, initialement conçu pour le service de LLM, pour gérer également l'exécution de modèles d'embedding. Pour minimiser l'effort nécessaire à la prise en charge des modèles d'embedding, ROSE réutilise agressivement le code entre les LLM et les embeddings. Par exemple, le service pplx-embed et le décodage LLM Qwen3.5 passent tous par les mêmes noyaux. Ce partage nous permet de servir facilement un modèle d'embedding qui a été initialement affiné à partir d'un LLM pour le prototypage, l'évaluation et l'inférence en production.
Pour les couches denses, l'inférence par embedding et par LLM est identique puisque les vecteurs de tokens sont traités indépendamment. Dans les couches d'attention, les différences sont gérées en ajoutant la prise en charge des entrées déchiquetées (ragged), ainsi que les configurations de préremplissage et de décodage paginées requises par les LLM. Lors du service d'un modèle d'embedding, nous n'instancions pas de cache KV et nous acheminons vers des variantes de noyaux d'attention qui prennent en charge le format déchiqueté pour éviter le remplissage (padding). Les routines de conversion et de calibrage associées sont également partagées avec les LLM.
Ivy
Ivy, notre couche de proxy HTTP d'inférence, joue également un rôle important dans les performances. Les charges utiles des requêtes variant en production, le routage de requêtes individuelles vers des réplicas individuels peut entraîner un déséquilibre de charge. Ivy divise les requêtes de grands lots en morceaux et équilibre leur charge entre les réplicas, améliorant ainsi l'utilisation et lissant la latence. Notre travail récent sur la tokenisation par unigrammes en interne, entièrement déployé dans Ivy, améliore considérablement les latences par rapport aux tokeniseurs prêts à l'emploi.
...mais les noyaux comptent toujours
ROSE prend en charge une variété de moteurs d'attention. Différents noyaux peuvent être adaptés à des tailles de problèmes spécifiques. Au fil du temps, nous avons intégré les noyaux FlashInfer 2, FlashInfer 3 et FlashAttention 4 pour implémenter l'attention déchiquetée.
En général, nous constatons que FlashAttention 4 est plus rapide. Cependant, FlashInfer 3 le surpasse sur les modèles basés sur Qwen pour des séquences de très grande longueur. Les performances et le réglage pouvant varier en fonction du nombre et de la dimension des têtes d'attention, nous maintenons la prise en charge de plusieurs configurations et prenons une décision au cas par cas lors du service.
Benchmarks
Nous effectuons des benchmarks par rapport à vLLM v0.22.0, en exécutant l'inférence en précision BF16 sur des poids de modèles réels et des entrées dérivées de jeux de données d'évaluation. Tous les lancements de chronométrage ont été précédés par des exécutions d'échauffement qui ont vérifié que la divergence de la similarité cosinus est inférieure à 0,1 %.
Embeddings à faible latence (p50 / p90 / p99 / max ms)
Nous rapportons les temps d'exécution pour une taille de lot de requêtes prétokenisées de 1, des requêtes entièrement séquentielles et des longueurs de séquence de 128, 512 et 4096 tokens.
Scoring à faible latence (p50 / p90 / p99 / max ms)
Tailles de lots de requêtes prétokenisées de 5, 25 et 50, longueur de séquence de 512 tokens.
Embeddings à haut débit (emb/s)
Taille de lot de requêtes de 100, quatre processus concurrents soumettant des requêtes, longueurs de séquence de 512, 1024 et 4096 tokens.
Embeddings à haute concurence (p50 / p90 / p99 / max ms)
Longueur de séquence de 512, taille de lot de 1, mais nous envoyons 1, 2, 4, 8 et 16 requêtes concurrentes. Ce benchmark inclut également les coûts de tokenisation via Ivy, ainsi que la surcharge réseau entre Ivy et Tulip.
Conclusion et travaux futurs
L'infrastructure de service composée d'Ivy, de Tulip et de ROSE nous permet de servir des embeddings pour Perplexity avec une latence plus faible et un meilleur débit, ce qui se traduit par une recherche plus précise à un coût réduit par rapport aux solutions prêtes à l'emploi.
En nous concentrant sur des modèles spécifiques et en prenant en charge l'ensemble de la pile, nous obtenons la liberté nécessaire pour trouver un équilibre efficace entre performances et flexibilité, en combinant des primitives Rust hautement réutilisables et performantes avec un code de modélisation Python plus générique. De nombreux moteurs d'inférence open-source, tels que vLLM, SGLang et TokenSpeed, intègrent des langages comme Rust et C++ dans leur pile. Nous avons investi dans Rust au cours des deux dernières années et en avons récolté de grands fruits en termes de performances et de maintenabilité. En partageant la majeure partie de l'implémentation des embeddings avec notre pile de service LLM, nous obtenons également des gains de débit, sans nécessiter d'efforts d'ingénierie importants consacrés à la maintenance des modèles d'embedding.
À mesure que les modèles évoluent, nous continuerons d'améliorer chaque couche de notre pile pour réduire les latences liées au CPU et au GPU. Nos protocoles personnalisés basés sur gRPC au sein d'Ivy et de Tulip nous permettent de modifier la communication pour réduire les latences réseau, tandis que ROSE fournit une base pour améliorer le débit de calcul. De plus, à mesure que la prise en charge de Python sans GIL (free-threaded) se développe dans tout l'écosystème, nous serons en mesure d'améliorer encore l'interopérabilité entre Python et Rust pour réduire les surcharges.