Vectorisation du Trie : Décodage Contraint Efficace pour la Récupération Génératrice basée sur LLM sur Accélérateurs
Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval on Accelerators
February 26, 2026
Auteurs: Zhengyang Su, Isay Katsman, Yueqi Wang, Ruining He, Lukasz Heldt, Raghunandan Keshavan, Shao-Chuan Wang, Xinyang Yi, Mingyan Gao, Onkar Dalal, Lichan Hong, Ed Chi, Ningren Han
cs.AI
Résumé
La récupération générative est devenue un paradigme puissant pour la recommandation basée sur les LLM. Cependant, les systèmes de recommandation industriels bénéficient souvent de la restriction de l'espace de sortie à un sous-ensemble contraint d'éléments basé sur la logique métier (par exemple, en imposant la fraîcheur du contenu ou une catégorie de produits), ce que le décodage autogressif standard ne peut pas prendre en charge nativement. De plus, les méthodes de décodage contraint existantes qui utilisent des arbres de préfixes (Tries) entraînent des pénalités de latence sévères sur les accélérateurs matériels (TPU/GPU). Dans ce travail, nous présentons STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding), une technique de décodage contraint efficace et évolutive conçue spécifiquement pour la récupération générative à haut débit basée sur les LLM sur les TPU/GPU. En aplatissant l'arbre de préfixes en une matrice statique au format Compressed Sparse Row (CSR), nous transformons les parcours d'arbres irréguliers en opérations entièrement vectorisées sur des matrices creuses, permettant des gains d'efficacité massifs sur les accélérateurs matériels. Nous déployons STATIC sur une plateforme de recommandation vidéo industrielle à grande échelle desservant des milliards d'utilisateurs. STATIC produit un impact significatif sur les métriques produits avec une surcharge de latence minimale (0,033 ms par étape et 0,25 % du temps d'inférence), obtenant une accélération de 948x par rapport à une implémentation de trie sur CPU et une accélération de 47 à 1033x par rapport à une baseline à recherche binaire accélérée matériellement. De plus, la surcharge d'exécution de STATIC reste extrêmement faible sur une large gamme de configurations pratiques. À notre connaissance, STATIC permet le premier déploiement à l'échelle de production de la récupération générative strictement contrainte. Par ailleurs, l'évaluation sur des benchmarks académiques démontre que STATIC peut considérablement améliorer les performances en cold-start pour la récupération générative. Notre code est disponible à l'adresse https://github.com/youtube/static-constraint-decoding.
English
Generative retrieval has emerged as a powerful paradigm for LLM-based recommendation. However, industrial recommender systems often benefit from restricting the output space to a constrained subset of items based on business logic (e.g. enforcing content freshness or product category), which standard autoregressive decoding cannot natively support. Moreover, existing constrained decoding methods that make use of prefix trees (Tries) incur severe latency penalties on hardware accelerators (TPUs/GPUs). In this work, we introduce STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding), an efficient and scalable constrained decoding technique designed specifically for high-throughput LLM-based generative retrieval on TPUs/GPUs. By flattening the prefix tree into a static Compressed Sparse Row (CSR) matrix, we transform irregular tree traversals into fully vectorized sparse matrix operations, unlocking massive efficiency gains on hardware accelerators. We deploy STATIC on a large-scale industrial video recommendation platform serving billions of users. STATIC produces significant product metric impact with minimal latency overhead (0.033 ms per step and 0.25% of inference time), achieving a 948x speedup over a CPU trie implementation and a 47-1033x speedup over a hardware-accelerated binary-search baseline. Furthermore, the runtime overhead of STATIC remains extremely low across a wide range of practical configurations. To the best of our knowledge, STATIC enables the first production-scale deployment of strictly constrained generative retrieval. In addition, evaluation on academic benchmarks demonstrates that STATIC can considerably improve cold-start performance for generative retrieval. Our code is available at https://github.com/youtube/static-constraint-decoding.