🤖 machine learning

FlashTrie: A GPU-Accelerated Constrained Beam Search for Generative Retrieval

FlashTrie est un système accéléré par GPU qui optimise la recherche par faisceau contrainte pour la recherche générative en employant une disposition de trie compressée par bits et des noyaux CUDA coopératifs pour éliminer les goulots d'étranglement du CPU, atteignant jusqu'à 24x de vitesse supplémentaire et une augmentation de revenu de 0,71 % dans les applications de recherche commerciales à grande échelle.

Auteurs originaux : Dakshitha Anandakumar, Anurag Mukkara, Wenxiang Hu, Jiusheng Chen, M Akash Kumar, Ting Ye, Qiang Lou, Jian Jiao

Publié 2026-07-14

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Dakshitha Anandakumar, Anurag Mukkara, Wenxiang Hu, Jiusheng Chen, M Akash Kumar, Ting Ye, Qiang Lou, Jian Jiao

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous êtes un robot super intelligent essayant d'écrire une liste de codes secrets (comme « DocID : 4592 ») basés sur une question que vous venez d'entendre. Mais il y a un piège : vous ne pouvez écrire que des codes qui existent réellement dans un immense annuaire pré-approuvé de 800 millions d'entrées valides. Si vous devinez un code qui n'est pas dans le livre, c'est un échec.

Pendant longtemps, les robots faisaient cela en demandant à un bibliothécaire très rapide et très organisé (fonctionnant sur une puce informatique standard, ou CPU) de vérifier chaque supposition. Mais à mesure que la liste des suppositions augmentait, le bibliothécaire était débordé. La vérification de l'annuaire devenait un embouteillage, ralentissant tout le processus. Le robot devait attendre son tour, étape par étape, pour voir si sa supposition était autorisée.

C'est là qu'entre en scène FlashTrie. Les chercheurs de Microsoft et Nvidia ont décidé de licencier le bibliothécaire et de déplacer l'intégralité de l'annuaire de 800 millions d'entrées directement dans la mémoire ultra-rapide du robot (le GPU). Mais ils n'ont pas seulement déplacé le livre ; ils l'ont reconstruit.

La magie de l'annuaire « Bit-Packed »

Imaginez l'ancien annuaire comme une immense bibliothèque où chaque livre est stocké dans une grande pièce vide avec beaucoup d'espace gaspillé. FlashTrie réduit la taille des livres. Il utilise une astuce ingénieuse appelée « compression de bits » pour serrer l'information, comme si l'on emballait une valise si efficacement que l'on peut faire tenir 800 millions de mots-clés dans seulement 3,1 Go d'espace. C'est assez petit pour tenir entièrement dans la mémoire haute vitesse du robot, de sorte qu'il n'ait jamais besoin d'attendre que le disque dur externe lent aille chercher une page.

La danse coopérative

Dans l'ancien système, le robot faisait une supposition, demandait au bibliothvia de la vérifier, attendait une réponse, faisait une autre supposition, et ainsi de suite. C'était un processus solitaire et séquentiel.

FlashTrie change la donne. Il utilise un « noyau CUDA coopératif », qui est comme une piste de danse massive avec 512 danseurs (threads) travaillant ensemble en parfaite synchronisation.

L'expansion : Au lieu d'une seule personne vérifiant une seule supposition, des centaines de danseurs vérifient des milliers de supposations exactement au même moment.
La validation : Ils utilisent une « recherche binaire parallèle » (une façon ultra-rapide de rechercher des informations) pour voir si les suppositions correspondent à l'annuaire.
L'élagage : Si une supposition est mauvaise, ils la rejettent immédiatement. Si elle est bonne, ils la conservent.

Parce que tout se passe sur la piste de danse (le GPU) sans que le robot n'ait à s'arrêter pour parler à l'ordinateur principal (le CPU) après chaque étape, le processus devient incroyablement rapide.

Les résultats : Vitesse et Intelligence

L'équipe a testé cela sur une bibliothèque de 800 millions de mots-clés.

