Vectorización del Trie: Decodificación Restringida Eficiente para Recuperación Generativa Basada en LLM en Aceleradores
Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval on Accelerators
February 26, 2026
Autores: 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
Resumen
La recuperación generativa ha surgido como un paradigma poderoso para la recomendación basada en LLM. Sin embargo, los sistemas de recomendación industriales a menudo se benefician de restringir el espacio de salida a un subconjunto limitado de elementos basado en lógica de negocio (por ejemplo, imponer frescura de contenido o categoría de producto), lo cual la decodificación autoregresiva estándar no puede soportar de forma nativa. Además, los métodos de decodificación restringida existentes que utilizan árboles de prefijos (Tries) incurren en penalizaciones de latencia severas en aceleradores de hardware (TPUs/GPUs). En este trabajo, presentamos STATIC (Matriz de Transición Dispersa-Acelerada con Índice Trie para Decodificación Restringida), una técnica de decodificación restringida eficiente y escalable diseñada específicamente para la recuperación generativa de alto rendimiento basada en LLM en TPUs/GPUs. Al aplanar el árbol de prefijos en una matriz estática de Filas Dispersas Comprimidas (CSR), transformamos los recorridos irregulares del árbol en operaciones de matriz dispersa completamente vectorizadas, desbloqueando ganancias masivas de eficiencia en aceleradores de hardware. Desplegamos STATIC en una plataforma de recomendación de videos a gran escala que sirve a miles de millones de usuarios. STATIC produce un impacto significativo en las métricas de producto con una sobrecarga de latencia mínima (0.033 ms por paso y 0.25% del tiempo de inferencia), logrando una aceleración de 948x sobre una implementación de trie en CPU y una aceleración de 47-1033x sobre una línea base de búsqueda binaria acelerada por hardware. Además, la sobrecarga en tiempo de ejecución de STATIC se mantiene extremadamente baja en una amplia gama de configuraciones prácticas. Hasta donde sabemos, STATIC permite el primer despliegue a escala de producción de recuperación generativa estrictamente restringida. Adicionalmente, la evaluación en benchmarks académicos demuestra que STATIC puede mejorar considerablemente el rendimiento en arranque en frío para la recuperación generativa. Nuestro código está disponible en 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.