Vitesse : Lorsqu'ils ont augmenté le nombre de suppositions (la « largeur de faisceau » ou beam width) à 1 000, l'ancien système CPU mettait environ 46 millisecondes et devenait de plus en plus lent à mesure que la liste augmentait. FlashTrie a maintenu le temps sous les 3 millisecondes (plus précisément, la moyenne était de 1,91 ms et les 1 % les plus lents étaient sous les 3,31 ms).
Le boost : Cela signifie que FlashTrie est jusqu'à 24 fois plus rapide que la version CPU hautement optimisée.
Qualité : Crucialement, être plus rapide ne signifiait pas être moins précis. FlashTrie a trouvé autant de codes corrects que le système lent. En fait, parce que FlashTrie est si rapide, le robot pouvait vérifier 600 suppositions au lieu de seulement 200 sans dépasser la limite de temps.

Impact dans le monde réel : Le test de l'argent

Les chercheurs ne se sont pas contentés de tests en laboratoire. Ils ont testé FlashTrie dans un moteur de recherche commercial réel (celui que vous utilisez peut-être pour chercher des choses en ligne). Ils ont mené une expérience pendant 16 jours à travers différents pays.

En utilisant FlashTire pour vérifier plus de suppositions, le moteur de recherche a affiché de meilleures publicités.
Cela a conduit à une augmentation de 0,71 % des revenus (l'argent généré par les publicités).
Cela a également augmenté les clics de 0,17 % pour les requêtes en anglais et de 0,20 % pour les requêtes non anglaises.
De manière importante, la qualité des publicités n'a pas baissé ; le « taux de défaut » (les mauvaises publicités affichées) est resté le même.

Ce que FlashTrie n'est PAS

Il est important de noter ce que cet article indique qui ne fonctionne pas ou n'est pas nécessaire ici. Les chercheurs ont explicitement exclu l'utilisation des bibliothèques de type « basé sur les pointeurs » sur le GPU car elles provoquent trop de confusion et ralentissent les danseurs. Ils ont également montré que le simple fait de déplacer l'ancien système vers le GPU sans redéfinir la structure des données (comme une méthode de « sondage linéaire » ou Linear-probe) serait 71 à 209 fois plus lent que leur nouvelle méthode. L'accélération provient de la conception spécifique de l'annuaire et de la danse, et non pas seulement de l'utilisation d'un matériel plus rapide.

L'essentiel

FlashTrie prouve que vous n'avez pas à choisir entre vitesse et précision. En redéfinissant la façon dont l'« annuaire » est stocké et la façon dont la « vérification » se déroule, ils ont transformé un goulot d'étranglement séquentiel et lent en une fête parallèle ultra-rapide. Cela permet aux robots de penser plus grand (vérifier plus d'options) et plus vite, tout en respectant les limites de temps strictes nécessaires pour les recherches sur Internet en temps réel. Le code de ce système sera rendu public après le processus de révision, afin que d'autres puissent essayer cette nouvelle façon de chercher.

Résumé Technique : FlashTrie

Problématique

La recherche générative reformule la recherche de documents en une tâche de séquence à séquence, mappant directement les requêtes vers des identifiants de documents (docIDs). Bien que ce paradigme offre une modélisation de bout en bout compacte, le service en ligne est contraint par des budgets de latence stricts. La procédure de décodage, particulièrement pour les décodeurs non autorégressifs (NAR), génère fréquemment des identifiants invalides qui nécessitent un filtrage via une recherche contrainte.

Les approches de pointe actuelles utilisent la recherche par faisceau contrainte par trie (trie-constrained beam search) pour garantir que les identifiants générés correspondent à une bibliothèque prédéfinie. Cependant, les implémentations existantes reposent sur des représentations de trie basées sur des pointeurs (par exemple, le standard MARISA) s'exécutant sur CPU. Ces conceptions souffrent de :

Accès mémoire irréguliers : Entraînant une faible utilisation du matériel et une latence élevée.
Parallélisme limité : À mesure que la largeur du faisceau (beam width) augmente, la traversée et la validation des candidats sur CPU deviennent le goulot d'étranglement dominant du service.
Inefficacité sur GPU : Les tries succinctes traditionnelles (comme MARISA) sont optimisées pour les clés au niveau du caractère et s'adaptent mal aux architectures GPU, provoquant une divergence de warp et des accès mémoire non coalescés.
Surcharge d'orchestration : L'application des contraintes est souvent implémentée en dehors du chemin de décodage principal du GPU, nécessissant des synchronisations répétées hôte-périphérique à chaque étape de décodage.

L'article répond à trois questions clés : Les tries succinctes peuvent-elles être adaptées aux vocabulaires de type jeton (token-level) sans sacrifier l'efficacité de l'espace ? La recherche par faisceau contrainte peut-elle être exécutée entièrement sur GPU ? De tels systèmes peuvent-ils respecter les contraintes de latence de production sans dégrader la qualité de la recherche ?

Méthodologie : FlashTrie

FlashTrie est un cadre de décodage contraint natif pour GPU qui optimise conjointement la représentation des contraintes et le calcul du décodage. Il remplace l'orchestration sur CPU par un noyau CUDA coopératif entièrement sur périphérique.

1. Structure de données : Trie succinct conscient des entiers (Integer-Aware Succinct Trie)

FlashTrie redessine la disposition du trie MARISA pour supporter de grands vocabulaires de jetons entiers (plus de 1 million de jetons) tout en maintenant une empreinte mémoire réduite, adaptée à la mémoire haute bande passante (HBM) du GPU.

LOUDS étroit (Narrow LOUDS) : Les auteurs étendent la représentation LOUDS (Level-Order Unary Degree Sequence) aux jetons entiers 32 bits.
Compression de bits : Pour réduire l'utilisation de la mémoire, la disposition sépare les étiquettes de nœuds et les indices de suffixe en une partie basse de 2 octets (stockée dans le tableau de base) et une partie haute compressée par bits (stockée dans un tableau secondaire pour les nœuds qui en nécessitent une).
Avantages : Cet « amincissement » réduit de moitié la taille du tableau dominant, rétrécit l'index de contrainte (faisant tenir 800 millions de mots-clés dans environ 3 Go) et améliore la localité de la mémoire en réduisant le nombre de bits extraits par accès.

2. Algorithme : Recherche par faisceau coopérative sur périphérique

FlashTrie exécute l'intégralité de la recherche par faisceau contrainte au sein d'un seul noyau CUDA persistant, éliminant l'orchestration hôte par étape.

Parallélisme à deux niveaux :
- Parallélisme de faisceau (Beam-Parallel) : Les faisceaux actifs à travers une requête sont mappés sur différents blocs de threads (CTAs).
- Parallélisme de proposition (Proposal-Parallel) : Au sein de chaque faisceau, les threads effectuent une validation parallèle. Chaque thread teste l'une des $K$ meilleures propositions du modèle de langage contre le trie en utilisant une recherche binaire parallèle sur les étiquettes d'enfants triées.
Flux de travail du noyau :
1. Expansion : Les faisaux actifs sont étendus à l'aide des $K$ meilleures propositions.
2. Validation : Les threads effectuent une recherche binaire sur les étiquettes d'enfants du trie pour valider les préfixes.
3. Élagage (Pruning) : Un noyau de sélection top- $B$ parallèle remplace le traditionnel tas (heap) CPU. Il utilise un tampon de candidats en mode ajout uniquement et un tri parallèle pour élaguer les faisceaux à la largeur cible, évitant ainsi la contention de verrouillage et la poursuite de pointeurs.
Synchronisation : Une barrière à l'échelle de la grille synchronise les étapes, permettant au noyau de s'exécuter pendant $T$ étapes de décodage en un seul lancement.

Contributions Clés

Index de contrainte résident sur GPU : Une nouvelle disposition de trie compressée par bits qui maintient des bibliothèques de contraintes à l'échelle du milliard entièrement dans la HBM du GPU, réduisant les blocages mémoire.
Noyau CUDA coopératif : Un modèle d'exécution à noyau unique qui effectue l'expansion du faisceau, la validation et l'élagage sur périphérique, supprimant le surcoût des lancements de noyaux répétés et de l'orchestration hôte.
Primitives conscientes du GPU : Le remplacement de la traversée irrégulière basée sur les pointeurs par une recherche binaire parallèle et une sélection top- $K$ sans tas, améliorant considérablement l'utilisation des warps et réduisant la divergence.
Scalabilité : Le système fait passer le décodage contraint à 800 millions de mots-clés sur un seul GPU, un régime auparavant réalisable uniquement en mode hors ligne.

Résultats Expérimentaux

Le système a été évalué sur un modèle NAR de production avec un vocabulaire de 2,2 millions de jetons et un trie de contrainte de 800 millions de mots-clés, tournant sur un GPU NVIDIA A100 80GB.

Métriques de Performance

Latence : FlashTrie réduit la latence de recherche dans le trie à moins de 3 ms (p99) même avec une largeur de faisceau ($BW$) de 1000.
- À $BW=1000$, FlashTrie atteint une latence moyenne de 1,91 ms, contre 46,30 ms pour la base de référence CPU optimisée (MARISA-Opt).
- Cela représente une accélération de 24× par rapport à la base de référence CPU multithreadée hautement optimisée.
Débit : Sous service par lots (batching $b=8$ ), FlashTrie culmine à 1 621 requêtes/seconde pour $BW=1000$, offrant un débit 32× à 71× plus élevé que la base de référence CPU.
Construction : FlashTrie construit l'index 3,4× plus vite et utilise 22 % d'espace disque en moins que MARISA-Opt.

Qualité de la Recherche

Parité avec le CPU : FlashTrie correspond à la qualité de recherche de la base de référence CPU (MARISA-Opt) à 0,001 de précision près pour toutes les largeurs de faisceau.
Gains de qualité via la largeur de faisceau : Parce que FlashTrie maintient une faible latence, il permet l'utilisation de largeurs de faisceau beaucoup plus importantes (jusqu'à 1000) qui sont irréalisables sur CPU. Augmenter la largeur du faisceau de 100 à 1000 améliore la précision@100 de 0,53 à 0,78.

Test A/B en ligne

Dans une expérience en ligne à grande échelle sur un moteur de recherche sponsorisé commercial :

Configuration : Le groupe de traitement utilisait FlashTrie avec $BW=600$, tandis que le groupe de contrôle utilisait MARISA-Opt avec $BW=200$ (le maximum réalisable pour le CPU dans les limites de latence).
Résultats : FlashTrie a délivré une hausse de revenus statistiquement significative de +0,71 % et une augmentation de 0,17–0,20 % du taux de clic sans dégrader la qualité des publicités (les taux de défaut sont restés non significatifs).
Latence : Le groupe de traitement a respecté des budgets de latence stricts, réduisant la latence p95 de 32 % par rapport à la limite du groupe de contrôle.

Signification et Revendications

L'article soutient que le décodage contraint ne doit pas rester un goulot d'étranglement du CPU. En redessinant les structures de données pour l'exécution parallèle et en déplaçant toute la logique d'application des contraintes sur le GPU, FlashTrie permet :

Mise à l'échelle en temps réel : Il permet aux systèmes de recherche générative d'augmenter la taille des faisceaux jusqu'à 5× dans les applications critiques en termes de latence (comme la recherche sponsorisée) sans violer les budgets de service.
Résolution du compromis Qualité-Performance : Il découple la qualité de la recherche de la latence, permettant aux systèmes de tirer parti de largeurs de faisceau plus importantes pour une meilleure sélection de candidats tout en maintenant une faible latence.
Applicabilité Générale : L'approche est applicable à toute tâche de génération contrainte avec de grands vocabulaires statiques, comme l'extraction d'entités ou la recherche de catalogues de produits.

Les auteurs concluent que FlashTrie démontre un schéma plus large où les structures de données traditionnellement implémentées sur CPU peuvent devenir hautement efficaces sur GPU lorsqu'elles sont redessinées pour une exécution par lots et parallèle.